581 lines
21 KiB
Python
581 lines
21 KiB
Python
"""Registered-root project knowledge index with deterministic lexical search."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import mimetypes
|
|
import re
|
|
from collections.abc import Iterable
|
|
from datetime import UTC, datetime
|
|
from enum import StrEnum
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
from urllib.parse import unquote, urlsplit
|
|
from uuid import UUID
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
|
|
|
from .pagination import decode_cursor, encode_cursor
|
|
from .redaction import safe_content, safe_snippet
|
|
|
|
|
|
class SourceType(StrEnum):
|
|
DOCUMENTATION = "documentation"
|
|
ADR = "adr"
|
|
RUNBOOK = "runbook"
|
|
README = "readme"
|
|
CURRENT_STATE = "current-state"
|
|
RELEASE_NOTE = "release-note"
|
|
|
|
|
|
class PathEscapeError(ValueError):
|
|
"""A registered root or symlink would escape its configured containment."""
|
|
|
|
|
|
class DocumentNotFound(LookupError):
|
|
"""The requested source-backed document is not indexed for this project."""
|
|
|
|
|
|
class RegisteredSource(BaseModel):
|
|
"""A control-plane registered root; MCP callers never provide this path."""
|
|
|
|
model_config = ConfigDict(extra="forbid", frozen=True, arbitrary_types_allowed=True)
|
|
|
|
source_id: UUID
|
|
project_id: UUID
|
|
project_slug: str = Field(pattern=r"^[a-z0-9][a-z0-9-]{0,62}$")
|
|
source_ref: str = Field(pattern=r"^[a-z0-9][a-z0-9-]{0,127}$")
|
|
source_type: SourceType
|
|
root: Path
|
|
repository_ref: str | None = Field(default=None, max_length=200)
|
|
repository_commit: str | None = Field(default=None, max_length=200)
|
|
allowed_extensions: frozenset[str] = frozenset(
|
|
{".md", ".markdown", ".txt", ".rst", ".adoc", ".json", ".yaml", ".yml"}
|
|
)
|
|
excluded_dirs: frozenset[str] = frozenset(
|
|
{
|
|
".git",
|
|
".hg",
|
|
".svn",
|
|
"node_modules",
|
|
"dist",
|
|
"build",
|
|
"coverage",
|
|
"__pycache__",
|
|
".venv",
|
|
"venv",
|
|
"backups",
|
|
"secrets",
|
|
}
|
|
)
|
|
max_file_bytes: int = Field(default=1_048_576, ge=1024, le=4 * 1024 * 1024)
|
|
enabled: bool = True
|
|
revision: str = Field(default="registered", max_length=200)
|
|
|
|
@model_validator(mode="after")
|
|
def validate_root(self) -> RegisteredSource:
|
|
if not self.root.is_absolute():
|
|
raise ValueError("registered source roots must be absolute")
|
|
if any(not extension.startswith(".") for extension in self.allowed_extensions):
|
|
raise ValueError("allowed extensions must start with a dot")
|
|
return self
|
|
|
|
@property
|
|
def uri_prefix(self) -> str:
|
|
return f"itworx-docs://{self.project_slug}/{self.source_ref}"
|
|
|
|
def canonical_path(self, relative_path: str | Path) -> Path:
|
|
relative = Path(relative_path)
|
|
if relative.is_absolute() or ".." in relative.parts:
|
|
raise PathEscapeError("source path must remain relative to the registered root")
|
|
root = self.root.resolve(strict=True)
|
|
candidate = (root / relative).resolve(strict=False)
|
|
if candidate != root and root not in candidate.parents:
|
|
raise PathEscapeError("source path escapes the registered root")
|
|
return candidate
|
|
|
|
def allows(self, relative_path: Path) -> bool:
|
|
if not relative_path.parts or any(
|
|
part in self.excluded_dirs for part in relative_path.parts
|
|
):
|
|
return False
|
|
lowered = relative_path.name.lower()
|
|
if lowered in {
|
|
".env",
|
|
".env.local",
|
|
".env.production",
|
|
"id_rsa",
|
|
"id_ed25519",
|
|
"credentials.json",
|
|
"secrets.yaml",
|
|
}:
|
|
return False
|
|
if any(token in lowered for token in ("secret", "password", "token", "private-key")):
|
|
return False
|
|
return relative_path.suffix.lower() in self.allowed_extensions
|
|
|
|
|
|
class ChunkRecord(BaseModel):
|
|
model_config = ConfigDict(extra="forbid", frozen=True)
|
|
|
|
index: int = Field(ge=0)
|
|
text: str = Field(min_length=1, max_length=200_000)
|
|
content_hash: str = Field(pattern=r"^[0-9a-f]{64}$")
|
|
|
|
|
|
class DocumentRecord(BaseModel):
|
|
model_config = ConfigDict(extra="forbid", frozen=True)
|
|
|
|
document_id: str = Field(pattern=r"^[0-9a-f]{64}$")
|
|
source_id: UUID
|
|
project_id: UUID
|
|
project_slug: str
|
|
source_ref: str
|
|
source_type: SourceType
|
|
external_ref: str = Field(min_length=1, max_length=1024)
|
|
title: str = Field(min_length=1, max_length=500)
|
|
uri: str = Field(min_length=1, max_length=2048)
|
|
mime_type: str = Field(min_length=1, max_length=100)
|
|
content_hash: str = Field(pattern=r"^[0-9a-f]{64}$")
|
|
repository_ref: str | None = None
|
|
repository_commit: str | None = None
|
|
source_revision: str = Field(max_length=200)
|
|
modified_at: datetime
|
|
content_bytes: int = Field(ge=0, le=4 * 1024 * 1024)
|
|
version: int = Field(ge=1)
|
|
state: Literal["registered", "deleted"] = "registered"
|
|
content: str = ""
|
|
chunks: tuple[ChunkRecord, ...] = ()
|
|
|
|
|
|
class SearchHit(BaseModel):
|
|
model_config = ConfigDict(extra="forbid", frozen=True)
|
|
|
|
uri: str
|
|
title: str
|
|
path: str
|
|
source_type: SourceType
|
|
source_revision: str
|
|
repository_ref: str | None = None
|
|
repository_commit: str | None = None
|
|
content_hash: str
|
|
modified_at: datetime
|
|
snippet: str
|
|
score: float = Field(ge=0)
|
|
truncated: bool = False
|
|
|
|
|
|
class SearchPage(BaseModel):
|
|
model_config = ConfigDict(extra="forbid", frozen=True)
|
|
|
|
items: tuple[SearchHit, ...]
|
|
next_cursor: str | None
|
|
truncated: bool
|
|
|
|
|
|
class SourceStatus(BaseModel):
|
|
model_config = ConfigDict(extra="forbid", frozen=True)
|
|
|
|
source_ref: str
|
|
source_type: SourceType
|
|
uri: str
|
|
status: Literal["ready", "indexing", "stale", "failed", "empty", "disabled"]
|
|
source_revision: str | None
|
|
last_indexed_at: datetime | None
|
|
document_count: int = Field(ge=0)
|
|
|
|
|
|
class IndexStatus(BaseModel):
|
|
model_config = ConfigDict(extra="forbid", frozen=True)
|
|
|
|
status: Literal["ready", "indexing", "stale", "failed", "empty"]
|
|
source_count: int = Field(ge=0)
|
|
document_count: int = Field(ge=0)
|
|
last_completed_at: datetime | None
|
|
last_error: str | None = Field(default=None, max_length=1000)
|
|
|
|
|
|
class IndexReport(BaseModel):
|
|
model_config = ConfigDict(extra="forbid", frozen=True)
|
|
|
|
source_ref: str
|
|
indexed: int = Field(ge=0)
|
|
skipped: int = Field(ge=0)
|
|
tombstoned: int = Field(ge=0)
|
|
completed_at: datetime
|
|
|
|
|
|
class ParserRegistry:
|
|
def __init__(self) -> None:
|
|
self._extensions: dict[str, str] = {}
|
|
|
|
def register(self, extension: str, parser_name: str) -> None:
|
|
normalized = extension.lower()
|
|
if not normalized.startswith(".") or not parser_name:
|
|
raise ValueError("parser registration is invalid")
|
|
self._extensions[normalized] = parser_name
|
|
|
|
def supports(self, path: Path) -> bool:
|
|
return path.suffix.lower() in self._extensions
|
|
|
|
def title(self, path: Path, content: str) -> str:
|
|
for line in content.splitlines()[:20]:
|
|
match = re.match(r"^#{1,6}\s+(.+?)\s*#*$", line.strip())
|
|
if match:
|
|
return match.group(1)[:500]
|
|
return path.stem.replace("_", " ").replace("-", " ").strip()[:500] or path.name
|
|
|
|
@classmethod
|
|
def default(cls) -> ParserRegistry:
|
|
registry = cls()
|
|
for extension in (".md", ".markdown", ".txt", ".rst", ".adoc", ".json", ".yaml", ".yml"):
|
|
registry.register(extension, "plain-text")
|
|
return registry
|
|
|
|
|
|
class DeterministicChunker:
|
|
def __init__(self, *, max_chars: int = 4000, overlap_chars: int = 200) -> None:
|
|
if not 256 <= max_chars <= 20_000 or not 0 <= overlap_chars < max_chars // 2:
|
|
raise ValueError("chunk bounds are invalid")
|
|
self.max_chars = max_chars
|
|
self.overlap_chars = overlap_chars
|
|
|
|
def chunk(self, content: str) -> tuple[ChunkRecord, ...]:
|
|
if not content:
|
|
return ()
|
|
chunks: list[ChunkRecord] = []
|
|
start = 0
|
|
index = 0
|
|
while start < len(content):
|
|
end = min(len(content), start + self.max_chars)
|
|
if end < len(content):
|
|
boundary = content.rfind("\n", start + self.max_chars // 2, end)
|
|
if boundary > start:
|
|
end = boundary
|
|
text = content[start:end].strip()
|
|
if text:
|
|
chunks.append(
|
|
ChunkRecord(
|
|
index=index,
|
|
text=text,
|
|
content_hash=hashlib.sha256(text.encode("utf-8")).hexdigest(),
|
|
)
|
|
)
|
|
index += 1
|
|
if end >= len(content):
|
|
break
|
|
start = max(end - self.overlap_chars, start + 1)
|
|
return tuple(chunks)
|
|
|
|
|
|
def _document_key(source: RegisteredSource, relative: str) -> str:
|
|
return hashlib.sha256(f"{source.source_id}:{relative}".encode()).hexdigest()
|
|
|
|
|
|
def _content_hash(content: str) -> str:
|
|
return hashlib.sha256(content.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _score_text(query: str, title: str, content: str) -> float:
|
|
terms = [term for term in re.findall(r"[\w-]{2,}", query.casefold())]
|
|
if not terms:
|
|
return 0.0
|
|
folded_title = title.casefold()
|
|
folded_content = content.casefold()
|
|
score = sum(folded_title.count(term) * 4 for term in terms)
|
|
score += sum(folded_content.count(term) for term in terms)
|
|
if query.casefold() in folded_content:
|
|
score += 2
|
|
return float(score)
|
|
|
|
|
|
class InMemoryKnowledgeIndex:
|
|
"""Deterministic index used by the connector and as a fake target for tests."""
|
|
|
|
def __init__(
|
|
self, *, parser: ParserRegistry | None = None, chunker: DeterministicChunker | None = None
|
|
) -> None:
|
|
self.parser = parser or ParserRegistry.default()
|
|
self.chunker = chunker or DeterministicChunker()
|
|
self.sources: dict[UUID, RegisteredSource] = {}
|
|
self.documents: dict[str, DocumentRecord] = {}
|
|
self._last_completed_at: datetime | None = None
|
|
self._last_error: str | None = None
|
|
self._indexing = False
|
|
|
|
def register_source(self, source: RegisteredSource) -> None:
|
|
if not source.enabled:
|
|
self.sources[source.source_id] = source
|
|
return
|
|
canonical_root = source.root.resolve(strict=True)
|
|
if not canonical_root.is_dir():
|
|
raise ValueError("registered source root must be a directory")
|
|
self.sources[source.source_id] = source.model_copy(update={"root": canonical_root})
|
|
|
|
def index_source(self, source_id: UUID) -> IndexReport:
|
|
source = self.sources.get(source_id)
|
|
if source is None:
|
|
raise KeyError("source is not registered")
|
|
if not source.enabled:
|
|
return IndexReport(
|
|
source_ref=source.source_ref,
|
|
indexed=0,
|
|
skipped=0,
|
|
tombstoned=0,
|
|
completed_at=datetime.now(UTC),
|
|
)
|
|
self._indexing = True
|
|
self._last_error = None
|
|
previous = {
|
|
doc.external_ref: doc
|
|
for doc in self.documents.values()
|
|
if doc.source_id == source.source_id and doc.state == "registered"
|
|
}
|
|
seen: set[str] = set()
|
|
indexed = 0
|
|
skipped = 0
|
|
try:
|
|
for path in self._iter_files(source):
|
|
relative = path.relative_to(source.root.resolve(strict=True)).as_posix()
|
|
seen.add(relative)
|
|
try:
|
|
text = path.read_text(encoding="utf-8")
|
|
if len(text.encode("utf-8")) > source.max_file_bytes:
|
|
skipped += 1
|
|
continue
|
|
parsed = safe_content(text)
|
|
if parsed is None or not self.parser.supports(path):
|
|
skipped += 1
|
|
continue
|
|
key = _document_key(source, relative)
|
|
old = previous.get(relative)
|
|
content_hash = _content_hash(parsed)
|
|
if old is not None and old.content_hash == content_hash:
|
|
continue
|
|
modified = datetime.fromtimestamp(path.stat().st_mtime, tz=UTC)
|
|
document = DocumentRecord(
|
|
document_id=key,
|
|
source_id=source.source_id,
|
|
project_id=source.project_id,
|
|
project_slug=source.project_slug,
|
|
source_ref=source.source_ref,
|
|
source_type=source.source_type,
|
|
external_ref=relative,
|
|
title=self.parser.title(path, parsed),
|
|
uri=f"{source.uri_prefix}/{key}",
|
|
mime_type=mimetypes.guess_type(path.name)[0] or "text/plain",
|
|
content_hash=content_hash,
|
|
repository_ref=source.repository_ref,
|
|
repository_commit=source.repository_commit,
|
|
source_revision=source.revision,
|
|
modified_at=modified,
|
|
content_bytes=len(parsed.encode("utf-8")),
|
|
version=(old.version + 1 if old is not None else 1),
|
|
content=parsed,
|
|
chunks=self.chunker.chunk(parsed),
|
|
)
|
|
self.documents[key] = document
|
|
indexed += 1
|
|
except (OSError, UnicodeDecodeError):
|
|
skipped += 1
|
|
tombstoned = 0
|
|
for relative, old in previous.items():
|
|
if relative not in seen:
|
|
self.documents[old.document_id] = old.model_copy(
|
|
update={
|
|
"state": "deleted",
|
|
"content": "",
|
|
"chunks": (),
|
|
"content_bytes": 0,
|
|
"version": old.version + 1,
|
|
}
|
|
)
|
|
tombstoned += 1
|
|
self._last_completed_at = datetime.now(UTC)
|
|
return IndexReport(
|
|
source_ref=source.source_ref,
|
|
indexed=indexed,
|
|
skipped=skipped,
|
|
tombstoned=tombstoned,
|
|
completed_at=self._last_completed_at,
|
|
)
|
|
except Exception as exc:
|
|
self._last_error = type(exc).__name__
|
|
raise
|
|
finally:
|
|
self._indexing = False
|
|
|
|
def _iter_files(self, source: RegisteredSource) -> Iterable[Path]:
|
|
root = source.root.resolve(strict=True)
|
|
for path in sorted(root.rglob("*"), key=lambda item: item.as_posix().casefold()):
|
|
if not path.is_file():
|
|
continue
|
|
relative = path.relative_to(root)
|
|
if not source.allows(relative):
|
|
continue
|
|
canonical = source.canonical_path(relative)
|
|
if canonical != path.resolve(strict=False) and not path.is_symlink():
|
|
continue
|
|
if not canonical.is_file():
|
|
continue
|
|
yield canonical
|
|
|
|
def get_document(self, project_id: UUID, uri: str) -> DocumentRecord:
|
|
parts = urlsplit(uri)
|
|
if parts.scheme != "itworx-docs" or parts.query or parts.fragment:
|
|
raise DocumentNotFound("document URI is invalid")
|
|
project = unquote(parts.netloc)
|
|
path_parts = [unquote(part) for part in parts.path.split("/") if part]
|
|
if len(path_parts) != 2 or any(
|
|
part in {".", ".."} or "/" in part or "\\" in part for part in path_parts
|
|
):
|
|
raise DocumentNotFound("document URI is invalid")
|
|
if project != next(
|
|
(
|
|
source.project_slug
|
|
for source in self.sources.values()
|
|
if source.project_id == project_id
|
|
),
|
|
"",
|
|
):
|
|
raise DocumentNotFound("document is outside the caller project")
|
|
document = self.documents.get(path_parts[1])
|
|
if (
|
|
document is None
|
|
or document.project_id != project_id
|
|
or document.state != "registered"
|
|
or document.source_ref != path_parts[0]
|
|
):
|
|
raise DocumentNotFound("document is not indexed")
|
|
return document
|
|
|
|
def search(
|
|
self,
|
|
project_id: UUID,
|
|
query: str,
|
|
*,
|
|
source_types: Iterable[SourceType] = (),
|
|
limit: int = 20,
|
|
cursor: str | None = None,
|
|
) -> SearchPage:
|
|
if not query.strip() or len(query) > 1000:
|
|
raise ValueError("query must be between 1 and 1000 characters")
|
|
if not 1 <= limit <= 20:
|
|
raise ValueError("limit must be between 1 and 20")
|
|
selected = tuple(sorted(str(value) for value in source_types))
|
|
page_cursor = decode_cursor(cursor, str(project_id), query, *selected)
|
|
hits: list[SearchHit] = []
|
|
for document in self.documents.values():
|
|
if (
|
|
document.project_id != project_id
|
|
or document.state != "registered"
|
|
or (selected and str(document.source_type) not in selected)
|
|
):
|
|
continue
|
|
for chunk in document.chunks:
|
|
score = _score_text(query, document.title, chunk.text)
|
|
if score <= 0:
|
|
continue
|
|
hits.append(
|
|
SearchHit(
|
|
uri=document.uri,
|
|
title=document.title,
|
|
path=document.external_ref,
|
|
source_type=document.source_type,
|
|
source_revision=document.source_revision,
|
|
repository_ref=document.repository_ref,
|
|
repository_commit=document.repository_commit,
|
|
content_hash=document.content_hash,
|
|
modified_at=document.modified_at,
|
|
snippet=safe_snippet(chunk.text),
|
|
score=score,
|
|
truncated=len(chunk.text.encode("utf-8")) > 2000,
|
|
)
|
|
)
|
|
hits.sort(key=lambda hit: (-hit.score, hit.uri, hit.path))
|
|
page = hits[page_cursor.offset : page_cursor.offset + limit]
|
|
next_offset = page_cursor.offset + len(page)
|
|
next_cursor = (
|
|
encode_cursor(next_offset, str(project_id), query, *selected)
|
|
if next_offset < len(hits)
|
|
else None
|
|
)
|
|
return SearchPage(
|
|
items=tuple(page), next_cursor=next_cursor, truncated=next_cursor is not None
|
|
)
|
|
|
|
def list_sources(self, project_id: UUID) -> tuple[SourceStatus, ...]:
|
|
return tuple(
|
|
SourceStatus(
|
|
source_ref=source.source_ref,
|
|
source_type=source.source_type,
|
|
uri=source.uri_prefix,
|
|
status="disabled"
|
|
if not source.enabled
|
|
else (
|
|
"ready"
|
|
if any(
|
|
doc.project_id == project_id
|
|
and doc.source_id == source.source_id
|
|
and doc.state == "registered"
|
|
for doc in self.documents.values()
|
|
)
|
|
else "empty"
|
|
),
|
|
source_revision=source.revision,
|
|
last_indexed_at=self._last_completed_at,
|
|
document_count=sum(
|
|
1
|
|
for doc in self.documents.values()
|
|
if doc.project_id == project_id
|
|
and doc.source_id == source.source_id
|
|
and doc.state == "registered"
|
|
),
|
|
)
|
|
for source in sorted(self.sources.values(), key=lambda item: item.source_ref)
|
|
if source.project_id == project_id
|
|
)
|
|
|
|
def status(self, project_id: UUID) -> IndexStatus:
|
|
sources = self.list_sources(project_id)
|
|
count = sum(source.document_count for source in sources)
|
|
status: Literal["ready", "indexing", "stale", "failed", "empty"] = (
|
|
"indexing"
|
|
if self._indexing
|
|
else ("failed" if self._last_error else ("empty" if count == 0 else "ready"))
|
|
)
|
|
return IndexStatus(
|
|
status=status,
|
|
source_count=len(sources),
|
|
document_count=count,
|
|
last_completed_at=self._last_completed_at,
|
|
last_error=self._last_error,
|
|
)
|
|
|
|
def documents_for_source(
|
|
self, project_id: UUID, source_type: SourceType | None = None
|
|
) -> tuple[DocumentRecord, ...]:
|
|
return tuple(
|
|
sorted(
|
|
(
|
|
doc
|
|
for doc in self.documents.values()
|
|
if doc.project_id == project_id
|
|
and doc.state == "registered"
|
|
and (source_type is None or doc.source_type == source_type)
|
|
),
|
|
key=lambda doc: (doc.external_ref, doc.document_id),
|
|
)
|
|
)
|
|
|
|
|
|
class KnowledgeRefreshJob:
|
|
"""Single bounded refresh operation; scheduling is owned by the control plane."""
|
|
|
|
def __init__(self, index: InMemoryKnowledgeIndex, source_ids: Iterable[UUID]) -> None:
|
|
self.index = index
|
|
self.source_ids = tuple(source_ids)
|
|
|
|
def run_once(self) -> tuple[IndexReport, ...]:
|
|
return tuple(self.index.index_source(source_id) for source_id in self.source_ids)
|