Initial curated ITWorx_MCP_Hub source from 368fba691ea8d23403b4d56b53b5ec670d7c5020

packages/connector_kit/knowledge.py

@@ -0,0 +1,580 @@
+"""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)