Complete local hybrid search and improve agent reliability

2026-10-02 23:38:02 +03:00
parent 3563a104a3
commit 140f6eb287
62 changed files with 7546 additions and 314 deletions
@@ -0,0 +1,372 @@
+"""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 = 2
+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"),
+}
+
+
+@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 = 5
+) -> 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, 5))
+    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 = 5
+) -> 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, 5))]]
+
+
+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 = 5) -> 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, 5))
+    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()