Import source: axon/services/index

Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.

axon/services/index.py

@@ -0,0 +1,161 @@
+import asyncio
+import json
+import sqlite3
+import uuid
+from pathlib import Path
+
+from axon.config import settings
+from axon.models import IndexDocument, IngestDocument, SearchResult
+
+_lock = asyncio.Lock()
+
+
+def _db_path() -> Path:
+    return Path(settings.index_db_path)
+
+
+def _connect() -> sqlite3.Connection:
+    conn = sqlite3.connect(_db_path())
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+async def init_index() -> None:
+    def _init() -> None:
+        with _connect() as conn:
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS documents (
+                    id TEXT PRIMARY KEY,
+                    title TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    url TEXT,
+                    metadata TEXT NOT NULL DEFAULT '{}',
+                    updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
+                )
+                """
+            )
+            conn.execute(
+                """
+                CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
+                    title,
+                    content,
+                    content='documents',
+                    content_rowid='rowid'
+                )
+                """
+            )
+            conn.execute(
+                """
+                CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN
+                    INSERT INTO documents_fts(rowid, title, content)
+                    VALUES (new.rowid, new.title, new.content);
+                END
+                """
+            )
+            conn.execute(
+                """
+                CREATE TRIGGER IF NOT EXISTS documents_ad AFTER DELETE ON documents BEGIN
+                    INSERT INTO documents_fts(documents_fts, rowid, title, content)
+                    VALUES ('delete', old.rowid, old.title, old.content);
+                END
+                """
+            )
+            conn.execute(
+                """
+                CREATE TRIGGER IF NOT EXISTS documents_au AFTER UPDATE ON documents BEGIN
+                    INSERT INTO documents_fts(documents_fts, rowid, title, content)
+                    VALUES ('delete', old.rowid, old.title, old.content);
+                    INSERT INTO documents_fts(rowid, title, content)
+                    VALUES (new.rowid, new.title, new.content);
+                END
+                """
+            )
+
+    async with _lock:
+        await asyncio.to_thread(_init)
+
+
+async def ingest_documents(documents: list[IngestDocument]) -> int:
+    await init_index()
+
+    def _ingest() -> int:
+        with _connect() as conn:
+            for doc in documents:
+                doc_id = doc.id or str(uuid.uuid4())
+                conn.execute(
+                    """
+                    INSERT INTO documents(id, title, content, url, metadata, updated_at)
+                    VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
+                    ON CONFLICT(id) DO UPDATE SET
+                        title=excluded.title,
+                        content=excluded.content,
+                        url=excluded.url,
+                        metadata=excluded.metadata,
+                        updated_at=CURRENT_TIMESTAMP
+                    """,
+                    (
+                        doc_id,
+                        doc.title,
+                        doc.content,
+                        doc.url,
+                        json.dumps(doc.metadata, sort_keys=True),
+                    ),
+                )
+            return len(documents)
+
+    async with _lock:
+        return await asyncio.to_thread(_ingest)
+
+
+async def search_index(query: str, max_results: int = 10) -> list[IndexDocument]:
+    await init_index()
+
+    def _search() -> list[IndexDocument]:
+        with _connect() as conn:
+            rows = conn.execute(
+                """
+                SELECT
+                    documents.id,
+                    documents.title,
+                    snippet(documents_fts, 1, '', '', '...', 40) AS content,
+                    documents.url,
+                    bm25(documents_fts) AS rank
+                FROM documents_fts
+                JOIN documents ON documents_fts.rowid = documents.rowid
+                WHERE documents_fts MATCH ?
+                ORDER BY rank
+                LIMIT ?
+                """,
+                (query, max_results),
+            ).fetchall()
+        results = []
+        for row in rows:
+            rank = float(row["rank"])
+            score = max(0.0, min(1.0, 1.0 / (1.0 + abs(rank))))
+            results.append(
+                IndexDocument(
+                    id=row["id"],
+                    title=row["title"],
+                    content=row["content"] or "",
+                    url=row["url"],
+                    score=round(score, 6),
+                )
+            )
+        return results
+
+    async with _lock:
+        return await asyncio.to_thread(_search)
+
+
+async def search_index_as_results(query: str, max_results: int = 5) -> list[SearchResult]:
+    docs = await search_index(query, max_results=max_results)
+    return [
+        SearchResult(
+            title=doc.title,
+            url=doc.url or f"private://{doc.id}",
+            content=doc.content,
+            score=doc.score,
+        )
+        for doc in docs
+    ]