Overhaul docs retrieval and web search quality

Replace the abandoned llms-txt-mcp/Chroma docs backend with an in-repo
MCP service: SQLite WAL + FTS5 + sentence-transformer embeddings,
transactional source replacement, persisted state across restarts,
singleflight refresh with conditional requests, hybrid lexical/semantic
ranking with exact-duplicate collapse, source/host filters, and
explicit-by-default content retrieval. Add docs_rebuild and a
docs-rebuild CLI command.

Add deterministic llms-full.txt snapshot generation for machine-local
menus with hash-validated provenance manifests; lifecycle commands
promote a local menu to its snapshot only when the manifest validates.
Switch public source profiles to content-bearing llms-full.txt feeds.

Improve web search: bounded provider fallback with per-attempt
diagnostics and cancellation, an optional Brave Search API provider,
strict SearXNG engine selection, capped link/media extraction, and a
real engine=browser renderer that routes every request through the
existing SSRF vetting while blocking WebSockets, non-GET traffic, and
private destinations.

Extend release checks with offline unit suites and isolated candidate
container tests for both images.
This commit is contained in:
2026-07-25 08:49:26 -07:00
parent 29bcb123fa
commit 51dceee224
60 changed files with 3207 additions and 107 deletions

View File

@@ -0,0 +1,3 @@
"""Maintained Context Kit documentation retrieval service."""
__version__ = "1.0.0"

View File

@@ -0,0 +1,4 @@
from .server import main
main()

View File

@@ -0,0 +1,75 @@
from __future__ import annotations
import asyncio
import hashlib
import numpy as np
class SentenceTransformerEmbedder:
def __init__(self, model_name: str):
self.model_name = model_name
self.fingerprint = f"sentence-transformers:{model_name}"
self._model = None
self._lock = asyncio.Lock()
@property
def ready(self) -> bool:
return self._model is not None
async def ensure_ready(self) -> None:
if self._model is not None:
return
async with self._lock:
if self._model is None:
self._model = await asyncio.to_thread(self._load)
def _load(self):
from sentence_transformers import SentenceTransformer
return SentenceTransformer(self.model_name, device="cpu")
async def encode_documents(self, texts: list[str]) -> np.ndarray:
await self.ensure_ready()
return await asyncio.to_thread(self._encode, texts)
async def encode_query(self, text: str) -> np.ndarray:
vectors = await self.encode_documents([text])
return vectors[0]
def _encode(self, texts: list[str]) -> np.ndarray:
return np.asarray(
self._model.encode(
texts,
batch_size=32,
show_progress_bar=False,
normalize_embeddings=True,
convert_to_numpy=True,
),
dtype=np.float32,
)
class LexicalFallbackEmbedder:
"""Deterministic fallback used only when a model cannot be loaded."""
fingerprint = "lexical-fallback-v1"
ready = True
async def ensure_ready(self) -> None:
return None
async def encode_documents(self, texts: list[str]) -> np.ndarray:
return np.asarray([self._encode(text) for text in texts], dtype=np.float32)
async def encode_query(self, text: str) -> np.ndarray:
return np.asarray(self._encode(text), dtype=np.float32)
@staticmethod
def _encode(text: str, dimensions: int = 384) -> np.ndarray:
vector = np.zeros(dimensions, dtype=np.float32)
for token in text.lower().split():
digest = hashlib.sha256(token.encode()).digest()
vector[int.from_bytes(digest[:4], "big") % dimensions] += 1.0
norm = np.linalg.norm(vector)
return vector / norm if norm else vector

View File

@@ -0,0 +1,51 @@
from __future__ import annotations
import httpx
from .models import FetchResponse, SourceState
class SourceFetcher:
def __init__(self, timeout_seconds: float = 30, max_bytes: int = 20_000_000):
self.timeout_seconds = timeout_seconds
self.max_bytes = max_bytes
self._client = httpx.AsyncClient(
follow_redirects=True,
timeout=httpx.Timeout(timeout_seconds),
headers={"User-Agent": "context-kit-docs/1.0"},
)
async def close(self) -> None:
await self._client.aclose()
async def fetch(self, source_url: str, state: SourceState | None = None) -> FetchResponse:
headers: dict[str, str] = {}
if state and state.etag:
headers["If-None-Match"] = state.etag
if state and state.last_modified:
headers["If-Modified-Since"] = state.last_modified
async with self._client.stream("GET", source_url, headers=headers) as response:
if response.status_code == 304:
return FetchResponse(
status=304,
requested_url=source_url,
resolved_url=str(response.url),
etag=response.headers.get("etag"),
last_modified=response.headers.get("last-modified"),
)
chunks: list[bytes] = []
size = 0
async for chunk in response.aiter_bytes():
size += len(chunk)
if size > self.max_bytes:
raise RuntimeError(f"source exceeds {self.max_bytes} byte limit")
chunks.append(chunk)
body = b"".join(chunks).decode(response.encoding or "utf-8", errors="replace")
return FetchResponse(
status=response.status_code,
requested_url=source_url,
resolved_url=str(response.url),
body=body,
etag=response.headers.get("etag"),
last_modified=response.headers.get("last-modified"),
)

View File

@@ -0,0 +1,122 @@
from __future__ import annotations
from dataclasses import dataclass, field
import numpy as np
@dataclass(frozen=True)
class ParsedDocument:
title: str
description: str
content: str
canonical_url: str
heading_path: str
chunk_index: int = 0
@dataclass(frozen=True)
class ParsedSource:
format: str
documents: list[ParsedDocument]
@dataclass(frozen=True)
class FetchResponse:
status: int
requested_url: str
resolved_url: str
body: str = ""
etag: str | None = None
last_modified: str | None = None
@dataclass(frozen=True)
class PreparedDocument:
id: str
configured_source: str
resolved_source: str
source_host: str
canonical_url: str
canonical_host: str
title: str
description: str
heading_path: str
content: str
content_hash: str
embedding: np.ndarray
@dataclass(frozen=True)
class SourceUpdate:
configured_source: str
resolved_source: str
etag: str | None
last_modified: str | None
body_hash: str
raw_body: str
parser_fingerprint: str
embedding_fingerprint: str
checked_at: float
indexed_at: float
stale_at: float
documents: list[PreparedDocument]
@dataclass(frozen=True)
class SourceState:
configured_source: str
resolved_source: str | None
active: bool
etag: str | None
last_modified: str | None
body_hash: str | None
raw_body: str | None
parser_fingerprint: str | None
embedding_fingerprint: str | None
checked_at: float | None
indexed_at: float | None
stale_at: float | None
last_error: str | None
doc_count: int
@dataclass(frozen=True)
class StoredDocument:
id: str
configured_source: str
resolved_source: str
source_host: str
canonical_url: str
canonical_host: str
title: str
description: str
heading_path: str
content: str
content_hash: str
embedding: np.ndarray
@dataclass(frozen=True)
class SearchResult:
id: str
configured_source: str
canonical_url: str
title: str
description: str
heading_path: str
content: str
content_hash: str
score: float
lexical_rank: int | None
semantic_rank: int | None
duplicate_count: int = 1
alternate_sources: list[dict[str, str]] = field(default_factory=list)
@dataclass(frozen=True)
class RefreshOutcome:
source: str
status: str
document_count: int
detail: str | None = None

View File

@@ -0,0 +1,154 @@
from __future__ import annotations
import re
from urllib.parse import urljoin
import yaml
from .models import ParsedDocument, ParsedSource
PARSER_FINGERPRINT = "context-docs-parser-v1"
_MENU_LINK = re.compile(r"^\s*[-*]\s+\[([^]]+)]\(([^)]+)\)(?::\s*(.*))?\s*$")
_HEADING = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
_FRONTMATTER = re.compile(r"(?m)^---\s*$")
def parse_llms_text(content: str, source_url: str, max_chunk_chars: int = 6_000) -> ParsedSource:
normalized = content.replace("\r\n", "\n").replace("\r", "\n").strip()
if not normalized:
return ParsedSource("empty", [])
yaml_documents = _parse_repeated_frontmatter(normalized, source_url, max_chunk_chars)
if yaml_documents is not None:
return ParsedSource("yaml-full", yaml_documents)
if source_url.split("?", 1)[0].endswith("/llms.txt"):
menu_documents = _parse_menu(normalized, source_url)
if menu_documents:
return ParsedSource("standard-menu", menu_documents)
return ParsedSource("markdown-full", _parse_markdown_bundle(normalized, source_url, max_chunk_chars))
def _parse_repeated_frontmatter(content: str, source_url: str, max_chunk_chars: int) -> list[ParsedDocument] | None:
if not content.startswith("---\n"):
return None
separators = [match.start() for match in _FRONTMATTER.finditer(content)]
if len(separators) < 2:
return None
documents: list[ParsedDocument] = []
cursor = 0
while cursor < len(content):
if not content.startswith("---", cursor):
return None
header_end = content.find("\n---", cursor + 3)
if header_end < 0:
return None
try:
metadata = yaml.safe_load(content[cursor + 3 : header_end]) or {}
except yaml.YAMLError:
return None
if not isinstance(metadata, dict) or not isinstance(metadata.get("title"), str):
return None
body_start = header_end + 4
if body_start < len(content) and content[body_start] == "\n":
body_start += 1
next_header = content.find("\n---\n", body_start)
body_end = len(content) if next_header < 0 else next_header
body = content[body_start:body_end].strip()
title = metadata["title"].strip()
description = str(metadata.get("description") or "").strip()
canonical = str(metadata.get("url") or metadata.get("canonical_url") or source_url)
documents.extend(_chunk_document(title, description, body, urljoin(source_url, canonical), title, max_chunk_chars))
if next_header < 0:
break
cursor = next_header + 1
return documents or None
def _parse_menu(content: str, source_url: str) -> list[ParsedDocument]:
documents: list[ParsedDocument] = []
headings: list[tuple[int, str]] = []
in_fence = False
for line in content.splitlines():
if line.lstrip().startswith(("```", "~~~")):
in_fence = not in_fence
continue
if in_fence:
continue
heading = _HEADING.match(line)
if heading:
level = len(heading.group(1))
headings = [entry for entry in headings if entry[0] < level]
headings.append((level, heading.group(2).strip()))
continue
link = _MENU_LINK.match(line)
if not link:
continue
title, target, description = link.group(1).strip(), link.group(2).strip(), (link.group(3) or "").strip()
canonical = urljoin(source_url, target)
path = " > ".join([name for _, name in headings] + [title])
rendered = f"{title}\n\n{description}\n\nSource: {canonical}".strip()
documents.append(ParsedDocument(title, description, rendered, canonical, path))
return documents
def _parse_markdown_bundle(content: str, source_url: str, max_chunk_chars: int) -> list[ParsedDocument]:
sections: list[tuple[str, str]] = []
current_title = "Documentation"
current_lines: list[str] = []
in_fence = False
for line in content.splitlines():
if line.lstrip().startswith(("```", "~~~")):
in_fence = not in_fence
heading = None if in_fence else _HEADING.match(line)
if heading and len(heading.group(1)) == 1:
if current_lines or sections:
sections.append((current_title, "\n".join(current_lines).strip()))
current_title = heading.group(2).strip()
current_lines = []
else:
current_lines.append(line)
if current_lines or not sections:
sections.append((current_title, "\n".join(current_lines).strip()))
documents: list[ParsedDocument] = []
for title, body in sections:
if not body and title == "Documentation":
continue
documents.extend(_chunk_document(title, "", body, source_url, title, max_chunk_chars))
return documents
def _chunk_document(
title: str,
description: str,
content: str,
canonical_url: str,
heading_path: str,
max_chunk_chars: int,
) -> list[ParsedDocument]:
if len(content) <= max_chunk_chars:
return [ParsedDocument(title, description, content, canonical_url, heading_path, 0)]
paragraphs = re.split(r"\n{2,}", content)
chunks: list[str] = []
current: list[str] = []
size = 0
for paragraph in paragraphs:
pieces = [paragraph[index : index + max_chunk_chars] for index in range(0, len(paragraph), max_chunk_chars)] or [""]
for piece in pieces:
added = len(piece) + (2 if current else 0)
if current and size + added > max_chunk_chars:
chunks.append("\n\n".join(current))
current, size = [], 0
current.append(piece)
size += len(piece) + (2 if len(current) > 1 else 0)
if current:
chunks.append("\n\n".join(current))
return [
ParsedDocument(title, description, chunk, canonical_url, heading_path, index)
for index, chunk in enumerate(chunks)
]

View File

@@ -0,0 +1,109 @@
from __future__ import annotations
import asyncio
import hashlib
from urllib.parse import urlparse
from .models import PreparedDocument, RefreshOutcome, SourceUpdate
from .parser import PARSER_FINGERPRINT
class RefreshCoordinator:
def __init__(self, store, fetcher, embedder, parser, ttl_seconds: float, now):
self.store = store
self.fetcher = fetcher
self.embedder = embedder
self.parser = parser
self.ttl_seconds = ttl_seconds
self.now = now
self._inflight: dict[str, asyncio.Task] = {}
self._lock = asyncio.Lock()
async def refresh(self, source: str, force: bool = False) -> RefreshOutcome:
state = self.store.get_source(source)
timestamp = self.now()
compatible = bool(
state
and state.parser_fingerprint == PARSER_FINGERPRINT
and state.embedding_fingerprint == self.embedder.fingerprint
)
if not force and compatible and state.doc_count and state.stale_at and state.stale_at > timestamp:
return RefreshOutcome(source, "fresh", state.doc_count)
async with self._lock:
task = self._inflight.get(source)
if task is None:
task = asyncio.create_task(self._refresh_once(source, timestamp))
self._inflight[source] = task
try:
return await task
finally:
async with self._lock:
if self._inflight.get(source) is task and task.done():
self._inflight.pop(source, None)
async def _refresh_once(self, source: str, timestamp: float) -> RefreshOutcome:
state = self.store.get_source(source)
try:
compatible = bool(
state
and state.parser_fingerprint == PARSER_FINGERPRINT
and state.embedding_fingerprint == self.embedder.fingerprint
)
response = await self.fetcher.fetch(source, state if compatible else None)
if response.status == 304:
count = state.doc_count if state else 0
self.store.mark_checked(source, timestamp, timestamp + self.ttl_seconds)
return RefreshOutcome(source, "not_modified", count)
if response.status != 200:
raise RuntimeError(f"source returned HTTP {response.status}")
parsed = self.parser(response.body, response.resolved_url)
if not parsed.documents:
raise RuntimeError("source parsed to zero documents; previous generation preserved")
texts = ["\n\n".join(filter(None, [doc.title, doc.description, doc.heading_path, doc.content])) for doc in parsed.documents]
vectors = await self.embedder.encode_documents(texts) if texts else []
documents: list[PreparedDocument] = []
source_host = (urlparse(response.resolved_url).hostname or "").lower()
for parsed_document, vector in zip(parsed.documents, vectors, strict=True):
content_hash = hashlib.sha256(parsed_document.content.encode()).hexdigest()
identity = "\0".join(
[source, parsed_document.canonical_url, parsed_document.heading_path, str(parsed_document.chunk_index)]
)
documents.append(
PreparedDocument(
id=hashlib.sha256(identity.encode()).hexdigest()[:24],
configured_source=source,
resolved_source=response.resolved_url,
source_host=source_host,
canonical_url=parsed_document.canonical_url,
canonical_host=(urlparse(parsed_document.canonical_url).hostname or source_host).lower(),
title=parsed_document.title,
description=parsed_document.description,
heading_path=parsed_document.heading_path,
content=parsed_document.content,
content_hash=content_hash,
embedding=vector,
)
)
body_hash = hashlib.sha256(response.body.encode()).hexdigest()
self.store.replace_source(
SourceUpdate(
configured_source=source,
resolved_source=response.resolved_url,
etag=response.etag,
last_modified=response.last_modified,
body_hash=body_hash,
raw_body=response.body,
parser_fingerprint=PARSER_FINGERPRINT,
embedding_fingerprint=self.embedder.fingerprint,
checked_at=timestamp,
indexed_at=timestamp,
stale_at=timestamp + self.ttl_seconds,
documents=documents,
)
)
return RefreshOutcome(source, "updated", len(documents), parsed.format)
except Exception as error:
self.store.mark_checked(source, timestamp, timestamp, str(error))
return RefreshOutcome(source, "error", state.doc_count if state else 0, str(error))

View File

@@ -0,0 +1,82 @@
from __future__ import annotations
from collections import defaultdict
import numpy as np
from .models import SearchResult, StoredDocument
from .store import IndexStore
class HybridSearch:
def __init__(self, store: IndexStore, embedder, rrf_k: int = 60):
self.store = store
self.embedder = embedder
self.rrf_k = rrf_k
async def search(
self,
query: str,
limit: int = 10,
sources: list[str] | None = None,
hosts: list[str] | None = None,
) -> list[SearchResult]:
pool_size = max(limit * 8, 40)
lexical = self.store.lexical_search(query, pool_size, sources, hosts)
candidates = self.store.semantic_candidates(sources, hosts)
semantic: list[StoredDocument] = []
if candidates:
query_vector = np.asarray(await self.embedder.encode_query(query), dtype=np.float32)
query_norm = np.linalg.norm(query_vector)
scored: list[tuple[float, StoredDocument]] = []
for document in candidates:
norm = np.linalg.norm(document.embedding) * query_norm
score = float(np.dot(document.embedding, query_vector) / norm) if norm else 0.0
scored.append((score, document))
semantic = [document for _, document in sorted(scored, key=lambda item: (-item[0], item[1].id))[:pool_size]]
lexical_ranks = {document.id: rank for rank, document in enumerate(lexical, 1)}
semantic_ranks = {document.id: rank for rank, document in enumerate(semantic, 1)}
documents = {document.id: document for document in [*lexical, *semantic]}
scores = defaultdict(float)
for identifier, rank in lexical_ranks.items():
scores[identifier] += 1.0 / (self.rrf_k + rank)
for identifier, rank in semantic_ranks.items():
scores[identifier] += 1.0 / (self.rrf_k + rank)
ordered = sorted(documents.values(), key=lambda item: (-scores[item.id], item.id))
groups: dict[str, list[StoredDocument]] = {}
group_order: list[str] = []
for document in ordered:
key = document.content_hash
if key not in groups:
groups[key] = []
group_order.append(key)
groups[key].append(document)
results: list[SearchResult] = []
for key in group_order[:limit]:
group = groups[key]
primary = group[0]
alternates = [
{"source": document.configured_source, "url": document.canonical_url}
for document in group[1:]
]
results.append(
SearchResult(
id=primary.id,
configured_source=primary.configured_source,
canonical_url=primary.canonical_url,
title=primary.title,
description=primary.description,
heading_path=primary.heading_path,
content=primary.content,
content_hash=primary.content_hash,
score=min(1.0, scores[primary.id] / (2.0 / (self.rrf_k + 1))),
lexical_rank=lexical_ranks.get(primary.id),
semantic_rank=semantic_ranks.get(primary.id),
duplicate_count=len(group),
alternate_sources=alternates,
)
)
return results

View File

@@ -0,0 +1,183 @@
from __future__ import annotations
import asyncio
import os
import re
import time
from contextlib import asynccontextmanager
from pathlib import Path
import uvicorn
from mcp.server.fastmcp import FastMCP
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Route
from .embedder import SentenceTransformerEmbedder
from .fetcher import SourceFetcher
from .parser import parse_llms_text
from .refresh import RefreshCoordinator
from .search import HybridSearch
from .service import DocsService
from .store import IndexStore
def parse_duration(value: str) -> float:
match = re.fullmatch(r"\s*(\d+(?:\.\d+)?)\s*([smhd]?)\s*", value)
if not match:
raise ValueError(f"invalid duration: {value}")
multiplier = {"": 1, "s": 1, "m": 60, "h": 3600, "d": 86400}[match.group(2)]
return float(match.group(1)) * multiplier
def read_sources(path: str | Path) -> list[str]:
sources: list[str] = []
for raw_line in Path(path).read_text().splitlines():
line = raw_line.strip()
if not line or line.startswith("#"):
continue
if not line.endswith(("/llms.txt", "/llms-full.txt")):
raise ValueError(f"source URL must end with /llms.txt or /llms-full.txt: {line}")
sources.append(line)
if not sources:
raise ValueError(f"no sources configured in {path}")
return list(dict.fromkeys(sources))
def build_server():
source_file = os.environ.get("DOCS_MCP_SOURCES_FILE", "/etc/context-kit/docs-sources.txt")
sources = read_sources(source_file)
store = IndexStore(os.environ.get("DOCS_MCP_STORE_PATH", "/data/docs.sqlite3"))
store.configure_sources(sources)
embedder = SentenceTransformerEmbedder(
os.environ.get("DOCS_MCP_EMBED_MODEL", "BAAI/bge-small-en-v1.5")
)
fetcher = SourceFetcher(
timeout_seconds=float(os.environ.get("DOCS_MCP_FETCH_TIMEOUT", "30")),
max_bytes=int(os.environ.get("DOCS_MCP_MAX_SOURCE_BYTES", "20000000")),
)
coordinator = RefreshCoordinator(
store=store,
fetcher=fetcher,
embedder=embedder,
parser=parse_llms_text,
ttl_seconds=parse_duration(os.environ.get("DOCS_MCP_TTL", "24h")),
now=time.time,
)
service = DocsService(
store,
HybridSearch(store, embedder),
coordinator,
max_get_bytes=int(os.environ.get("DOCS_MCP_MAX_GET_BYTES", "75000")),
)
mcp = FastMCP(
"Context Kit Docs",
instructions="Search and retrieve configured documentation using persisted hybrid retrieval.",
host=os.environ.get("DOCS_MCP_HTTP_HOST", "0.0.0.0"),
port=int(os.environ.get("DOCS_MCP_HTTP_PORT", "8000")),
streamable_http_path="/mcp",
stateless_http=True,
)
@mcp.tool()
async def docs_query(
query: str,
limit: int = 10,
auto_retrieve: bool = False,
auto_retrieve_threshold: float = 0.55,
auto_retrieve_limit: int = 5,
retrieve_ids: list[str] | None = None,
max_bytes: int | None = None,
merge: bool = False,
sources: list[str] | None = None,
hosts: list[str] | None = None,
) -> dict:
"""Search docs. Content retrieval is explicit by default; optionally filter source URLs or hosts."""
return await service.query(
query, limit, auto_retrieve, auto_retrieve_threshold, auto_retrieve_limit,
retrieve_ids, max_bytes, merge, sources, hosts,
)
@mcp.tool()
async def docs_refresh(
source: str | None = None,
sources: list[str] | None = None,
force: bool = False,
) -> dict:
"""Refresh configured sources transactionally; concurrent requests are coalesced."""
if source and sources:
raise ValueError("pass source or sources, not both")
if source:
sources = [source]
return await service.refresh(sources, force)
@mcp.tool()
async def docs_sources() -> dict:
"""Report configured-source freshness, errors, and document counts."""
return service.source_status()
@mcp.tool()
async def docs_rebuild(
source: str | None = None,
sources: list[str] | None = None,
) -> dict:
"""Force a safe source rebuild without deleting the last good generation first."""
if source and sources:
raise ValueError("pass source or sources, not both")
if source:
sources = [source]
return await service.refresh(sources, force=True)
app = mcp.streamable_http_app()
mcp_lifespan = app.router.lifespan_context
@asynccontextmanager
async def application_lifespan(application):
preindex_task = None
async with mcp_lifespan(application):
if os.environ.get("DOCS_MCP_PREINDEX", "0") == "1":
preindex_task = asyncio.create_task(service.refresh())
try:
yield
finally:
if preindex_task:
await preindex_task
await fetcher.close()
store.close()
app.router.lifespan_context = application_lifespan
async def status(_request: Request) -> JSONResponse:
state = service.source_status()
errors = sum(1 for source in state["sources"] if source["last_error"])
return JSONResponse(
{
"status": "ok" if state["document_count"] or not errors else "degraded",
"ready": True,
"model_ready": embedder.ready,
"source_count": state["source_count"],
"document_count": state["document_count"],
"source_errors": errors,
}
)
app.routes.insert(0, Route("/status", status, methods=["GET"]))
origins = os.environ.get("DOCS_MCP_ALLOW_ORIGIN", "").split()
if origins:
app = CORSMiddleware(app, allow_origins=origins, allow_methods=["POST", "GET", "DELETE"], allow_headers=["*"])
return app
def main() -> None:
uvicorn.run(
build_server(),
host=os.environ.get("DOCS_MCP_HTTP_HOST", "0.0.0.0"),
port=int(os.environ.get("DOCS_MCP_HTTP_PORT", "8000")),
log_level=os.environ.get("DOCS_MCP_LOG_LEVEL", "info").lower(),
)
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,124 @@
from __future__ import annotations
import asyncio
from dataclasses import asdict
class DocsService:
def __init__(self, store, search, refresh, max_get_bytes: int = 75_000):
self.store = store
self.search_engine = search
self.refresh_coordinator = refresh
self.max_get_bytes = max_get_bytes
async def refresh(self, sources: list[str] | None = None, force: bool = False) -> dict:
configured = [state.configured_source for state in self.store.list_sources()]
selected = configured if sources is None else sources
unknown = sorted(set(selected) - set(configured))
if unknown:
raise ValueError(f"unconfigured sources: {', '.join(unknown)}")
outcomes = await asyncio.gather(
*(self.refresh_coordinator.refresh(source, force=force) for source in selected)
)
return {"sources": [asdict(outcome) for outcome in outcomes]}
async def query(
self,
query: str,
limit: int = 10,
auto_retrieve: bool = False,
auto_retrieve_threshold: float = 0.55,
auto_retrieve_limit: int = 5,
retrieve_ids: list[str] | None = None,
max_bytes: int | None = None,
merge: bool = False,
sources: list[str] | None = None,
hosts: list[str] | None = None,
) -> dict:
if not query.strip():
raise ValueError("query must not be empty")
if not 1 <= limit <= 100:
raise ValueError("limit must be between 1 and 100")
if not 0 <= auto_retrieve_threshold <= 1:
raise ValueError("auto_retrieve_threshold must be between 0 and 1")
if not 0 <= auto_retrieve_limit <= 25:
raise ValueError("auto_retrieve_limit must be between 0 and 25")
await self._refresh_missing_or_stale(sources)
results = await self.search_engine.search(query, limit, sources, hosts)
search_results = [
{
"id": item.id,
"source": item.configured_source,
"url": item.canonical_url,
"host": item.canonical_url.split("/", 3)[2] if "://" in item.canonical_url else "",
"title": item.title,
"description": item.description,
"heading_path": item.heading_path,
"score": round(item.score, 6),
"snippet": item.content[:500],
"duplicate_count": item.duplicate_count,
"alternate_sources": item.alternate_sources,
}
for item in results
]
selected_ids = list(dict.fromkeys(retrieve_ids or []))
if auto_retrieve:
selected_ids.extend(
item.id
for item in results[:auto_retrieve_limit]
if item.score >= auto_retrieve_threshold and item.id not in selected_ids
)
byte_budget = min(max_bytes or self.max_get_bytes, self.max_get_bytes)
retrieved: dict[str, dict] = {}
used = 0
for identifier in selected_ids[:25]:
document = self.store.get_document(identifier, sources, hosts)
if not document:
continue
encoded = document.content.encode()
remaining = max(0, byte_budget - used)
if remaining == 0:
break
content = encoded[:remaining].decode(errors="ignore")
used += len(content.encode())
retrieved[identifier] = {
"id": identifier,
"source": document.configured_source,
"url": document.canonical_url,
"title": document.title,
"content": content,
"truncated": len(content.encode()) < len(encoded),
}
merged = ""
if merge:
merged = "\n\n".join(
f"# {item['title']}\n\nSource: {item['url']}\n\n{item['content']}"
for item in retrieved.values()
)
return {
"search_results": search_results,
"retrieved_content": retrieved,
"merged_content": merged,
"auto_retrieved_count": len(retrieved) - len([item for item in retrieve_ids or [] if item in retrieved]),
"total_results": len(search_results),
}
async def _refresh_missing_or_stale(self, sources: list[str] | None) -> None:
states = self.store.list_sources()
selected = [state for state in states if sources is None or state.configured_source in sources]
await asyncio.gather(
*(self.refresh_coordinator.refresh(state.configured_source) for state in selected)
)
def source_status(self) -> dict:
states = [asdict(state) for state in self.store.list_sources()]
for state in states:
state.pop("raw_body", None)
return {
"sources": states,
"source_count": len(states),
"document_count": sum(state["doc_count"] for state in states),
}

View File

@@ -0,0 +1,282 @@
from __future__ import annotations
import re
import sqlite3
import threading
from pathlib import Path
import numpy as np
from .models import PreparedDocument, SourceState, SourceUpdate, StoredDocument
class IndexStore:
def __init__(self, path: str | Path):
self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True)
self.connection = sqlite3.connect(self.path, check_same_thread=False, isolation_level=None)
self.connection.row_factory = sqlite3.Row
self._lock = threading.RLock()
self._initialize()
def _initialize(self) -> None:
with self.connection:
self.connection.execute("PRAGMA journal_mode=WAL")
self.connection.execute("PRAGMA foreign_keys=ON")
self.connection.execute("PRAGMA busy_timeout=5000")
self.connection.executescript(
"""
CREATE TABLE IF NOT EXISTS sources (
configured_source TEXT PRIMARY KEY,
resolved_source TEXT,
active INTEGER NOT NULL DEFAULT 1,
etag TEXT,
last_modified TEXT,
body_hash TEXT,
raw_body TEXT,
parser_fingerprint TEXT,
embedding_fingerprint TEXT,
checked_at REAL,
indexed_at REAL,
stale_at REAL,
last_error TEXT
);
CREATE TABLE IF NOT EXISTS documents (
id TEXT PRIMARY KEY,
configured_source TEXT NOT NULL REFERENCES sources(configured_source) ON DELETE CASCADE,
resolved_source TEXT NOT NULL,
source_host TEXT NOT NULL,
canonical_url TEXT NOT NULL,
canonical_host TEXT NOT NULL,
title TEXT NOT NULL,
description TEXT NOT NULL,
heading_path TEXT NOT NULL,
content TEXT NOT NULL,
content_hash TEXT NOT NULL,
embedding BLOB NOT NULL,
embedding_dim INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS documents_source ON documents(configured_source);
CREATE INDEX IF NOT EXISTS documents_hash ON documents(content_hash);
CREATE INDEX IF NOT EXISTS documents_hosts ON documents(source_host, canonical_host);
CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
doc_id UNINDEXED, title, description, heading_path, content, canonical_url,
tokenize='unicode61 remove_diacritics 2 tokenchars ''_-'''
);
"""
)
def close(self) -> None:
self.connection.close()
def configure_sources(self, sources: list[str]) -> None:
with self._lock, self.connection:
self.connection.execute("UPDATE sources SET active = 0")
self.connection.executemany(
"INSERT INTO sources(configured_source, active) VALUES(?, 1) "
"ON CONFLICT(configured_source) DO UPDATE SET active = 1",
[(source,) for source in dict.fromkeys(sources)],
)
def replace_source(self, update: SourceUpdate) -> None:
with self._lock:
self.connection.execute("BEGIN IMMEDIATE")
try:
self._replace_source(update)
except Exception:
self.connection.rollback()
raise
else:
self.connection.commit()
def _replace_source(self, update: SourceUpdate) -> None:
self.connection.execute(
"""INSERT INTO sources(
configured_source, resolved_source, active, etag, last_modified,
body_hash, raw_body, parser_fingerprint, embedding_fingerprint,
checked_at, indexed_at, stale_at, last_error
) VALUES(?, ?, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL)
ON CONFLICT(configured_source) DO UPDATE SET
resolved_source=excluded.resolved_source, etag=excluded.etag,
last_modified=excluded.last_modified, body_hash=excluded.body_hash,
raw_body=excluded.raw_body, parser_fingerprint=excluded.parser_fingerprint,
embedding_fingerprint=excluded.embedding_fingerprint,
checked_at=excluded.checked_at, indexed_at=excluded.indexed_at,
stale_at=excluded.stale_at, last_error=NULL""",
(
update.configured_source,
update.resolved_source,
update.etag,
update.last_modified,
update.body_hash,
update.raw_body,
update.parser_fingerprint,
update.embedding_fingerprint,
update.checked_at,
update.indexed_at,
update.stale_at,
),
)
old_ids = [row[0] for row in self.connection.execute("SELECT id FROM documents WHERE configured_source=?", (update.configured_source,))]
if old_ids:
self.connection.executemany("DELETE FROM documents_fts WHERE doc_id=?", [(identifier,) for identifier in old_ids])
self.connection.execute("DELETE FROM documents WHERE configured_source=?", (update.configured_source,))
for document in update.documents:
vector = np.asarray(document.embedding, dtype=np.float32)
self.connection.execute(
"""INSERT INTO documents VALUES(
?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
)""",
(
document.id,
document.configured_source,
document.resolved_source,
document.source_host,
document.canonical_url,
document.canonical_host,
document.title,
document.description,
document.heading_path,
document.content,
document.content_hash,
vector.tobytes(),
vector.size,
),
)
self.connection.execute(
"INSERT INTO documents_fts VALUES(?, ?, ?, ?, ?, ?)",
(
document.id,
document.title,
document.description,
document.heading_path,
document.content,
document.canonical_url,
),
)
def mark_checked(self, source: str, checked_at: float, stale_at: float, error: str | None = None) -> None:
with self._lock, self.connection:
self.connection.execute(
"UPDATE sources SET checked_at=?, stale_at=?, last_error=? WHERE configured_source=?",
(checked_at, stale_at, error, source),
)
def list_sources(self, include_inactive: bool = False) -> list[SourceState]:
condition = "" if include_inactive else "WHERE s.active=1"
rows = self.connection.execute(
f"""SELECT s.*, COUNT(d.id) AS doc_count FROM sources s
LEFT JOIN documents d ON d.configured_source=s.configured_source
{condition} GROUP BY s.configured_source ORDER BY s.configured_source"""
).fetchall()
return [self._source(row) for row in rows]
def get_source(self, source: str) -> SourceState | None:
row = self.connection.execute(
"""SELECT s.*, COUNT(d.id) AS doc_count FROM sources s
LEFT JOIN documents d ON d.configured_source=s.configured_source
WHERE s.configured_source=? GROUP BY s.configured_source""",
(source,),
).fetchone()
return self._source(row) if row else None
def get_document(
self,
identifier: str,
sources: list[str] | None = None,
hosts: list[str] | None = None,
) -> StoredDocument | None:
where, parameters = self._filters(sources, hosts, alias="d")
row = self.connection.execute(
f"SELECT d.* FROM documents d JOIN sources s ON s.configured_source=d.configured_source "
f"WHERE s.active=1 AND d.id=? {where}",
[identifier, *parameters],
).fetchone()
return self._document(row) if row else None
def lexical_search(
self,
query: str,
limit: int,
sources: list[str] | None = None,
hosts: list[str] | None = None,
) -> list[StoredDocument]:
terms = re.findall(r"[\w.-]+", query, flags=re.UNICODE)
if not terms:
return []
expression = " AND ".join(f'"{term.replace(chr(34), chr(34) * 2)}"' for term in terms)
where, parameters = self._filters(sources, hosts, alias="d")
rows = self.connection.execute(
f"""SELECT d.* FROM documents_fts f
JOIN documents d ON d.id=f.doc_id
JOIN sources s ON s.configured_source=d.configured_source
WHERE documents_fts MATCH ? AND s.active=1 {where}
ORDER BY bm25(documents_fts, 0, 8, 3, 5, 1, 2) LIMIT ?""",
[expression, *parameters, limit],
).fetchall()
return [self._document(row) for row in rows]
def semantic_candidates(
self,
sources: list[str] | None = None,
hosts: list[str] | None = None,
) -> list[StoredDocument]:
where, parameters = self._filters(sources, hosts, alias="d")
rows = self.connection.execute(
f"SELECT d.* FROM documents d JOIN sources s ON s.configured_source=d.configured_source "
f"WHERE s.active=1 {where}",
parameters,
).fetchall()
return [self._document(row) for row in rows]
@staticmethod
def _filters(sources: list[str] | None, hosts: list[str] | None, alias: str) -> tuple[str, list[str]]:
clauses: list[str] = []
parameters: list[str] = []
if sources:
clauses.append(f"{alias}.configured_source IN ({','.join('?' for _ in sources)})")
parameters.extend(sources)
if hosts:
clauses.append(
f"({alias}.source_host IN ({','.join('?' for _ in hosts)}) OR "
f"{alias}.canonical_host IN ({','.join('?' for _ in hosts)}))"
)
parameters.extend(hosts)
parameters.extend(hosts)
return (" AND " + " AND ".join(clauses) if clauses else "", parameters)
@staticmethod
def _source(row: sqlite3.Row) -> SourceState:
return SourceState(
configured_source=row["configured_source"],
resolved_source=row["resolved_source"],
active=bool(row["active"]),
etag=row["etag"],
last_modified=row["last_modified"],
body_hash=row["body_hash"],
raw_body=row["raw_body"],
parser_fingerprint=row["parser_fingerprint"],
embedding_fingerprint=row["embedding_fingerprint"],
checked_at=row["checked_at"],
indexed_at=row["indexed_at"],
stale_at=row["stale_at"],
last_error=row["last_error"],
doc_count=row["doc_count"],
)
@staticmethod
def _document(row: sqlite3.Row) -> StoredDocument:
return StoredDocument(
id=row["id"],
configured_source=row["configured_source"],
resolved_source=row["resolved_source"],
source_host=row["source_host"],
canonical_url=row["canonical_url"],
canonical_host=row["canonical_host"],
title=row["title"],
description=row["description"],
heading_path=row["heading_path"],
content=row["content"],
content_hash=row["content_hash"],
embedding=np.frombuffer(row["embedding"], dtype=np.float32, count=row["embedding_dim"]).copy(),
)