src/knowledge/sqlite_store.py - master - Ilmer/tandemBotLLM

src/knowledge/sqlite_store.py
88 строк3 KB
vkuzin
base proect
14 май 2026, 07:36
14 май 2026, 07:360893926
from __future__ import annotations

import re
from datetime import datetime, timezone
from typing import Any

import aiosqlite

from persistence.repositories import Database


def _iso(dt: datetime) -> str:
    return dt.astimezone(timezone.utc).isoformat()


def _fts_terms(q: str) -> str:
    parts = re.findall(r"[\w\-]+", q, flags=re.UNICODE)[:10]
    if not parts:
        return ""
    return " AND ".join(f'"{p}"' for p in parts if len(p) <= 80)


class SqliteKnowledgeStore:
    """SQLite + FTS5 (`knowledge_docs` + `knowledge_fts`)."""

    def __init__(self, db: Database) -> None:
        self._db = db

    async def ingest(
        self,
        *,
        kind: str,
        title: str | None,
        body: str,
        source_path: str | None = None,
    ) -> int:
        now = _iso(datetime.now(timezone.utc))
        async with await self._db.connect() as conn:
            cur = await conn.execute(
                """
                INSERT INTO knowledge_docs (kind, title, body, source_path, created_at)
                VALUES (?, ?, ?, ?, ?)
                """,
                (kind, title, body, source_path, now),
            )
            doc_id = int(cur.lastrowid)
            await conn.execute(
                """
                INSERT INTO knowledge_fts (doc_id, title, body)
                VALUES (?, ?, ?)
                """,
                (doc_id, title or "", body),
            )
            await conn.commit()
        return doc_id

    async def search(self, query: str, *, limit: int = 5) -> list[dict[str, Any]]:
        fts = _fts_terms(query.strip())
        if not fts:
            return []
        async with await self._db.connect() as conn:
            try:
                cur = await conn.execute(
                    """
                    SELECT d.id AS id, d.title AS title, d.body AS body
                    FROM knowledge_fts AS f
                    JOIN knowledge_docs AS d ON d.id = f.doc_id
                    WHERE f MATCH ?
                    ORDER BY bm25(f)
                    LIMIT ?
                    """,
                    (fts, limit),
                )
                rows = await cur.fetchall()
            except aiosqlite.OperationalError:
                return []
        out: list[dict[str, Any]] = []
        for r in rows:
            body = r["body"] or ""
            out.append(
                {
                    "id": r["id"],
                    "title": r["title"],
                    "excerpt": body[:800],
                    "body": body,
                }
            )
        return out