"""Hub persistence adapter for project knowledge and PostgreSQL lexical search.""" from __future__ import annotations import re from collections.abc import Iterable from datetime import UTC, datetime from typing import Any from uuid import NAMESPACE_URL, UUID, uuid5 from sqlalchemy import delete, desc, func, select from sqlalchemy.orm import Session from packages.connector_kit.knowledge import ( DocumentRecord, RegisteredSource, SearchHit, SearchPage, SourceType, ) from packages.connector_kit.redaction import safe_snippet from ..persistence.models import KnowledgeChunkRow, KnowledgeDocumentRow, KnowledgeSourceRow class SqlAlchemyKnowledgeIndex: """Scoped read/write index adapter; all rows remain Hub-owned and project-filtered.""" def __init__(self, session: Session, *, workspace_id: UUID, project_id: UUID) -> None: self.session = session self.workspace_id = workspace_id self.project_id = project_id def sync_source(self, source: RegisteredSource, documents: Iterable[DocumentRecord]) -> None: if source.project_id != self.project_id: raise ValueError("source project does not match index scope") source_row = self.session.get(KnowledgeSourceRow, source.source_id) if source_row is None: source_row = KnowledgeSourceRow( id=source.source_id, workspace_id=self.workspace_id, project_id=self.project_id, slug=source.source_ref, source_type=source.source_type.value, uri=source.uri_prefix, visibility="project", state="registered" if source.enabled else "disabled", configuration={ "rootRegistered": True, "repositoryRef": source.repository_ref, "repositoryCommit": source.repository_commit, }, last_indexed_at=datetime.now(UTC), ) self.session.add(source_row) self.session.flush() else: if ( source_row.workspace_id != self.workspace_id or source_row.project_id != self.project_id ): raise ValueError("source row is outside index scope") source_row.state = "registered" if source.enabled else "disabled" source_row.last_indexed_at = datetime.now(UTC) active = { document.external_ref: document for document in documents if document.project_id == self.project_id and document.state == "registered" } existing = { row.external_ref: row for row in self.session.scalars( select(KnowledgeDocumentRow).where( KnowledgeDocumentRow.workspace_id == self.workspace_id, KnowledgeDocumentRow.project_id == self.project_id, KnowledgeDocumentRow.source_id == source.source_id, ) ).all() } for external_ref, document in active.items(): row = existing.get(external_ref) metadata = { "uri": document.uri, "path": document.external_ref, "sourceType": document.source_type.value, "mimeType": document.mime_type, "sourceRevision": document.source_revision, "repositoryRef": document.repository_ref, "repositoryCommit": document.repository_commit, "modifiedAt": document.modified_at.isoformat(), } if row is None: row = KnowledgeDocumentRow( id=stable_document_id(document.document_id), workspace_id=self.workspace_id, project_id=self.project_id, source_id=source.source_id, external_ref=external_ref, title=document.title, content_hash=document.content_hash, document_version=document.version, state="registered", content_bytes=document.content_bytes, record_metadata=metadata, ) self.session.add(row) self.session.flush() else: row.title = document.title row.content_hash = document.content_hash row.document_version = document.version row.state = "registered" row.content_bytes = document.content_bytes row.record_metadata = metadata self.session.execute( delete(KnowledgeChunkRow).where(KnowledgeChunkRow.document_id == row.id) ) self.session.add_all( [ KnowledgeChunkRow( id=hashlib_id(document.document_id, chunk.index), workspace_id=self.workspace_id, project_id=self.project_id, document_id=row.id, chunk_index=chunk.index, text=chunk.text, content_hash=chunk.content_hash, token_count=len(re.findall(r"\w+", chunk.text)), record_metadata={}, ) for chunk in document.chunks ] ) for external_ref, row in existing.items(): if external_ref in active: continue row.state = "deleted" row.content_bytes = 0 row.document_version += 1 row.record_metadata = {"tombstone": True, "path": external_ref} self.session.execute( delete(KnowledgeChunkRow).where(KnowledgeChunkRow.document_id == row.id) ) self.session.flush() def search(self, query: str, *, limit: int = 20) -> SearchPage: if not query.strip() or not 1 <= limit <= 20: raise ValueError("query or limit is invalid") document = KnowledgeDocumentRow chunk = KnowledgeChunkRow if self.session.bind is not None and self.session.bind.dialect.name == "postgresql": vector = func.to_tsvector("simple", chunk.text) tsquery = func.plainto_tsquery("simple", query) score = ( func.ts_rank_cd(vector, tsquery) + func.similarity(document.title, query) + func.similarity(chunk.text, query) ) statement = ( select(document, chunk, score.label("score")) .join(chunk, chunk.document_id == document.id) .where( document.workspace_id == self.workspace_id, document.project_id == self.project_id, document.state == "registered", vector.op("@@")(tsquery) | (func.similarity(document.title, query) > 0.1) | (func.similarity(chunk.text, query) > 0.1), ) .order_by(desc("score"), document.id, chunk.chunk_index) .limit(limit) ) else: statement = ( select(document, chunk) .join(chunk, chunk.document_id == document.id) .where( document.workspace_id == self.workspace_id, document.project_id == self.project_id, document.state == "registered", ) .order_by(document.id, chunk.chunk_index) ) hits: list[SearchHit] = [] terms = [term.casefold() for term in re.findall(r"[\w-]{2,}", query)] for result in self.session.execute(statement).all(): row = result[0] chunk_row = result[1] raw_score = ( float(result[2]) if len(result) > 2 else float( sum(chunk_row.text.casefold().count(term) for term in terms) + row.title.casefold().count(query.casefold()) * 4 ) ) if raw_score <= 0: continue metadata: dict[str, Any] = ( row.record_metadata if isinstance(row.record_metadata, dict) else {} ) hits.append( SearchHit( uri=str(metadata.get("uri", "")), title=row.title, path=row.external_ref, source_type=SourceType(metadata.get("sourceType", "documentation")), source_revision=str(metadata.get("sourceRevision", "registered")), repository_ref=metadata.get("repositoryRef"), repository_commit=metadata.get("repositoryCommit"), content_hash=row.content_hash, modified_at=datetime.fromisoformat(str(metadata.get("modifiedAt"))) if metadata.get("modifiedAt") else datetime.now(UTC), snippet=safe_snippet(chunk_row.text), score=raw_score, truncated=len(chunk_row.text.encode("utf-8")) > 2000, ) ) hits.sort(key=lambda hit: (-hit.score, hit.uri, hit.path)) return SearchPage(items=tuple(hits[:limit]), next_cursor=None, truncated=False) def stable_document_id(document_id: str) -> UUID: return uuid5(NAMESPACE_URL, "knowledge-document:" + document_id) def hashlib_id(document_id: str, chunk_index: int) -> UUID: """Stable UUID for a persisted chunk row.""" from uuid import uuid5 return uuid5(UUID(document_id[:32]), str(chunk_index))