386 lines
16 KiB
Python
386 lines
16 KiB
Python
"""Bounded, local SQLite FTS index for user-selected workspace documents."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import sqlite3
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import struct
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Any, Iterator
|
|
|
|
from app import database
|
|
from app import workspace
|
|
|
|
CHUNK_SIZE = 1000
|
|
CHUNK_OVERLAP = 120
|
|
MAX_CHUNKS_PER_DOCUMENT = 4
|
|
MAX_SEARCH_RESULTS = 8
|
|
MAX_FILES_PER_INDEX = 20
|
|
MAX_TOTAL_CHARS = 2_000_000
|
|
SEARCH_STOP_WORDS = {
|
|
"ابحث", "البحث", "بحث", "في", "عن", "فهرس", "الفهرس", "المعرفة", "المحلي",
|
|
"المحلية", "قاعدة", "المرحلة", "مرحلة", "ثم", "لخص", "استشهد", "مسار", "الملف",
|
|
"أي", "ما", "كيف", "أين", "وين", "خدمة", "الذي", "التي", "إلى", "من", "على",
|
|
"the", "and", "for", "from", "with", "search", "knowledge", "index", "indexed",
|
|
}
|
|
QUERY_SYNONYMS = {
|
|
"تفرغ": ("تفريغ", "التفريغ", "للتفريغ", "transcription"),
|
|
"يفرغ": ("تفريغ", "التفريغ", "للتفريغ", "transcription"),
|
|
"التسجيلات": ("التسجيل", "التسجيلات"),
|
|
"الصوتية": ("الصوت", "الصوتية", "audio"),
|
|
"permission": ("allowed_tools", "tool_name"),
|
|
"permissions": ("allowed_tools", "tool_name"),
|
|
"صلاحية": ("allowed_tools", "tool_name", "مسموح"),
|
|
"صلاحيات": ("allowed_tools", "tool_name", "مسموح"),
|
|
}
|
|
|
|
|
|
@contextmanager
|
|
def _connect() -> Iterator[sqlite3.Connection]:
|
|
connection = sqlite3.connect(database.DATABASE_PATH, timeout=10)
|
|
try:
|
|
with connection:
|
|
yield connection
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
def initialize() -> None:
|
|
with _connect() as connection:
|
|
connection.execute("PRAGMA foreign_keys = ON")
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS knowledge_documents (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id TEXT NOT NULL,
|
|
workspace_path TEXT NOT NULL,
|
|
relative_path TEXT NOT NULL,
|
|
content_hash TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(user_id, workspace_path, relative_path)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS knowledge_chunks (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
document_id INTEGER NOT NULL REFERENCES knowledge_documents(id) ON DELETE CASCADE,
|
|
chunk_index INTEGER NOT NULL,
|
|
body TEXT NOT NULL,
|
|
UNIQUE(document_id, chunk_index)
|
|
);
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS knowledge_fts USING fts5(
|
|
body, tokenize = 'unicode61'
|
|
);
|
|
CREATE TABLE IF NOT EXISTS knowledge_embeddings (
|
|
chunk_id INTEGER PRIMARY KEY REFERENCES knowledge_chunks(id) ON DELETE CASCADE,
|
|
model TEXT NOT NULL,
|
|
dimensions INTEGER NOT NULL,
|
|
vector BLOB NOT NULL
|
|
);
|
|
"""
|
|
)
|
|
|
|
|
|
def _chunks(text: str) -> list[str]:
|
|
cleaned = text.strip()
|
|
if not cleaned:
|
|
return []
|
|
step = CHUNK_SIZE - CHUNK_OVERLAP
|
|
result = [cleaned[start : start + CHUNK_SIZE] for start in range(0, len(cleaned), step)]
|
|
return [chunk for chunk in result if chunk.strip()]
|
|
|
|
|
|
def index_document(
|
|
*, user_id: str, workspace_path: Path, relative_path: str, text: str, content_hash: str
|
|
) -> dict[str, Any]:
|
|
chunks = _chunks(text)
|
|
if not chunks:
|
|
raise ValueError("الملف لا يحتوي على نص قابل للفهرسة.")
|
|
if sum(map(len, chunks)) > MAX_TOTAL_CHARS + CHUNK_SIZE:
|
|
raise ValueError("نص الملف يتجاوز حد الفهرسة.")
|
|
|
|
root = str(workspace_path.resolve(strict=True))
|
|
with _connect() as connection:
|
|
connection.execute("PRAGMA foreign_keys = ON")
|
|
existing = connection.execute(
|
|
"SELECT id FROM knowledge_documents WHERE user_id=? AND workspace_path=? AND relative_path=?",
|
|
(user_id, root, relative_path),
|
|
).fetchone()
|
|
if existing:
|
|
document_id = int(existing[0])
|
|
rows = connection.execute(
|
|
"SELECT id FROM knowledge_chunks WHERE document_id=?", (document_id,)
|
|
).fetchall()
|
|
connection.executemany("DELETE FROM knowledge_fts WHERE rowid=?", rows)
|
|
connection.executemany("DELETE FROM knowledge_embeddings WHERE chunk_id=?", rows)
|
|
connection.execute("DELETE FROM knowledge_chunks WHERE document_id=?", (document_id,))
|
|
connection.execute(
|
|
"UPDATE knowledge_documents SET content_hash=?, updated_at=CURRENT_TIMESTAMP WHERE id=?",
|
|
(content_hash, document_id),
|
|
)
|
|
else:
|
|
cursor = connection.execute(
|
|
"INSERT INTO knowledge_documents(user_id,workspace_path,relative_path,content_hash) VALUES(?,?,?,?)",
|
|
(user_id, root, relative_path, content_hash),
|
|
)
|
|
document_id = int(cursor.lastrowid)
|
|
|
|
for index, body in enumerate(chunks):
|
|
cursor = connection.execute(
|
|
"INSERT INTO knowledge_chunks(document_id,chunk_index,body) VALUES(?,?,?)",
|
|
(document_id, index, body),
|
|
)
|
|
connection.execute(
|
|
"INSERT INTO knowledge_fts(rowid,body) VALUES(?,?)",
|
|
(int(cursor.lastrowid), body),
|
|
)
|
|
return {"path": relative_path, "chunks": len(chunks), "characters": len(text)}
|
|
|
|
|
|
def store_embeddings(
|
|
*, user_id: str, workspace_path: Path, relative_path: str, model: str, vectors: list[list[float]]
|
|
) -> int:
|
|
"""Replace a document's vectors, requiring one finite vector per indexed chunk."""
|
|
with _connect() as connection:
|
|
connection.row_factory = sqlite3.Row
|
|
document = connection.execute(
|
|
"SELECT id FROM knowledge_documents WHERE user_id=? AND workspace_path=? AND relative_path=?",
|
|
(user_id, str(workspace_path.resolve()), relative_path),
|
|
).fetchone()
|
|
if document is None:
|
|
return 0
|
|
chunks = connection.execute(
|
|
"SELECT id FROM knowledge_chunks WHERE document_id=? ORDER BY chunk_index",
|
|
(document["id"],),
|
|
).fetchall()
|
|
if len(chunks) != len(vectors):
|
|
raise ValueError("يجب أن يقابل كل مقطع مفهرس متجه تضمين واحد.")
|
|
prepared: list[tuple[int, str, int, bytes]] = []
|
|
dimensions: int | None = None
|
|
for chunk, vector in zip(chunks, vectors, strict=True):
|
|
if not vector or (dimensions is not None and len(vector) != dimensions):
|
|
raise ValueError("أبعاد متجهات التضمين غير صالحة أو غير متطابقة.")
|
|
dimensions = len(vector)
|
|
numeric = [float(value) for value in vector]
|
|
if not all(math.isfinite(value) for value in numeric) or not any(numeric):
|
|
raise ValueError("متجه التضمين يحتوي قيمة غير صالحة.")
|
|
packed = struct.pack(f"<{dimensions}f", *numeric)
|
|
prepared.append((int(chunk["id"]), model, dimensions, packed))
|
|
connection.execute(
|
|
"DELETE FROM knowledge_embeddings WHERE chunk_id IN "
|
|
"(SELECT id FROM knowledge_chunks WHERE document_id=?)",
|
|
(document["id"],),
|
|
)
|
|
connection.executemany(
|
|
"INSERT INTO knowledge_embeddings(chunk_id,model,dimensions,vector) VALUES(?,?,?,?)",
|
|
prepared,
|
|
)
|
|
return len(prepared)
|
|
|
|
|
|
def has_embeddings(*, user_id: str, workspace_path: Path, model: str) -> bool:
|
|
with _connect() as connection:
|
|
row = connection.execute(
|
|
"""SELECT 1 FROM knowledge_embeddings AS e
|
|
JOIN knowledge_chunks AS c ON c.id=e.chunk_id
|
|
JOIN knowledge_documents AS d ON d.id=c.document_id
|
|
WHERE d.user_id=? AND d.workspace_path=? AND e.model=? LIMIT 1""",
|
|
(user_id, str(workspace_path.resolve()), model),
|
|
).fetchone()
|
|
return row is not None
|
|
|
|
|
|
def search_by_embedding(
|
|
vector: list[float], *, user_id: str, workspace_path: Path, model: str,
|
|
limit: int = MAX_SEARCH_RESULTS,
|
|
) -> list[dict[str, Any]]:
|
|
if not vector or not all(math.isfinite(float(value)) for value in vector):
|
|
return []
|
|
query_norm = math.sqrt(sum(float(value) ** 2 for value in vector))
|
|
if query_norm == 0:
|
|
return []
|
|
with _connect() as connection:
|
|
connection.row_factory = sqlite3.Row
|
|
rows = connection.execute(
|
|
"""SELECT d.id AS document_id, d.relative_path, d.content_hash,
|
|
c.chunk_index, c.body, e.dimensions, e.vector
|
|
FROM knowledge_embeddings AS e
|
|
JOIN knowledge_chunks AS c ON c.id=e.chunk_id
|
|
JOIN knowledge_documents AS d ON d.id=c.document_id
|
|
WHERE d.user_id=? AND d.workspace_path=? AND e.model=?""",
|
|
(user_id, str(workspace_path.resolve()), model),
|
|
).fetchall()
|
|
ranked: list[tuple[float, sqlite3.Row]] = []
|
|
for row in rows:
|
|
dimensions = int(row["dimensions"])
|
|
if dimensions != len(vector) or len(row["vector"]) != dimensions * 4:
|
|
continue
|
|
values = struct.unpack(f"<{dimensions}f", row["vector"])
|
|
norm = math.sqrt(sum(value * value for value in values))
|
|
if norm == 0:
|
|
continue
|
|
similarity = sum(float(left) * right for left, right in zip(vector, values, strict=True)) / (query_norm * norm)
|
|
ranked.append((similarity, row))
|
|
ranked.sort(key=lambda item: (-item[0], item[1]["relative_path"], item[1]["chunk_index"]))
|
|
bounded_limit = max(1, min(limit, MAX_SEARCH_RESULTS))
|
|
selected: list[tuple[float, sqlite3.Row]] = []
|
|
per_document: dict[int, int] = {}
|
|
for score, row in ranked:
|
|
document_id = int(row["document_id"])
|
|
if per_document.get(document_id, 0) >= MAX_CHUNKS_PER_DOCUMENT:
|
|
continue
|
|
per_document[document_id] = per_document.get(document_id, 0) + 1
|
|
selected.append((score, row))
|
|
if len(selected) >= bounded_limit:
|
|
break
|
|
valid: dict[int, bool] = {}
|
|
results: list[dict[str, Any]] = []
|
|
for score, row in selected:
|
|
document_id = int(row["document_id"])
|
|
if document_id not in valid:
|
|
try:
|
|
path = workspace.relative_knowledge_file(workspace_path, str(row["relative_path"]))
|
|
valid[document_id] = hashlib.sha256(path.read_bytes()).hexdigest() == row["content_hash"]
|
|
except (OSError, ValueError):
|
|
valid[document_id] = False
|
|
if not valid[document_id]:
|
|
delete_document(
|
|
user_id=user_id,
|
|
workspace_path=workspace_path,
|
|
relative_path=str(row["relative_path"]),
|
|
)
|
|
if valid[document_id]:
|
|
body = str(row["body"])
|
|
results.append(
|
|
{
|
|
"path": row["relative_path"],
|
|
"chunk": row["chunk_index"],
|
|
"text": body,
|
|
"excerpt": body[:220],
|
|
"similarity": round(score, 4),
|
|
}
|
|
)
|
|
return results
|
|
|
|
|
|
def merge_search_results(
|
|
lexical: list[dict[str, Any]], semantic: list[dict[str, Any]],
|
|
*, limit: int = MAX_SEARCH_RESULTS,
|
|
) -> list[dict[str, Any]]:
|
|
"""Combine lexical and semantic ranks with reciprocal-rank fusion."""
|
|
combined: dict[tuple[str, int], dict[str, Any]] = {}
|
|
scores: dict[tuple[str, int], float] = {}
|
|
for rank, result in enumerate(lexical, 1):
|
|
key = (str(result["path"]), int(result["chunk"]))
|
|
combined[key] = dict(result)
|
|
scores[key] = scores.get(key, 0.0) + 1 / (60 + rank)
|
|
for rank, result in enumerate(semantic, 1):
|
|
key = (str(result["path"]), int(result["chunk"]))
|
|
if key not in combined:
|
|
combined[key] = dict(result)
|
|
else:
|
|
combined[key]["similarity"] = result.get("similarity")
|
|
scores[key] = scores.get(key, 0.0) + 1 / (60 + rank)
|
|
ordered = sorted(combined, key=lambda key: (-scores[key], key[0], key[1]))
|
|
return [combined[key] for key in ordered[: max(1, min(limit, MAX_SEARCH_RESULTS))]]
|
|
|
|
|
|
def delete_document(*, user_id: str, workspace_path: Path, relative_path: str) -> bool:
|
|
root = str(workspace_path.resolve())
|
|
with _connect() as connection:
|
|
connection.execute("PRAGMA foreign_keys = ON")
|
|
row = connection.execute(
|
|
"SELECT id FROM knowledge_documents WHERE user_id=? AND workspace_path=? AND relative_path=?",
|
|
(user_id, root, relative_path),
|
|
).fetchone()
|
|
if row is None:
|
|
return False
|
|
document_id = int(row[0])
|
|
chunks = connection.execute(
|
|
"SELECT id FROM knowledge_chunks WHERE document_id=?", (document_id,)
|
|
).fetchall()
|
|
connection.executemany("DELETE FROM knowledge_fts WHERE rowid=?", chunks)
|
|
connection.executemany("DELETE FROM knowledge_embeddings WHERE chunk_id=?", chunks)
|
|
connection.execute("DELETE FROM knowledge_documents WHERE id=?", (document_id,))
|
|
return True
|
|
|
|
|
|
def search(
|
|
query: str,
|
|
*,
|
|
user_id: str,
|
|
workspace_path: Path,
|
|
limit: int = MAX_SEARCH_RESULTS,
|
|
) -> list[dict[str, Any]]:
|
|
base_terms = [
|
|
term
|
|
for term in dict.fromkeys(re.findall(r"[\w\u0600-\u06ff]{2,}", query.casefold()))
|
|
if term not in SEARCH_STOP_WORDS
|
|
]
|
|
terms = list(
|
|
dict.fromkeys(
|
|
term
|
|
for base in base_terms
|
|
for term in (base, *QUERY_SYNONYMS.get(base, ()))
|
|
)
|
|
)[:20]
|
|
if not terms:
|
|
return []
|
|
match_query = " OR ".join('"' + term.replace('"', '""') + '"' for term in terms)
|
|
bounded_limit = max(1, min(limit, MAX_SEARCH_RESULTS))
|
|
with _connect() as connection:
|
|
connection.row_factory = sqlite3.Row
|
|
rows = connection.execute(
|
|
"""
|
|
SELECT d.id AS document_id, d.relative_path, d.content_hash,
|
|
c.chunk_index, c.body,
|
|
snippet(knowledge_fts, 0, '[', ']', ' … ', 24) AS excerpt,
|
|
bm25(knowledge_fts) AS rank
|
|
FROM knowledge_fts
|
|
JOIN knowledge_chunks AS c ON c.id=knowledge_fts.rowid
|
|
JOIN knowledge_documents AS d ON d.id=c.document_id
|
|
WHERE knowledge_fts MATCH ? AND d.user_id=? AND d.workspace_path=?
|
|
ORDER BY rank, d.relative_path, c.chunk_index
|
|
LIMIT ?
|
|
""",
|
|
(match_query, user_id, str(workspace_path.resolve()), bounded_limit * 8),
|
|
).fetchall()
|
|
selected_rows: list[sqlite3.Row] = []
|
|
chunks_per_document: dict[int, int] = {}
|
|
for row in rows:
|
|
document_id = int(row["document_id"])
|
|
if chunks_per_document.get(document_id, 0) >= MAX_CHUNKS_PER_DOCUMENT:
|
|
continue
|
|
chunks_per_document[document_id] = chunks_per_document.get(document_id, 0) + 1
|
|
selected_rows.append(row)
|
|
if len(selected_rows) >= bounded_limit:
|
|
break
|
|
current_documents: dict[int, bool] = {}
|
|
for row in selected_rows:
|
|
document_id = int(row["document_id"])
|
|
if document_id in current_documents:
|
|
continue
|
|
relative_path = str(row["relative_path"])
|
|
try:
|
|
source = workspace.relative_knowledge_file(workspace_path, relative_path)
|
|
actual_hash = hashlib.sha256(source.read_bytes()).hexdigest()
|
|
current_documents[document_id] = actual_hash == row["content_hash"]
|
|
except (OSError, ValueError):
|
|
current_documents[document_id] = False
|
|
if not current_documents[document_id]:
|
|
delete_document(
|
|
user_id=user_id,
|
|
workspace_path=workspace_path,
|
|
relative_path=relative_path,
|
|
)
|
|
return [
|
|
{"path": row["relative_path"], "chunk": row["chunk_index"], "text": row["body"], "excerpt": row["excerpt"]}
|
|
for row in selected_rows
|
|
if current_documents.get(int(row["document_id"]), False)
|
|
]
|
|
|
|
|
|
initialize()
|