From 86c61ffb0f110a6821729d2c7ec6405d8f8f9bb0 Mon Sep 17 00:00:00 2001 From: arsam Date: Fri, 18 Sep 2026 23:06:30 +0330 Subject: [PATCH] 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 | 161 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 axon/services/index.py diff --git a/axon/services/index.py b/axon/services/index.py new file mode 100644 index 0000000..2ac9e6d --- /dev/null +++ b/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 + ]