b60439152b
- harness/es.py: append-only event store + проекции (образец — кооп-портал) - harness/levels.py: уровни 1с/1м/1ч/1д, контекст ∝ масштабу - harness/scheduler.py: уровневые потоки, идемпотентные тики, absorbed-стабы - harness/isolated_runner.py + _llm_worker.py: LLM-тик в сабпроцессе (глобальный asyncio-рантайм Strands не переживает 2+ create_harness) - cli.py: init/task/tasks/status/events + requeue (написан самим минутным агентом) - 9router: отдельный ключ sk-simbiot, модель auto (glm-5.3 первым)
341 lines
10 KiB
Python
341 lines
10 KiB
Python
"""BM25 full-text search strategy powered by SQLite FTS5.
|
|
|
|
Maintains a SQLite-backed inverted index over storage contents and uses BM25
|
|
scoring for relevance ranking. Accounts for term frequency, inverse document
|
|
frequency, and document length normalization.
|
|
|
|
Indexes entries at write time via :meth:`Bm25SearchStrategy.index` so searches
|
|
do not need to re-read storage contents. Requires a ``base_dir`` property on
|
|
the storage instance only to locate the SQLite database file on the host
|
|
filesystem.
|
|
|
|
Zero external dependencies -- uses Python's stdlib ``sqlite3`` module which
|
|
ships FTS5 support on all modern platforms.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import os
|
|
import sqlite3
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Any, Literal
|
|
|
|
from ..storage import StorageSearchResult
|
|
from .keyword import tokenize
|
|
|
|
if TYPE_CHECKING:
|
|
from ..local_file_storage import LocalFileStorage
|
|
|
|
STOP_WORDS: frozenset[str] = frozenset(
|
|
{
|
|
"a",
|
|
"about",
|
|
"above",
|
|
"after",
|
|
"again",
|
|
"all",
|
|
"although",
|
|
"an",
|
|
"and",
|
|
"any",
|
|
"are",
|
|
"as",
|
|
"at",
|
|
"be",
|
|
"because",
|
|
"been",
|
|
"before",
|
|
"being",
|
|
"below",
|
|
"between",
|
|
"both",
|
|
"but",
|
|
"by",
|
|
"can",
|
|
"could",
|
|
"dare",
|
|
"did",
|
|
"do",
|
|
"does",
|
|
"during",
|
|
"each",
|
|
"either",
|
|
"every",
|
|
"few",
|
|
"for",
|
|
"from",
|
|
"further",
|
|
"had",
|
|
"has",
|
|
"have",
|
|
"he",
|
|
"her",
|
|
"him",
|
|
"his",
|
|
"how",
|
|
"i",
|
|
"if",
|
|
"in",
|
|
"into",
|
|
"is",
|
|
"it",
|
|
"its",
|
|
"just",
|
|
"may",
|
|
"me",
|
|
"might",
|
|
"more",
|
|
"most",
|
|
"my",
|
|
"need",
|
|
"neither",
|
|
"no",
|
|
"nor",
|
|
"not",
|
|
"of",
|
|
"off",
|
|
"on",
|
|
"once",
|
|
"only",
|
|
"or",
|
|
"other",
|
|
"ought",
|
|
"our",
|
|
"out",
|
|
"over",
|
|
"own",
|
|
"same",
|
|
"shall",
|
|
"she",
|
|
"should",
|
|
"since",
|
|
"so",
|
|
"some",
|
|
"such",
|
|
"than",
|
|
"that",
|
|
"the",
|
|
"their",
|
|
"them",
|
|
"then",
|
|
"these",
|
|
"they",
|
|
"this",
|
|
"those",
|
|
"though",
|
|
"through",
|
|
"to",
|
|
"too",
|
|
"under",
|
|
"unless",
|
|
"until",
|
|
"up",
|
|
"us",
|
|
"used",
|
|
"very",
|
|
"was",
|
|
"we",
|
|
"were",
|
|
"what",
|
|
"when",
|
|
"where",
|
|
"which",
|
|
"while",
|
|
"who",
|
|
"whom",
|
|
"why",
|
|
"will",
|
|
"with",
|
|
"would",
|
|
"yet",
|
|
"you",
|
|
"your",
|
|
}
|
|
)
|
|
|
|
|
|
def _build_query(query: str) -> str | None:
|
|
"""Build an FTS5 match expression from a natural-language query.
|
|
|
|
Filters stop words and single-character tokens, appends a prefix
|
|
wildcard to each term, then joins with implicit AND so all terms
|
|
must be present as prefixes (ranked by BM25).
|
|
|
|
Args:
|
|
query: Natural-language search query.
|
|
|
|
Returns:
|
|
FTS5 match expression, or None if no meaningful terms remain.
|
|
"""
|
|
terms = sorted(term for term in tokenize(query) if len(term) > 1 and term not in STOP_WORDS)
|
|
if not terms:
|
|
return None
|
|
return " ".join(f"{term}*" for term in terms)
|
|
|
|
|
|
@dataclass
|
|
class Bm25SearchStrategyConfig:
|
|
"""Configuration for :class:`Bm25SearchStrategy`.
|
|
|
|
Attributes:
|
|
db_path: Path to the SQLite database file for the FTS5 index.
|
|
Defaults to ``.<dir>-fts5.sqlite`` alongside the storage directory.
|
|
"""
|
|
|
|
db_path: str | None = field(default=None)
|
|
|
|
|
|
class Bm25SearchStrategy:
|
|
"""BM25 full-text search strategy powered by SQLite FTS5.
|
|
|
|
Maintains a SQLite-backed inverted index over the storage contents and uses
|
|
BM25 scoring for relevance ranking. Accounts for term frequency, inverse
|
|
document frequency, and document length normalization.
|
|
|
|
Indexes entries at write time via :meth:`index` -- consumers call
|
|
``strategy.index(storage, key, data)`` on each write, and ``search()``
|
|
queries the pre-built index. Only entries passed through ``index()`` are
|
|
searchable; pre-existing storage contents are not backfilled automatically.
|
|
|
|
The FTS5 index is persisted on the host filesystem via the storage's
|
|
``base_dir`` and is not sandbox-aware.
|
|
|
|
Example:
|
|
```python
|
|
from strands.storage import LocalFileStorage
|
|
from strands.storage.search.bm25 import Bm25SearchStrategy
|
|
|
|
storage = LocalFileStorage("./memory/")
|
|
strategy = Bm25SearchStrategy()
|
|
|
|
await strategy.index(storage, "auth.md", b"OAuth2 authentication flow")
|
|
results = await strategy.search(storage, "authentication flow")
|
|
await strategy.close()
|
|
```
|
|
"""
|
|
|
|
requires_host_fs: Literal[True] = True
|
|
|
|
def __init__(self, config: Bm25SearchStrategyConfig | None = None) -> None:
|
|
"""Initialize the BM25 search strategy.
|
|
|
|
Args:
|
|
config: Optional configuration. See :class:`Bm25SearchStrategyConfig`.
|
|
"""
|
|
self._config = config or Bm25SearchStrategyConfig()
|
|
self._conn: sqlite3.Connection | None = None
|
|
self._storage_path: str | None = None
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def index(self, storage: LocalFileStorage, key: str, data: bytes, **kwargs: Any) -> None:
|
|
"""Index a single entry for future searches.
|
|
|
|
Skips hidden files (keys whose final segment starts with '.').
|
|
Uses content hashing to avoid redundant re-indexing when the
|
|
data has not changed.
|
|
|
|
Args:
|
|
storage: The storage backend (used to locate the index db).
|
|
key: The storage key being written.
|
|
data: The raw bytes being stored.
|
|
**kwargs: Unused; accepted for protocol compatibility.
|
|
"""
|
|
if key.rsplit("/", maxsplit=1)[-1].startswith("."):
|
|
return
|
|
|
|
content = data.decode("utf-8", errors="replace")
|
|
content_hash = hashlib.md5(data, usedforsecurity=False).hexdigest()
|
|
|
|
async with self._lock:
|
|
conn = self._ensure_connection(storage)
|
|
await asyncio.shield(asyncio.to_thread(self._upsert, conn, key, content, content_hash))
|
|
|
|
async def search(self, storage: LocalFileStorage, query: str, **kwargs: Any) -> list[StorageSearchResult]:
|
|
"""Search the index using BM25 full-text search.
|
|
|
|
Args:
|
|
storage: The storage backend (used to locate the index db).
|
|
query: Natural-language search query.
|
|
**kwargs: Unused; accepted for protocol compatibility.
|
|
|
|
Returns:
|
|
Matched keys with BM25 relevance scores, ranked best-first.
|
|
|
|
Raises:
|
|
RuntimeError: If the SQLite build lacks FTS5 support.
|
|
"""
|
|
fts_query = _build_query(query)
|
|
if not fts_query:
|
|
return []
|
|
|
|
async with self._lock:
|
|
conn = self._ensure_connection(storage)
|
|
return await asyncio.shield(asyncio.to_thread(self._query, conn, fts_query))
|
|
|
|
async def close(self) -> None:
|
|
"""Close the SQLite connection and release resources."""
|
|
if self._conn is not None:
|
|
self._conn.close()
|
|
self._conn = None
|
|
self._storage_path = None
|
|
|
|
def _ensure_connection(self, storage: LocalFileStorage) -> sqlite3.Connection:
|
|
"""Lazily initialize the SQLite connection from the storage backend."""
|
|
storage_path = self._resolve_storage_path(storage)
|
|
|
|
if self._conn is not None and self._storage_path == storage_path:
|
|
return self._conn
|
|
|
|
if self._conn is not None:
|
|
self._conn.close()
|
|
|
|
os.makedirs(storage_path, exist_ok=True)
|
|
db_path = self._config.db_path
|
|
if db_path is None:
|
|
parent = os.path.dirname(os.path.abspath(storage_path))
|
|
basename = os.path.basename(os.path.abspath(storage_path))
|
|
db_path = os.path.join(parent, f".__strands_{basename}-fts5.sqlite")
|
|
|
|
conn = sqlite3.connect(db_path, check_same_thread=False)
|
|
try:
|
|
conn.execute("CREATE VIRTUAL TABLE IF NOT EXISTS documents USING fts5(key, content)")
|
|
except sqlite3.OperationalError as error:
|
|
conn.close()
|
|
raise RuntimeError("SQLite build lacks FTS5 support") from error
|
|
conn.execute("CREATE TABLE IF NOT EXISTS doc_hashes (key TEXT PRIMARY KEY, hash TEXT NOT NULL)")
|
|
conn.commit()
|
|
|
|
self._conn = conn
|
|
self._storage_path = storage_path
|
|
return conn
|
|
|
|
def _upsert(self, conn: sqlite3.Connection, key: str, content: str, content_hash: str) -> None:
|
|
"""Insert or replace a single document in the FTS5 index."""
|
|
row = conn.execute("SELECT hash FROM doc_hashes WHERE key = ?", (key,)).fetchone()
|
|
if row and row[0] == content_hash:
|
|
return
|
|
conn.execute("DELETE FROM documents WHERE key = ?", (key,))
|
|
conn.execute("INSERT INTO documents (key, content) VALUES (?, ?)", (key, content))
|
|
conn.execute("INSERT OR REPLACE INTO doc_hashes (key, hash) VALUES (?, ?)", (key, content_hash))
|
|
conn.commit()
|
|
|
|
@staticmethod
|
|
def _query(conn: sqlite3.Connection, fts_query: str) -> list[StorageSearchResult]:
|
|
"""Execute the FTS5 query and return scored results."""
|
|
cursor = conn.execute(
|
|
"SELECT key, rank FROM documents WHERE documents MATCH ? ORDER BY rank",
|
|
(fts_query,),
|
|
)
|
|
results: list[StorageSearchResult] = []
|
|
for row in cursor.fetchall():
|
|
raw_score = abs(row[1])
|
|
score = raw_score / (1.0 + raw_score)
|
|
results.append(StorageSearchResult(key=row[0], score=score))
|
|
return results
|
|
|
|
@staticmethod
|
|
def _resolve_storage_path(storage: LocalFileStorage) -> str:
|
|
"""Extract base_dir from the storage instance."""
|
|
return storage.base_dir
|