171 lines
6.9 KiB
Python
171 lines
6.9 KiB
Python
"""Project knowledge-base store: indexed docs searchable per project.
|
|
|
|
Zero third-party dependencies; Python standard library only.
|
|
|
|
Design notes:
|
|
- Every document belongs to exactly one ``project_id``. Every search
|
|
filters on ``project_id`` first: results from another project are
|
|
never returned. This project isolation is a hard rule, not a default.
|
|
- Full-text search uses an FTS5 virtual table ``kb_fts(content)`` when
|
|
the local sqlite build supports it; otherwise search falls back to
|
|
``LIKE``. The fallback is chosen once at construction time (try to
|
|
create the FTS5 table, record whether it worked) and the search path
|
|
follows that branch, so a build without FTS5 never crashes.
|
|
- Re-indexing the same ``project_id`` + ``path`` replaces the old
|
|
document (old row and its FTS row are deleted first), so repeated
|
|
indexing never creates duplicates.
|
|
"""
|
|
|
|
import sqlite3
|
|
from cloud.db import connect as _db_connect
|
|
import time
|
|
|
|
_KINDS = ("file", "readme", "doc", "task", "review")
|
|
|
|
|
|
class KBStore:
|
|
"""SQLite-backed, per-project knowledge base with keyword search."""
|
|
|
|
def __init__(self, db_path):
|
|
self.db_path = db_path
|
|
self.conn = _db_connect(db_path)
|
|
self.conn.row_factory = sqlite3.Row
|
|
self.conn.execute(
|
|
"""CREATE TABLE IF NOT EXISTS kb_docs(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
project_id TEXT,
|
|
path TEXT,
|
|
content TEXT,
|
|
kind TEXT,
|
|
created_at REAL
|
|
)"""
|
|
)
|
|
# Try FTS5; fall back to LIKE if this sqlite build lacks it.
|
|
self._fts = False
|
|
try:
|
|
self.conn.execute(
|
|
"CREATE VIRTUAL TABLE IF NOT EXISTS kb_fts USING fts5(content)"
|
|
)
|
|
self._fts = True
|
|
except sqlite3.OperationalError:
|
|
self._fts = False
|
|
self.conn.commit()
|
|
|
|
def index_file(self, project_id, path, content, kind="file"):
|
|
"""Index one document; returns the new row id.
|
|
|
|
If the same ``project_id`` + ``path`` was already indexed, the
|
|
old document (and its FTS row, when FTS5 is in use) is deleted
|
|
first, so re-indexing replaces rather than duplicates.
|
|
"""
|
|
if kind not in _KINDS:
|
|
raise ValueError("kind must be one of %s, got %r" % (_KINDS, kind))
|
|
old_ids = [
|
|
row["id"]
|
|
for row in self.conn.execute(
|
|
"SELECT id FROM kb_docs WHERE project_id = ? AND path = ?",
|
|
(project_id, path),
|
|
).fetchall()
|
|
]
|
|
for old_id in old_ids:
|
|
if self._fts:
|
|
self.conn.execute(
|
|
"DELETE FROM kb_fts WHERE rowid = ?", (old_id,)
|
|
)
|
|
self.conn.execute("DELETE FROM kb_docs WHERE id = ?", (old_id,))
|
|
cur = self.conn.execute(
|
|
"""INSERT INTO kb_docs
|
|
(project_id, path, content, kind, created_at)
|
|
VALUES (?, ?, ?, ?, ?)""",
|
|
(project_id, path, content, kind, time.time()),
|
|
)
|
|
doc_id = cur.lastrowid
|
|
if self._fts:
|
|
self.conn.execute(
|
|
"INSERT INTO kb_fts(rowid, content) VALUES (?, ?)",
|
|
(doc_id, content),
|
|
)
|
|
self.conn.commit()
|
|
return doc_id
|
|
|
|
def index_task(self, project_id, task_id, prompt, conclusion, kind="task"):
|
|
"""Index a task's prompt + conclusion as one document.
|
|
|
|
The document's path is ``"task:<task_id>"`` and its content is
|
|
the prompt and conclusion joined together.
|
|
"""
|
|
content = "%s\n%s" % (prompt, conclusion)
|
|
return self.index_file(project_id, "task:%s" % task_id, content, kind=kind)
|
|
|
|
def search(self, project_id, query, top_k=5):
|
|
"""Search one project's docs by keyword; returns list[dict].
|
|
|
|
Each dict has ``path``, ``kind`` and ``snippet`` (first 200
|
|
characters of content), plus ``id`` and ``project_id``.
|
|
The ``project_id`` filter is always applied: docs from other
|
|
projects are never returned.
|
|
"""
|
|
if not query:
|
|
return []
|
|
if self._fts:
|
|
try:
|
|
rows = self.conn.execute(
|
|
"""SELECT d.id AS id, d.project_id AS project_id,
|
|
d.path AS path, d.kind AS kind,
|
|
d.content AS content
|
|
FROM kb_docs d JOIN kb_fts f ON f.rowid = d.id
|
|
WHERE d.project_id = ? AND kb_fts MATCH ?
|
|
ORDER BY d.created_at DESC, d.id DESC
|
|
LIMIT ?""",
|
|
(project_id, self._fts_query(query), top_k),
|
|
).fetchall()
|
|
if rows:
|
|
return [self._row_to_result(row) for row in rows]
|
|
# Empty FTS result falls through to LIKE below. This is
|
|
# not just a safety net: FTS5's default unicode61
|
|
# tokenizer does not split CJK text into searchable
|
|
# words, so a Chinese keyword that is present verbatim
|
|
# in a document (verified: "登录" in a doc containing
|
|
# "用户登录流程") returns zero FTS hits. LIKE substring
|
|
# search recovers those cases; the project_id filter
|
|
# still applies, so isolation is unchanged.
|
|
except sqlite3.OperationalError:
|
|
# Query syntax FTS5 cannot parse, or a broken FTS table:
|
|
# fall through to the LIKE path rather than crashing.
|
|
pass
|
|
# LIKE fallback (also the path when FTS5 is unavailable).
|
|
pattern = "%" + query.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + "%"
|
|
rows = self.conn.execute(
|
|
"""SELECT id, project_id, path, kind, content FROM kb_docs
|
|
WHERE project_id = ? AND content LIKE ? ESCAPE '\\'
|
|
ORDER BY created_at DESC, id DESC
|
|
LIMIT ?""",
|
|
(project_id, pattern, top_k),
|
|
).fetchall()
|
|
return [self._row_to_result(row) for row in rows]
|
|
|
|
def list_docs(self, project_id):
|
|
"""List all documents of one project, newest first."""
|
|
rows = self.conn.execute(
|
|
"""SELECT id, project_id, path, kind, content, created_at
|
|
FROM kb_docs WHERE project_id = ?
|
|
ORDER BY created_at DESC, id DESC""",
|
|
(project_id,),
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
@staticmethod
|
|
def _fts_query(query):
|
|
"""Turn a raw keyword query into a safe FTS5 phrase query."""
|
|
# Phrase-quoting escapes special FTS5 syntax in the user query;
|
|
# a multi-word query is searched as one phrase, which is the
|
|
# conservative reading for keyword lookup.
|
|
return '"%s"' % query.replace('"', '""')
|
|
|
|
@staticmethod
|
|
def _row_to_result(row):
|
|
d = dict(row)
|
|
content = d.pop("content", "") or ""
|
|
d["snippet"] = content[:200]
|
|
return d
|