234 lines
9.7 KiB
Python
234 lines
9.7 KiB
Python
"""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))
|