998 lines
39 KiB
Python
998 lines
39 KiB
Python
|
"""Scope and freshness policy for derived canonical search."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import base64
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
import secrets
|
||
|
|
import sqlite3
|
||
|
|
from collections import defaultdict
|
||
|
|
from collections.abc import Callable, Collection, Iterator
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from datetime import UTC, date, datetime
|
||
|
|
from typing import Literal, Protocol
|
||
|
|
|
||
|
|
from gewu.application.identity import Principal
|
||
|
|
from gewu.application.ports import EmbeddingProvider, VectorGeneration, VectorIndex
|
||
|
|
from gewu.domain.contracts import (
|
||
|
|
BudgetProfile,
|
||
|
|
Confidence,
|
||
|
|
EdgeEndpoint,
|
||
|
|
EvidenceReference,
|
||
|
|
Finding,
|
||
|
|
KnowledgeNode,
|
||
|
|
NeighborRequest,
|
||
|
|
ReadEnvelope,
|
||
|
|
Scope,
|
||
|
|
SearchRequest,
|
||
|
|
SourceListRequest,
|
||
|
|
SourceRevisionRecord,
|
||
|
|
TypedEdge,
|
||
|
|
)
|
||
|
|
from gewu.domain.validation import canonical_json, evidence_digest
|
||
|
|
|
||
|
|
|
||
|
|
class SearchError(RuntimeError):
|
||
|
|
"""Safe search failure with an HTTP-compatible status."""
|
||
|
|
|
||
|
|
def __init__(self, status_code: int, message: str) -> None:
|
||
|
|
super().__init__(message)
|
||
|
|
self.status_code = status_code
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True, slots=True)
|
||
|
|
class IndexedNode:
|
||
|
|
node: KnowledgeNode
|
||
|
|
score: float
|
||
|
|
expanded_from: str | None = None
|
||
|
|
expansion_edge: TypedEdge | None = None
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True, slots=True)
|
||
|
|
class SearchIndexResult:
|
||
|
|
index_revision: str
|
||
|
|
nodes: tuple[IndexedNode, ...]
|
||
|
|
has_more: bool = False
|
||
|
|
retrieval_mode: Literal["keyword", "hybrid"] = "keyword"
|
||
|
|
warnings: tuple[Finding, ...] = ()
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True, slots=True)
|
||
|
|
class SourceIndexResult:
|
||
|
|
index_revision: str
|
||
|
|
sources: tuple[SourceRevisionRecord, ...]
|
||
|
|
has_more: bool = False
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True, slots=True)
|
||
|
|
class SearchEdgeResult:
|
||
|
|
index_revision: str
|
||
|
|
origin: KnowledgeNode | None
|
||
|
|
edges: tuple[SearchEdgeHit, ...]
|
||
|
|
has_more: bool = False
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True, slots=True)
|
||
|
|
class SearchEdgeHit:
|
||
|
|
edge: TypedEdge
|
||
|
|
neighbor_endpoint: EdgeEndpoint
|
||
|
|
neighbor_node: KnowledgeNode | None
|
||
|
|
neighbor_source: SourceRevisionRecord | None
|
||
|
|
|
||
|
|
|
||
|
|
class SearchIndex(Protocol):
|
||
|
|
def diagnostic_status(self) -> tuple[str | None, int, bool]: ...
|
||
|
|
|
||
|
|
def search(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
query: str,
|
||
|
|
project_ids: tuple[str, ...],
|
||
|
|
include_shared: bool,
|
||
|
|
limit: int,
|
||
|
|
node_types: tuple[str, ...] = (),
|
||
|
|
tags: tuple[str, ...] = (),
|
||
|
|
audience: tuple[str, ...] = (),
|
||
|
|
use_for: tuple[str, ...] = (),
|
||
|
|
domain: tuple[str, ...] = (),
|
||
|
|
confidences: tuple[Confidence, ...] = (),
|
||
|
|
as_of_from: date | None = None,
|
||
|
|
as_of_to: date | None = None,
|
||
|
|
stale_only: bool = False,
|
||
|
|
candidate_node_ids: tuple[str, ...] | None = None,
|
||
|
|
offset: int = 0,
|
||
|
|
) -> SearchIndexResult: ...
|
||
|
|
|
||
|
|
def neighbors(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
node_id: str,
|
||
|
|
edge_types: tuple[str, ...],
|
||
|
|
project_ids: tuple[str, ...],
|
||
|
|
include_shared: bool,
|
||
|
|
limit: int,
|
||
|
|
offset: int = 0,
|
||
|
|
) -> SearchEdgeResult: ...
|
||
|
|
|
||
|
|
def list_sources(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
project_ids: tuple[str, ...],
|
||
|
|
include_shared: bool,
|
||
|
|
limit: int,
|
||
|
|
offset: int = 0,
|
||
|
|
) -> SourceIndexResult: ...
|
||
|
|
|
||
|
|
|
||
|
|
_VECTOR_CANDIDATE_LIMIT = 100
|
||
|
|
_VECTOR_BUILD_BATCH_SIZE = 32
|
||
|
|
_GRAPH_SEED_LIMIT = 10
|
||
|
|
_GRAPH_EDGE_LIMIT = 50
|
||
|
|
_GRAPH_EXPANSION_SCORE = 1 / 160.5
|
||
|
|
_NODE_CONTENT_PROFILE = "knowledge-node-title-summary-claims-tags-v1"
|
||
|
|
|
||
|
|
|
||
|
|
class HybridSearchIndex:
|
||
|
|
"""Fuse scoped lexical and vector candidates over a pinned vector generation."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
lexical: SearchIndex,
|
||
|
|
vectors: VectorIndex,
|
||
|
|
embeddings: EmbeddingProvider,
|
||
|
|
) -> None:
|
||
|
|
self.lexical = lexical
|
||
|
|
self.vectors = vectors
|
||
|
|
self.embeddings = embeddings
|
||
|
|
self._failed_profile: tuple[str, str] | None = None
|
||
|
|
|
||
|
|
def diagnostic_status(self) -> tuple[str | None, int, bool]:
|
||
|
|
return self.lexical.diagnostic_status()
|
||
|
|
|
||
|
|
def cursor_profile(self) -> str:
|
||
|
|
return hashlib.sha256(
|
||
|
|
"\0".join(
|
||
|
|
(
|
||
|
|
"hybrid-rrf-v1",
|
||
|
|
"typed-graph-one-hop-v1",
|
||
|
|
str(_VECTOR_CANDIDATE_LIMIT),
|
||
|
|
str(_GRAPH_SEED_LIMIT),
|
||
|
|
str(_GRAPH_EDGE_LIMIT),
|
||
|
|
self.embeddings.model_id,
|
||
|
|
self.embeddings.artifact_sha256,
|
||
|
|
str(self.embeddings.dimensions),
|
||
|
|
_NODE_CONTENT_PROFILE,
|
||
|
|
)
|
||
|
|
).encode()
|
||
|
|
).hexdigest()
|
||
|
|
|
||
|
|
def search(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
query: str,
|
||
|
|
project_ids: tuple[str, ...],
|
||
|
|
include_shared: bool,
|
||
|
|
limit: int,
|
||
|
|
node_types: tuple[str, ...] = (),
|
||
|
|
tags: tuple[str, ...] = (),
|
||
|
|
audience: tuple[str, ...] = (),
|
||
|
|
use_for: tuple[str, ...] = (),
|
||
|
|
domain: tuple[str, ...] = (),
|
||
|
|
confidences: tuple[Confidence, ...] = (),
|
||
|
|
as_of_from: date | None = None,
|
||
|
|
as_of_to: date | None = None,
|
||
|
|
stale_only: bool = False,
|
||
|
|
candidate_node_ids: tuple[str, ...] | None = None,
|
||
|
|
offset: int = 0,
|
||
|
|
) -> SearchIndexResult:
|
||
|
|
if not query.strip() or candidate_node_ids is not None:
|
||
|
|
return self.lexical.search(
|
||
|
|
query=query,
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=limit,
|
||
|
|
node_types=node_types,
|
||
|
|
tags=tags,
|
||
|
|
audience=audience,
|
||
|
|
use_for=use_for,
|
||
|
|
domain=domain,
|
||
|
|
confidences=confidences,
|
||
|
|
as_of_from=as_of_from,
|
||
|
|
as_of_to=as_of_to,
|
||
|
|
stale_only=stale_only,
|
||
|
|
candidate_node_ids=candidate_node_ids,
|
||
|
|
offset=offset,
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
lexical_head = self.lexical.search(
|
||
|
|
query=query,
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=_VECTOR_CANDIDATE_LIMIT,
|
||
|
|
node_types=node_types,
|
||
|
|
tags=tags,
|
||
|
|
audience=audience,
|
||
|
|
use_for=use_for,
|
||
|
|
domain=domain,
|
||
|
|
confidences=confidences,
|
||
|
|
as_of_from=as_of_from,
|
||
|
|
as_of_to=as_of_to,
|
||
|
|
stale_only=stale_only,
|
||
|
|
)
|
||
|
|
generation = self.vectors.active_generation()
|
||
|
|
if not self._matches_generation(generation, lexical_head.index_revision):
|
||
|
|
raise RuntimeError("Vector generation is unavailable")
|
||
|
|
assert generation is not None
|
||
|
|
query_vector = self.embeddings.embed_queries((query,))[0]
|
||
|
|
vector_hits = self.vectors.search(
|
||
|
|
query=query_vector,
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=_VECTOR_CANDIDATE_LIMIT,
|
||
|
|
expected_generation_id=generation.generation_id,
|
||
|
|
)
|
||
|
|
vector_nodes = self.lexical.search(
|
||
|
|
query="",
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=_VECTOR_CANDIDATE_LIMIT,
|
||
|
|
node_types=node_types,
|
||
|
|
tags=tags,
|
||
|
|
audience=audience,
|
||
|
|
use_for=use_for,
|
||
|
|
domain=domain,
|
||
|
|
confidences=confidences,
|
||
|
|
as_of_from=as_of_from,
|
||
|
|
as_of_to=as_of_to,
|
||
|
|
stale_only=stale_only,
|
||
|
|
candidate_node_ids=tuple(hit.node_id for hit in vector_hits),
|
||
|
|
)
|
||
|
|
vector_node_by_id = {hit.node.node_id: hit.node for hit in vector_nodes.nodes}
|
||
|
|
fused_scores: defaultdict[str, float] = defaultdict(float)
|
||
|
|
nodes_by_id: dict[str, KnowledgeNode] = {}
|
||
|
|
for rank, hit in enumerate(lexical_head.nodes, start=1):
|
||
|
|
fused_scores[hit.node.node_id] += 1 / (60 + rank)
|
||
|
|
nodes_by_id[hit.node.node_id] = hit.node
|
||
|
|
for rank, vector_hit in enumerate(vector_hits, start=1):
|
||
|
|
node = vector_node_by_id.get(vector_hit.node_id)
|
||
|
|
if node is not None:
|
||
|
|
fused_scores[vector_hit.node_id] += 1 / (60 + rank)
|
||
|
|
nodes_by_id[vector_hit.node_id] = node
|
||
|
|
ranked = sorted(
|
||
|
|
fused_scores,
|
||
|
|
key=lambda node_id: (-fused_scores[node_id], node_id),
|
||
|
|
)
|
||
|
|
graph_candidates: list[tuple[str, TypedEdge, KnowledgeNode]] = []
|
||
|
|
graph_node_ids: set[str] = set()
|
||
|
|
fused_node_ids = set(ranked)
|
||
|
|
examined_edge_count = 0
|
||
|
|
for seed_id in ranked[:_GRAPH_SEED_LIMIT]:
|
||
|
|
remaining = _GRAPH_EDGE_LIMIT - examined_edge_count
|
||
|
|
if remaining <= 0:
|
||
|
|
break
|
||
|
|
neighbors = self.lexical.neighbors(
|
||
|
|
node_id=seed_id,
|
||
|
|
edge_types=(),
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=remaining,
|
||
|
|
)
|
||
|
|
examined_edge_count += len(neighbors.edges)
|
||
|
|
for neighbor in neighbors.edges:
|
||
|
|
node = neighbor.neighbor_node
|
||
|
|
if (
|
||
|
|
node is None
|
||
|
|
or node.node_id in fused_node_ids
|
||
|
|
or node.node_id in graph_node_ids
|
||
|
|
):
|
||
|
|
continue
|
||
|
|
graph_node_ids.add(node.node_id)
|
||
|
|
graph_candidates.append((seed_id, neighbor.edge, node))
|
||
|
|
if graph_candidates:
|
||
|
|
eligible_graph_nodes = self.lexical.search(
|
||
|
|
query="",
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=len(graph_candidates),
|
||
|
|
node_types=node_types,
|
||
|
|
tags=tags,
|
||
|
|
audience=audience,
|
||
|
|
use_for=use_for,
|
||
|
|
domain=domain,
|
||
|
|
confidences=confidences,
|
||
|
|
as_of_from=as_of_from,
|
||
|
|
as_of_to=as_of_to,
|
||
|
|
stale_only=stale_only,
|
||
|
|
candidate_node_ids=tuple(node.node_id for _, _, node in graph_candidates),
|
||
|
|
)
|
||
|
|
eligible_ids = {hit.node.node_id for hit in eligible_graph_nodes.nodes}
|
||
|
|
expanded_nodes = tuple(
|
||
|
|
IndexedNode(
|
||
|
|
node=node,
|
||
|
|
score=_GRAPH_EXPANSION_SCORE,
|
||
|
|
expanded_from=seed_id,
|
||
|
|
expansion_edge=edge,
|
||
|
|
)
|
||
|
|
for seed_id, edge, node in graph_candidates
|
||
|
|
if node.node_id in eligible_ids
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
expanded_nodes = ()
|
||
|
|
fixed_head = tuple(
|
||
|
|
[IndexedNode(nodes_by_id[node_id], fused_scores[node_id]) for node_id in ranked]
|
||
|
|
+ list(expanded_nodes)
|
||
|
|
)
|
||
|
|
head_start = min(offset, len(fixed_head))
|
||
|
|
head_page = [*fixed_head[head_start : offset + limit]]
|
||
|
|
tail_offset = max(0, offset - len(fixed_head))
|
||
|
|
tail_limit = limit - len(head_page)
|
||
|
|
lexical_tail = SearchIndexResult(lexical_head.index_revision, ())
|
||
|
|
if lexical_head.has_more and tail_limit > 0:
|
||
|
|
lexical_tail = self.lexical.search(
|
||
|
|
query=query,
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=tail_limit,
|
||
|
|
node_types=node_types,
|
||
|
|
tags=tags,
|
||
|
|
audience=audience,
|
||
|
|
use_for=use_for,
|
||
|
|
domain=domain,
|
||
|
|
confidences=confidences,
|
||
|
|
as_of_from=as_of_from,
|
||
|
|
as_of_to=as_of_to,
|
||
|
|
stale_only=stale_only,
|
||
|
|
offset=_VECTOR_CANDIDATE_LIMIT + tail_offset,
|
||
|
|
)
|
||
|
|
lexical_tail = SearchIndexResult(
|
||
|
|
index_revision=lexical_tail.index_revision,
|
||
|
|
nodes=tuple(
|
||
|
|
IndexedNode(
|
||
|
|
hit.node,
|
||
|
|
1 / (60 + _VECTOR_CANDIDATE_LIMIT + tail_offset + rank),
|
||
|
|
)
|
||
|
|
for rank, hit in enumerate(lexical_tail.nodes, start=1)
|
||
|
|
),
|
||
|
|
has_more=lexical_tail.has_more,
|
||
|
|
)
|
||
|
|
page = (*head_page, *lexical_tail.nodes)
|
||
|
|
has_more = (
|
||
|
|
offset + len(head_page) < len(fixed_head)
|
||
|
|
or (tail_limit == 0 and lexical_head.has_more)
|
||
|
|
or (tail_limit > 0 and lexical_tail.has_more)
|
||
|
|
)
|
||
|
|
return SearchIndexResult(
|
||
|
|
index_revision=lexical_head.index_revision,
|
||
|
|
nodes=page,
|
||
|
|
has_more=has_more,
|
||
|
|
retrieval_mode="hybrid",
|
||
|
|
)
|
||
|
|
except (OSError, RuntimeError, ValueError, IndexError):
|
||
|
|
warning = Finding(
|
||
|
|
code="HYBRID_RETRIEVAL_DEGRADED",
|
||
|
|
message=(
|
||
|
|
"Vector retrieval is unavailable for this index revision; "
|
||
|
|
"keyword results are shown."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
fallback = self.lexical.search(
|
||
|
|
query=query,
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=limit,
|
||
|
|
node_types=node_types,
|
||
|
|
tags=tags,
|
||
|
|
audience=audience,
|
||
|
|
use_for=use_for,
|
||
|
|
domain=domain,
|
||
|
|
confidences=confidences,
|
||
|
|
as_of_from=as_of_from,
|
||
|
|
as_of_to=as_of_to,
|
||
|
|
stale_only=stale_only,
|
||
|
|
offset=offset,
|
||
|
|
)
|
||
|
|
return SearchIndexResult(
|
||
|
|
index_revision=fallback.index_revision,
|
||
|
|
nodes=fallback.nodes,
|
||
|
|
has_more=fallback.has_more,
|
||
|
|
warnings=(warning,),
|
||
|
|
)
|
||
|
|
|
||
|
|
def vector_generation_is_current(self, revision: str) -> bool:
|
||
|
|
active = self.vectors.active_generation()
|
||
|
|
return self._matches_generation(active, revision) or self._failed_profile == (
|
||
|
|
revision,
|
||
|
|
self.cursor_profile(),
|
||
|
|
)
|
||
|
|
|
||
|
|
def build_vector_generation(
|
||
|
|
self,
|
||
|
|
connection: sqlite3.Connection,
|
||
|
|
revision: str,
|
||
|
|
nodes: tuple[KnowledgeNode, ...],
|
||
|
|
) -> None:
|
||
|
|
"""Stage vectors and their active pointer inside the canonical index transaction."""
|
||
|
|
|
||
|
|
ordered_nodes = tuple(sorted(nodes, key=lambda node: (node.project_id, node.node_id)))
|
||
|
|
generation_id = self._generation_id(revision)
|
||
|
|
active = self.vectors.active_generation_in_transaction(connection)
|
||
|
|
generation = self.vectors.generation_in_transaction(connection, generation_id)
|
||
|
|
if not self._matches_generation(generation, revision):
|
||
|
|
connection.execute("SAVEPOINT vector_generation_stage")
|
||
|
|
try:
|
||
|
|
if generation is None:
|
||
|
|
self.vectors.build_generation_in_transaction(
|
||
|
|
connection,
|
||
|
|
generation_id=generation_id,
|
||
|
|
revision=revision,
|
||
|
|
model_id=self.embeddings.model_id,
|
||
|
|
artifact_sha256=self.embeddings.artifact_sha256,
|
||
|
|
content_profile=_NODE_CONTENT_PROFILE,
|
||
|
|
dimensions=self.embeddings.dimensions,
|
||
|
|
vectors=self._vector_records(ordered_nodes),
|
||
|
|
)
|
||
|
|
generation = self.vectors.generation_in_transaction(connection, generation_id)
|
||
|
|
if not self._matches_generation(generation, revision):
|
||
|
|
raise RuntimeError("Vector generation metadata does not match the model")
|
||
|
|
self.vectors.activate_in_transaction(
|
||
|
|
connection,
|
||
|
|
generation_id,
|
||
|
|
expected_active=active.generation_id if active is not None else None,
|
||
|
|
)
|
||
|
|
connection.execute("RELEASE SAVEPOINT vector_generation_stage")
|
||
|
|
self._failed_profile = None
|
||
|
|
except (RuntimeError, ValueError, sqlite3.Error):
|
||
|
|
connection.execute("ROLLBACK TO SAVEPOINT vector_generation_stage")
|
||
|
|
connection.execute("RELEASE SAVEPOINT vector_generation_stage")
|
||
|
|
self.vectors.deactivate_in_transaction(connection)
|
||
|
|
self._failed_profile = (revision, self.cursor_profile())
|
||
|
|
return
|
||
|
|
elif active is None or active.generation_id != generation_id:
|
||
|
|
self.vectors.activate_in_transaction(
|
||
|
|
connection,
|
||
|
|
generation_id,
|
||
|
|
expected_active=active.generation_id if active is not None else None,
|
||
|
|
)
|
||
|
|
self._failed_profile = None
|
||
|
|
|
||
|
|
def _generation_id(self, revision: str) -> str:
|
||
|
|
return hashlib.sha256(
|
||
|
|
"\0".join(
|
||
|
|
(
|
||
|
|
revision,
|
||
|
|
self.embeddings.model_id,
|
||
|
|
self.embeddings.artifact_sha256,
|
||
|
|
str(self.embeddings.dimensions),
|
||
|
|
_NODE_CONTENT_PROFILE,
|
||
|
|
)
|
||
|
|
).encode()
|
||
|
|
).hexdigest()
|
||
|
|
|
||
|
|
def _vector_records(
|
||
|
|
self, nodes: tuple[KnowledgeNode, ...]
|
||
|
|
) -> Iterator[tuple[str, str, tuple[float, ...]]]:
|
||
|
|
for start in range(0, len(nodes), _VECTOR_BUILD_BATCH_SIZE):
|
||
|
|
batch = nodes[start : start + _VECTOR_BUILD_BATCH_SIZE]
|
||
|
|
texts = tuple(self._node_text(node) for node in batch)
|
||
|
|
embeddings = self.embeddings.embed_passages(texts)
|
||
|
|
yield from (
|
||
|
|
(node.project_id, node.node_id, vector)
|
||
|
|
for node, vector in zip(batch, embeddings, strict=True)
|
||
|
|
)
|
||
|
|
|
||
|
|
def _matches_generation(self, generation: VectorGeneration | None, revision: str) -> bool:
|
||
|
|
return bool(
|
||
|
|
generation is not None
|
||
|
|
and generation.revision == revision
|
||
|
|
and generation.model_id == self.embeddings.model_id
|
||
|
|
and generation.artifact_sha256 == self.embeddings.artifact_sha256
|
||
|
|
and generation.content_profile == _NODE_CONTENT_PROFILE
|
||
|
|
and generation.dimensions == self.embeddings.dimensions
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _node_text(node: KnowledgeNode) -> str:
|
||
|
|
return "\n".join(
|
||
|
|
(
|
||
|
|
node.title,
|
||
|
|
node.summary,
|
||
|
|
*(claim.text for claim in node.evidence_claims),
|
||
|
|
*node.tags,
|
||
|
|
*node.aliases,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def neighbors(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
node_id: str,
|
||
|
|
edge_types: tuple[str, ...],
|
||
|
|
project_ids: tuple[str, ...],
|
||
|
|
include_shared: bool,
|
||
|
|
limit: int,
|
||
|
|
offset: int = 0,
|
||
|
|
) -> SearchEdgeResult:
|
||
|
|
return self.lexical.neighbors(
|
||
|
|
node_id=node_id,
|
||
|
|
edge_types=edge_types,
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=limit,
|
||
|
|
offset=offset,
|
||
|
|
)
|
||
|
|
|
||
|
|
def list_sources(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
project_ids: tuple[str, ...],
|
||
|
|
include_shared: bool,
|
||
|
|
limit: int,
|
||
|
|
offset: int = 0,
|
||
|
|
) -> SourceIndexResult:
|
||
|
|
return self.lexical.list_sources(
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=include_shared,
|
||
|
|
limit=limit,
|
||
|
|
offset=offset,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _budgeted_envelope(
|
||
|
|
*,
|
||
|
|
request_id: str,
|
||
|
|
scope: Scope,
|
||
|
|
rows: list[dict[str, object]],
|
||
|
|
citation_groups: list[dict[str, EvidenceReference]],
|
||
|
|
warnings: list[Finding],
|
||
|
|
canonical_revision: str,
|
||
|
|
index_revision: str,
|
||
|
|
budget_tokens: int,
|
||
|
|
retrieval_mode: Literal["keyword", "hybrid", "canonical"],
|
||
|
|
next_cursor_for: Callable[[int], str | None] | None = None,
|
||
|
|
) -> ReadEnvelope:
|
||
|
|
for included_count in range(len(rows), -1, -1):
|
||
|
|
evidence: dict[str, EvidenceReference] = {}
|
||
|
|
for group in citation_groups[:included_count]:
|
||
|
|
evidence.update(group)
|
||
|
|
omitted = len(rows) - included_count
|
||
|
|
omissions = (
|
||
|
|
(
|
||
|
|
Finding(
|
||
|
|
code="BUDGET_TRUNCATED",
|
||
|
|
message=f"Omitted {omitted} lower-ranked result(s) to fit the budget.",
|
||
|
|
),
|
||
|
|
)
|
||
|
|
if omitted
|
||
|
|
else ()
|
||
|
|
)
|
||
|
|
envelope = ReadEnvelope(
|
||
|
|
request_id=request_id,
|
||
|
|
resolved_scope=scope,
|
||
|
|
data=tuple(rows[:included_count]),
|
||
|
|
citations=tuple(evidence[key] for key in sorted(evidence)),
|
||
|
|
warnings=tuple(warnings),
|
||
|
|
omissions=omissions,
|
||
|
|
next_cursor=next_cursor_for(included_count) if next_cursor_for is not None else None,
|
||
|
|
canonical_revision=canonical_revision,
|
||
|
|
index_revision=index_revision,
|
||
|
|
retrieval_mode=retrieval_mode,
|
||
|
|
budget_profile=BudgetProfile(
|
||
|
|
tokenizer="unicode-codepoints-v1", encoding="canonical-json-v1"
|
||
|
|
),
|
||
|
|
tokens_used=0,
|
||
|
|
)
|
||
|
|
envelope = CanonicalSearchService._count_complete_payload(envelope)
|
||
|
|
if envelope.tokens_used <= budget_tokens:
|
||
|
|
if len(canonical_json(envelope.model_dump(mode="json"))) > 1_000_000:
|
||
|
|
raise SearchError(413, "Search response exceeds the one-megabyte response budget")
|
||
|
|
return envelope
|
||
|
|
raise SearchError(413, "Search budget is too small for the required response envelope")
|
||
|
|
|
||
|
|
|
||
|
|
class CanonicalSearchService:
|
||
|
|
"""Authorize explicit project scope, then search a rebuildable index."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
index: SearchIndex,
|
||
|
|
project_ids_provider: Callable[[], Collection[str]],
|
||
|
|
canonical_head: Callable[[], str],
|
||
|
|
default_includes_provider: Callable[[str], tuple[str, ...]] | None = None,
|
||
|
|
) -> None:
|
||
|
|
self.index = index
|
||
|
|
self.project_ids_provider = project_ids_provider
|
||
|
|
self.canonical_head = canonical_head
|
||
|
|
self.default_includes_provider = default_includes_provider
|
||
|
|
|
||
|
|
def diagnostic_index_status(self) -> tuple[str, str | None, str, int | None]:
|
||
|
|
canonical_revision = self.canonical_head()
|
||
|
|
try:
|
||
|
|
index_revision, schema_version, available = self.index.diagnostic_status()
|
||
|
|
except (OSError, RuntimeError):
|
||
|
|
return canonical_revision, None, "unavailable", None
|
||
|
|
if not available:
|
||
|
|
return canonical_revision, None, "unavailable", schema_version
|
||
|
|
if index_revision is None:
|
||
|
|
return canonical_revision, None, "not_built", schema_version
|
||
|
|
state = "current" if index_revision == canonical_revision else "lagging"
|
||
|
|
return canonical_revision, index_revision, state, schema_version
|
||
|
|
|
||
|
|
def search(self, request: SearchRequest, *, principal: Principal) -> ReadEnvelope:
|
||
|
|
project_ids = self._resolve_scope(request.scope, principal)
|
||
|
|
fingerprint = hashlib.sha256(
|
||
|
|
canonical_json(
|
||
|
|
{
|
||
|
|
"query": request.query,
|
||
|
|
"projects": project_ids,
|
||
|
|
"include_shared": request.scope.include_shared,
|
||
|
|
"node_types": tuple(sorted(item.value for item in request.node_types)),
|
||
|
|
"tags": tuple(sorted(request.tags)),
|
||
|
|
"audience": tuple(sorted(request.audience)),
|
||
|
|
"use_for": tuple(sorted(request.use_for)),
|
||
|
|
"domain": tuple(sorted(request.domain)),
|
||
|
|
"confidences": tuple(sorted(item.value for item in request.confidences)),
|
||
|
|
"as_of_from": request.as_of_from.isoformat() if request.as_of_from else None,
|
||
|
|
"as_of_to": request.as_of_to.isoformat() if request.as_of_to else None,
|
||
|
|
"stale_only": request.stale_only,
|
||
|
|
"limit": request.limit,
|
||
|
|
"retrieval_profile": (
|
||
|
|
self.index.cursor_profile()
|
||
|
|
if isinstance(self.index, HybridSearchIndex)
|
||
|
|
else "keyword-v1"
|
||
|
|
),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
).hexdigest()
|
||
|
|
offset, expected_index_revision = self._cursor_offset(request.cursor, fingerprint)
|
||
|
|
|
||
|
|
result = self.index.search(
|
||
|
|
query=request.query,
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=request.scope.include_shared,
|
||
|
|
node_types=tuple(item.value for item in request.node_types),
|
||
|
|
tags=request.tags,
|
||
|
|
audience=request.audience,
|
||
|
|
use_for=request.use_for,
|
||
|
|
domain=request.domain,
|
||
|
|
confidences=request.confidences,
|
||
|
|
as_of_from=request.as_of_from,
|
||
|
|
as_of_to=request.as_of_to,
|
||
|
|
stale_only=request.stale_only,
|
||
|
|
limit=request.limit,
|
||
|
|
offset=offset,
|
||
|
|
)
|
||
|
|
if expected_index_revision is not None and expected_index_revision != result.index_revision:
|
||
|
|
raise SearchError(409, "Search cursor expired because the index revision changed")
|
||
|
|
|
||
|
|
canonical_revision = self.canonical_head()
|
||
|
|
warnings: list[Finding] = list(result.warnings)
|
||
|
|
now = datetime.now(UTC)
|
||
|
|
evidence_by_node: list[dict[str, EvidenceReference]] = []
|
||
|
|
rows: list[dict[str, object]] = []
|
||
|
|
for hit in result.nodes:
|
||
|
|
node = hit.node
|
||
|
|
row: dict[str, object] = {
|
||
|
|
"node": node.model_dump(mode="json"),
|
||
|
|
"retrieval_score": hit.score,
|
||
|
|
}
|
||
|
|
if hit.expanded_from is not None and hit.expansion_edge is not None:
|
||
|
|
row["retrieval_path"] = "graph_expansion"
|
||
|
|
row["expanded_from"] = hit.expanded_from
|
||
|
|
row["expansion_edge"] = hit.expansion_edge.model_dump(mode="json")
|
||
|
|
rows.append(row)
|
||
|
|
if node.stale_after is not None and node.stale_after <= now:
|
||
|
|
warnings.append(
|
||
|
|
Finding(
|
||
|
|
code="STALE_NODE",
|
||
|
|
message=f"Knowledge node {node.node_id} is past its freshness date.",
|
||
|
|
path=node.node_id,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
node_evidence: dict[str, EvidenceReference] = {}
|
||
|
|
for claim in node.evidence_claims:
|
||
|
|
for item in claim.evidence:
|
||
|
|
key = evidence_digest(item)
|
||
|
|
node_evidence[key] = item
|
||
|
|
if hit.expansion_edge is not None:
|
||
|
|
for item in hit.expansion_edge.evidence:
|
||
|
|
key = evidence_digest(item)
|
||
|
|
node_evidence[key] = item
|
||
|
|
evidence_by_node.append(node_evidence)
|
||
|
|
if result.index_revision != canonical_revision:
|
||
|
|
warnings.append(
|
||
|
|
Finding(
|
||
|
|
code="INDEX_LAG",
|
||
|
|
message=(
|
||
|
|
"Search results reflect an older accepted revision; "
|
||
|
|
"canonical reads remain available."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return _budgeted_envelope(
|
||
|
|
request_id="search-" + secrets.token_hex(12),
|
||
|
|
scope=Scope(projects=project_ids, include_shared=request.scope.include_shared),
|
||
|
|
rows=rows,
|
||
|
|
citation_groups=evidence_by_node,
|
||
|
|
warnings=warnings,
|
||
|
|
canonical_revision=canonical_revision,
|
||
|
|
index_revision=result.index_revision,
|
||
|
|
budget_tokens=request.budget_tokens,
|
||
|
|
retrieval_mode=result.retrieval_mode,
|
||
|
|
next_cursor_for=(
|
||
|
|
lambda included: (
|
||
|
|
(
|
||
|
|
self._encode_cursor(offset + included, fingerprint, result.index_revision)
|
||
|
|
if result.has_more or included < len(rows)
|
||
|
|
else None
|
||
|
|
)
|
||
|
|
if not result.warnings
|
||
|
|
else None
|
||
|
|
)
|
||
|
|
),
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _encode_cursor(offset: int, fingerprint: str, index_revision: str) -> str:
|
||
|
|
payload = canonical_json(
|
||
|
|
{"offset": offset, "fingerprint": fingerprint, "index_revision": index_revision}
|
||
|
|
)
|
||
|
|
return base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=")
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _cursor_offset(
|
||
|
|
cursor: str | None, fingerprint: str, *, max_offset: int = 100_000
|
||
|
|
) -> tuple[int, str | None]:
|
||
|
|
if cursor is None:
|
||
|
|
return 0, None
|
||
|
|
try:
|
||
|
|
padded = cursor + "=" * (-len(cursor) % 4)
|
||
|
|
payload = json.loads(base64.b64decode(padded, altchars=b"-_", validate=True))
|
||
|
|
offset = payload["offset"]
|
||
|
|
if (
|
||
|
|
not isinstance(payload, dict)
|
||
|
|
or set(payload) != {"offset", "fingerprint", "index_revision"}
|
||
|
|
or type(offset) is not int
|
||
|
|
or not 0 <= offset <= max_offset
|
||
|
|
or payload["fingerprint"] != fingerprint
|
||
|
|
or not isinstance(payload["index_revision"], str)
|
||
|
|
or len(payload["index_revision"]) not in {40, 64}
|
||
|
|
or any(
|
||
|
|
character not in "0123456789abcdef" for character in payload["index_revision"]
|
||
|
|
)
|
||
|
|
):
|
||
|
|
raise ValueError("cursor does not match this search")
|
||
|
|
return offset, payload["index_revision"]
|
||
|
|
except (ValueError, KeyError, TypeError, json.JSONDecodeError) as error:
|
||
|
|
raise SearchError(422, "Search cursor is invalid for this query and scope") from error
|
||
|
|
|
||
|
|
def list_projects(self, *, principal: Principal) -> tuple[str, ...]:
|
||
|
|
"""List only projects visible to the authenticated principal."""
|
||
|
|
|
||
|
|
return tuple(
|
||
|
|
sorted(
|
||
|
|
project_id
|
||
|
|
for project_id in set(self.project_ids_provider())
|
||
|
|
if principal.can_read_project(project_id)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def _resolve_scope(self, scope: Scope, principal: Principal) -> tuple[str, ...]:
|
||
|
|
registered_projects = set(self.project_ids_provider())
|
||
|
|
if scope.all_projects:
|
||
|
|
if "admin" not in principal.roles:
|
||
|
|
raise SearchError(403, "All-project search requires an administrator")
|
||
|
|
project_ids = tuple(sorted(registered_projects))
|
||
|
|
else:
|
||
|
|
if scope.include_projects is not None:
|
||
|
|
includes = scope.include_projects
|
||
|
|
elif len(scope.projects) == 1 and self.default_includes_provider is not None:
|
||
|
|
includes = self.default_includes_provider(scope.projects[0])
|
||
|
|
else:
|
||
|
|
includes = ()
|
||
|
|
includes = tuple(project_id for project_id in includes if project_id != "shared")
|
||
|
|
project_ids = tuple(dict.fromkeys((*scope.projects, *includes)))
|
||
|
|
unknown = set(project_ids) - registered_projects - {"shared"}
|
||
|
|
if unknown:
|
||
|
|
raise SearchError(404, "Search scope contains an unknown project")
|
||
|
|
if not project_ids:
|
||
|
|
raise SearchError(422, "Search scope resolved to no projects")
|
||
|
|
if any(not principal.can_read_project(project_id) for project_id in project_ids):
|
||
|
|
raise SearchError(403, "Credential cannot read every requested project")
|
||
|
|
if scope.include_shared and not principal.can_read_project("shared"):
|
||
|
|
raise SearchError(403, "Credential cannot read shared knowledge")
|
||
|
|
return project_ids
|
||
|
|
|
||
|
|
def neighbors(self, request: NeighborRequest, *, principal: Principal) -> ReadEnvelope:
|
||
|
|
project_ids = self._resolve_scope(request.scope, principal)
|
||
|
|
fingerprint = hashlib.sha256(
|
||
|
|
canonical_json(
|
||
|
|
{
|
||
|
|
"node_id": request.node_id,
|
||
|
|
"projects": project_ids,
|
||
|
|
"include_shared": request.scope.include_shared,
|
||
|
|
"edge_types": request.edge_types,
|
||
|
|
"limit": request.limit,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
).hexdigest()
|
||
|
|
offset, expected_index_revision = self._cursor_offset(
|
||
|
|
request.cursor, fingerprint, max_offset=500_000
|
||
|
|
)
|
||
|
|
result = self.index.neighbors(
|
||
|
|
node_id=request.node_id,
|
||
|
|
edge_types=tuple(request.edge_types),
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=request.scope.include_shared,
|
||
|
|
limit=request.limit,
|
||
|
|
offset=offset,
|
||
|
|
)
|
||
|
|
if expected_index_revision is not None and expected_index_revision != result.index_revision:
|
||
|
|
raise SearchError(409, "Neighbor cursor expired because the index revision changed")
|
||
|
|
origin = result.origin
|
||
|
|
allowed = set(project_ids)
|
||
|
|
if request.scope.include_shared:
|
||
|
|
allowed.add("shared")
|
||
|
|
if origin is None or origin.project_id not in allowed:
|
||
|
|
raise SearchError(404, "Knowledge node was not found within the requested scope")
|
||
|
|
|
||
|
|
canonical_revision = self.canonical_head()
|
||
|
|
warnings: list[Finding] = []
|
||
|
|
stale_ids: set[str] = set()
|
||
|
|
if origin.stale_after is not None and origin.stale_after <= datetime.now(UTC):
|
||
|
|
stale_ids.add(origin.node_id)
|
||
|
|
rows: list[dict[str, object]] = []
|
||
|
|
citation_groups: list[dict[str, EvidenceReference]] = []
|
||
|
|
for hit in result.edges:
|
||
|
|
endpoint = hit.neighbor_endpoint
|
||
|
|
if endpoint.kind == "node":
|
||
|
|
neighbor = hit.neighbor_node
|
||
|
|
if neighbor is None or neighbor.project_id not in allowed:
|
||
|
|
continue
|
||
|
|
neighbor_value: dict[str, object] = {
|
||
|
|
"kind": "node",
|
||
|
|
"id": endpoint.id,
|
||
|
|
"record": neighbor.model_dump(mode="json"),
|
||
|
|
}
|
||
|
|
if neighbor.stale_after is not None and neighbor.stale_after <= datetime.now(UTC):
|
||
|
|
stale_ids.add(neighbor.node_id)
|
||
|
|
elif endpoint.kind == "project":
|
||
|
|
if endpoint.id not in allowed:
|
||
|
|
continue
|
||
|
|
neighbor_value = {"kind": "project", "id": endpoint.id}
|
||
|
|
neighbor = None
|
||
|
|
else:
|
||
|
|
source = hit.neighbor_source
|
||
|
|
if source is None or not (set(source.ingested_for) & allowed):
|
||
|
|
continue
|
||
|
|
neighbor_value = {
|
||
|
|
"kind": "source",
|
||
|
|
"id": endpoint.id,
|
||
|
|
"record": source.model_dump(mode="json"),
|
||
|
|
}
|
||
|
|
neighbor = None
|
||
|
|
rows.append(
|
||
|
|
{
|
||
|
|
"from_node": {
|
||
|
|
"node_id": origin.node_id,
|
||
|
|
"project_id": origin.project_id,
|
||
|
|
"title": origin.title,
|
||
|
|
},
|
||
|
|
"edge": hit.edge.model_dump(mode="json"),
|
||
|
|
"neighbor": neighbor_value,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
group: dict[str, EvidenceReference] = {}
|
||
|
|
for item in hit.edge.evidence:
|
||
|
|
group[evidence_digest(item)] = item
|
||
|
|
if neighbor is not None:
|
||
|
|
for claim in neighbor.claims:
|
||
|
|
for item in claim.evidence:
|
||
|
|
group[evidence_digest(item)] = item
|
||
|
|
citation_groups.append(group)
|
||
|
|
warnings.extend(
|
||
|
|
Finding(
|
||
|
|
code="STALE_NODE",
|
||
|
|
message=f"Knowledge node {node_id} is past its freshness date.",
|
||
|
|
path=node_id,
|
||
|
|
)
|
||
|
|
for node_id in sorted(stale_ids)
|
||
|
|
)
|
||
|
|
if result.index_revision != canonical_revision:
|
||
|
|
warnings.append(
|
||
|
|
Finding(
|
||
|
|
code="INDEX_LAG",
|
||
|
|
message=(
|
||
|
|
"Neighbor results reflect an older accepted revision; "
|
||
|
|
"canonical reads remain available."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return _budgeted_envelope(
|
||
|
|
request_id="neighbors-" + secrets.token_hex(12),
|
||
|
|
scope=Scope(projects=project_ids, include_shared=request.scope.include_shared),
|
||
|
|
rows=rows,
|
||
|
|
citation_groups=citation_groups,
|
||
|
|
warnings=warnings,
|
||
|
|
canonical_revision=canonical_revision,
|
||
|
|
index_revision=result.index_revision,
|
||
|
|
budget_tokens=request.budget_tokens,
|
||
|
|
retrieval_mode="keyword",
|
||
|
|
next_cursor_for=(
|
||
|
|
lambda included: (
|
||
|
|
self._encode_cursor(offset + included, fingerprint, result.index_revision)
|
||
|
|
if result.has_more or included < len(rows)
|
||
|
|
else None
|
||
|
|
)
|
||
|
|
),
|
||
|
|
)
|
||
|
|
|
||
|
|
def list_sources(self, request: SourceListRequest, *, principal: Principal) -> ReadEnvelope:
|
||
|
|
project_ids = self._resolve_scope(request.scope, principal)
|
||
|
|
fingerprint = hashlib.sha256(
|
||
|
|
canonical_json(
|
||
|
|
{
|
||
|
|
"projects": project_ids,
|
||
|
|
"include_shared": request.scope.include_shared,
|
||
|
|
"limit": request.limit,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
).hexdigest()
|
||
|
|
offset, expected_index_revision = self._cursor_offset(
|
||
|
|
request.cursor, fingerprint, max_offset=500_000
|
||
|
|
)
|
||
|
|
result = self.index.list_sources(
|
||
|
|
project_ids=project_ids,
|
||
|
|
include_shared=request.scope.include_shared,
|
||
|
|
limit=request.limit,
|
||
|
|
offset=offset,
|
||
|
|
)
|
||
|
|
if expected_index_revision is not None and expected_index_revision != result.index_revision:
|
||
|
|
raise SearchError(409, "Source cursor expired because the index revision changed")
|
||
|
|
|
||
|
|
canonical_revision = self.canonical_head()
|
||
|
|
rows: list[dict[str, object]] = [
|
||
|
|
{"source": source.model_dump(mode="json")} for source in result.sources
|
||
|
|
]
|
||
|
|
warnings: list[Finding] = []
|
||
|
|
if result.index_revision != canonical_revision:
|
||
|
|
warnings.append(
|
||
|
|
Finding(
|
||
|
|
code="INDEX_LAG",
|
||
|
|
message=(
|
||
|
|
"Source results reflect an older accepted revision; "
|
||
|
|
"canonical source reads remain available."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return _budgeted_envelope(
|
||
|
|
request_id="sources-" + secrets.token_hex(12),
|
||
|
|
scope=Scope(projects=project_ids, include_shared=request.scope.include_shared),
|
||
|
|
rows=rows,
|
||
|
|
citation_groups=[{} for _ in rows],
|
||
|
|
warnings=warnings,
|
||
|
|
canonical_revision=canonical_revision,
|
||
|
|
index_revision=result.index_revision,
|
||
|
|
budget_tokens=request.budget_tokens,
|
||
|
|
retrieval_mode="keyword",
|
||
|
|
next_cursor_for=(
|
||
|
|
lambda included: (
|
||
|
|
self._encode_cursor(offset + included, fingerprint, result.index_revision)
|
||
|
|
if result.has_more or included < len(rows)
|
||
|
|
else None
|
||
|
|
)
|
||
|
|
),
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _count_complete_payload(envelope: ReadEnvelope) -> ReadEnvelope:
|
||
|
|
counted = envelope
|
||
|
|
for _ in range(4):
|
||
|
|
serialized = canonical_json(counted.model_dump(mode="json"))
|
||
|
|
count = len(serialized.decode("utf-8"))
|
||
|
|
if count == counted.tokens_used:
|
||
|
|
return counted
|
||
|
|
counted = counted.model_copy(update={"tokens_used": count})
|
||
|
|
return counted
|