"""Governed retrieval contracts over the asset registry.""" from __future__ import annotations from dataclasses import dataclass, field from time import perf_counter from typing import Any from kontextual_engine.core import ( AuditEvent, AuditOutcome, AssetRepresentation, ContextEntity, ContextEntityType, CoreRelationship, KnowledgeAsset, LifecycleState, MetadataRecord, OperationContext, PolicyDecision, RelationshipTargetKind, RepresentationKind, RetrievalFeedbackLabel, RetrievalFeedbackRecord, Sensitivity, ) from kontextual_engine.errors import Diagnostic from kontextual_engine.ports import AllowAllPolicyGateway, AssetRegistryRepository, PolicyGateway SUPPORTED_SORT_KEYS = { "asset_id", "asset_type", "created_at", "lifecycle", "title", "updated_at", } SUPPORTED_CONTEXT_ENTITY_SORT_KEYS = { "entity_id", "entity_type", "external_ref", "name", } SUPPORTED_RELATIONSHIP_SORT_KEYS = { "created_at", "predicate", "relationship_id", "source_id", "target_id", "target_kind", } @dataclass(frozen=True) class AssetQueryRequest: text: str | None = None asset_type: str | None = None lifecycle: LifecycleState | str | None = None sensitivity: Sensitivity | str | None = None owner: str | None = None topic: str | None = None tags: tuple[str, ...] = () collection: str | None = None review_state: str | None = None metadata_filters: dict[str, Any] = field(default_factory=dict) confirmed_metadata_only: bool = False source_system: str | None = None source_path: str | None = None context_entity_id: str | None = None context_entity_type: ContextEntityType | str | None = None context_entity_name: str | None = None context_entity_external_ref: str | None = None workflow_run_id: str | None = None related_asset_id: str | None = None relationship_predicate: str | None = None relationship_direction: str = "both" include_relationships: bool = False include_snippets: bool = False max_snippets: int = 3 snippet_radius: int = 80 created_after: str | None = None created_before: str | None = None updated_after: str | None = None updated_before: str | None = None representation_kind: RepresentationKind | str | None = None sort_by: str = "title" sort_order: str = "asc" limit: int = 50 offset: int = 0 def to_dict(self) -> dict[str, Any]: return { "text": self.text, "asset_type": self.asset_type, "lifecycle": self.lifecycle.value if isinstance(self.lifecycle, LifecycleState) else self.lifecycle, "sensitivity": self.sensitivity.value if isinstance(self.sensitivity, Sensitivity) else self.sensitivity, "owner": self.owner, "topic": self.topic, "tags": list(self.tags), "collection": self.collection, "review_state": self.review_state, "metadata_filters": dict(self.metadata_filters), "confirmed_metadata_only": self.confirmed_metadata_only, "source_system": self.source_system, "source_path": self.source_path, "context_entity_id": self.context_entity_id, "context_entity_type": self.context_entity_type.value if isinstance(self.context_entity_type, ContextEntityType) else self.context_entity_type, "context_entity_name": self.context_entity_name, "context_entity_external_ref": self.context_entity_external_ref, "workflow_run_id": self.workflow_run_id, "related_asset_id": self.related_asset_id, "relationship_predicate": self.relationship_predicate, "relationship_direction": self.relationship_direction, "include_relationships": self.include_relationships, "include_snippets": self.include_snippets, "max_snippets": self.max_snippets, "snippet_radius": self.snippet_radius, "created_after": self.created_after, "created_before": self.created_before, "updated_after": self.updated_after, "updated_before": self.updated_before, "representation_kind": self.representation_kind.value if isinstance(self.representation_kind, RepresentationKind) else self.representation_kind, "sort_by": self.sort_by, "sort_order": self.sort_order, "limit": self.limit, "offset": self.offset, } @dataclass(frozen=True) class AssetQueryItem: asset: KnowledgeAsset representations: tuple[AssetRepresentation, ...] = () metadata_records: tuple[MetadataRecord, ...] = () relationships: tuple[CoreRelationship, ...] = () context_entities: tuple[ContextEntity, ...] = () snippets: tuple["RetrievalSnippet", ...] = () diagnostics: tuple[Diagnostic, ...] = () relevance: dict[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: object.__setattr__(self, "representations", tuple(self.representations)) object.__setattr__(self, "metadata_records", tuple(self.metadata_records)) object.__setattr__(self, "relationships", tuple(self.relationships)) object.__setattr__(self, "context_entities", tuple(self.context_entities)) object.__setattr__(self, "snippets", tuple(self.snippets)) object.__setattr__(self, "diagnostics", tuple(self.diagnostics)) def to_dict(self) -> dict[str, Any]: return { "asset_id": self.asset.id, "title": self.asset.title, "lifecycle": self.asset.lifecycle.value, "classification": self.asset.classification.to_dict(), "current_version_id": self.asset.current_version_id, "aliases": list(self.asset.aliases), "asset_metadata": dict(self.asset.metadata), "source_refs": [source_ref.to_dict() for source_ref in self.asset.source_refs], "representations": [representation.to_dict() for representation in self.representations], "metadata_records": [record.to_dict() for record in self.metadata_records], "relationships": [relationship.to_dict() for relationship in self.relationships], "context_entities": [entity.to_dict() for entity in self.context_entities], "snippets": [snippet.to_dict() for snippet in self.snippets], "relevance": dict(self.relevance), "diagnostics": [diagnostic.to_dict() for diagnostic in self.diagnostics], } @dataclass(frozen=True) class AssetQueryResult: request: AssetQueryRequest correlation_id: str total: int items: tuple[AssetQueryItem, ...] = () diagnostics: tuple[Diagnostic, ...] = () metadata: dict[str, Any] = field(default_factory=dict) success: bool = True def __post_init__(self) -> None: object.__setattr__(self, "items", tuple(self.items)) object.__setattr__(self, "diagnostics", tuple(self.diagnostics)) @property def result_count(self) -> int: return len(self.items) @property def next_offset(self) -> int | None: next_offset = self.request.offset + self.result_count return next_offset if next_offset < self.total else None def to_dict(self) -> dict[str, Any]: return { "query": self.request.to_dict(), "correlation_id": self.correlation_id, "success": self.success, "total": self.total, "result_count": self.result_count, "limit": self.request.limit, "offset": self.request.offset, "next_offset": self.next_offset, "sort": { "by": self.request.sort_by, "order": self.request.sort_order, }, "metadata": dict(self.metadata), "results": [item.to_dict() for item in self.items], "diagnostics": [diagnostic.to_dict() for diagnostic in self.diagnostics], } @dataclass(frozen=True) class LexicalIndexRefreshResult: indexed_assets: int indexed_representations: int def to_dict(self) -> dict[str, int]: return { "indexed_assets": self.indexed_assets, "indexed_representations": self.indexed_representations, } @dataclass(frozen=True) class RetrievalSnippet: asset_id: str representation_id: str text: str start_offset: int end_offset: int match_text: str media_type: str source_ref_id: str | None = None storage_ref: str | None = None provenance: dict[str, Any] = field(default_factory=dict) def to_dict(self) -> dict[str, Any]: return { "asset_id": self.asset_id, "representation_id": self.representation_id, "source_ref_id": self.source_ref_id, "storage_ref": self.storage_ref, "media_type": self.media_type, "text": self.text, "start_offset": self.start_offset, "end_offset": self.end_offset, "match_text": self.match_text, "provenance": dict(self.provenance), } @dataclass(frozen=True) class ContextEntityQueryRequest: entity_id: str | None = None entity_type: ContextEntityType | str | None = None name: str | None = None external_ref: str | None = None metadata_filters: dict[str, Any] = field(default_factory=dict) sort_by: str = "name" sort_order: str = "asc" limit: int = 50 offset: int = 0 def to_dict(self) -> dict[str, Any]: return { "entity_id": self.entity_id, "entity_type": self.entity_type.value if isinstance(self.entity_type, ContextEntityType) else self.entity_type, "name": self.name, "external_ref": self.external_ref, "metadata_filters": dict(self.metadata_filters), "sort_by": self.sort_by, "sort_order": self.sort_order, "limit": self.limit, "offset": self.offset, } @dataclass(frozen=True) class ContextEntityQueryItem: entity: ContextEntity asset_ids: tuple[str, ...] = () relationship_count: int = 0 def __post_init__(self) -> None: object.__setattr__(self, "asset_ids", tuple(self.asset_ids)) def to_dict(self) -> dict[str, Any]: return { **self.entity.to_dict(), "asset_ids": list(self.asset_ids), "relationship_count": self.relationship_count, } @dataclass(frozen=True) class ContextEntityQueryResult: request: ContextEntityQueryRequest correlation_id: str total: int items: tuple[ContextEntityQueryItem, ...] = () diagnostics: tuple[Diagnostic, ...] = () success: bool = True def __post_init__(self) -> None: object.__setattr__(self, "items", tuple(self.items)) object.__setattr__(self, "diagnostics", tuple(self.diagnostics)) @property def result_count(self) -> int: return len(self.items) @property def next_offset(self) -> int | None: next_offset = self.request.offset + self.result_count return next_offset if next_offset < self.total else None def to_dict(self) -> dict[str, Any]: return { "query": self.request.to_dict(), "correlation_id": self.correlation_id, "success": self.success, "total": self.total, "result_count": self.result_count, "limit": self.request.limit, "offset": self.request.offset, "next_offset": self.next_offset, "sort": { "by": self.request.sort_by, "order": self.request.sort_order, }, "results": [item.to_dict() for item in self.items], "diagnostics": [diagnostic.to_dict() for diagnostic in self.diagnostics], } @dataclass(frozen=True) class RelationshipQueryRequest: source_id: str | None = None target_id: str | None = None asset_id: str | None = None context_entity_id: str | None = None context_entity_type: ContextEntityType | str | None = None context_entity_name: str | None = None context_entity_external_ref: str | None = None workflow_run_id: str | None = None target_kind: RelationshipTargetKind | str | None = None predicate: str | None = None direction: str = "both" sort_by: str = "source_id" sort_order: str = "asc" limit: int = 50 offset: int = 0 def to_dict(self) -> dict[str, Any]: return { "source_id": self.source_id, "target_id": self.target_id, "asset_id": self.asset_id, "context_entity_id": self.context_entity_id, "context_entity_type": self.context_entity_type.value if isinstance(self.context_entity_type, ContextEntityType) else self.context_entity_type, "context_entity_name": self.context_entity_name, "context_entity_external_ref": self.context_entity_external_ref, "workflow_run_id": self.workflow_run_id, "target_kind": self.target_kind.value if isinstance(self.target_kind, RelationshipTargetKind) else self.target_kind, "predicate": self.predicate, "direction": self.direction, "sort_by": self.sort_by, "sort_order": self.sort_order, "limit": self.limit, "offset": self.offset, } @dataclass(frozen=True) class RelationshipQueryItem: relationship: CoreRelationship source_asset: KnowledgeAsset | None = None target_asset: KnowledgeAsset | None = None target_entity: ContextEntity | None = None def to_dict(self) -> dict[str, Any]: payload = self.relationship.to_dict() if self.source_asset is not None: payload["source_asset"] = self.source_asset.to_dict() if self.target_asset is not None: payload["target_asset"] = self.target_asset.to_dict() if self.target_entity is not None: payload["target_entity"] = self.target_entity.to_dict() return payload @dataclass(frozen=True) class RelationshipQueryResult: request: RelationshipQueryRequest correlation_id: str total: int items: tuple[RelationshipQueryItem, ...] = () diagnostics: tuple[Diagnostic, ...] = () success: bool = True def __post_init__(self) -> None: object.__setattr__(self, "items", tuple(self.items)) object.__setattr__(self, "diagnostics", tuple(self.diagnostics)) @property def result_count(self) -> int: return len(self.items) @property def next_offset(self) -> int | None: next_offset = self.request.offset + self.result_count return next_offset if next_offset < self.total else None def to_dict(self) -> dict[str, Any]: return { "query": self.request.to_dict(), "correlation_id": self.correlation_id, "success": self.success, "total": self.total, "result_count": self.result_count, "limit": self.request.limit, "offset": self.request.offset, "next_offset": self.next_offset, "sort": { "by": self.request.sort_by, "order": self.request.sort_order, }, "results": [item.to_dict() for item in self.items], "diagnostics": [diagnostic.to_dict() for diagnostic in self.diagnostics], } @dataclass(frozen=True) class RetrievalFeedbackRequest: label: RetrievalFeedbackLabel | str query: dict[str, Any] result_ref: dict[str, Any] = field(default_factory=dict) notes: str | None = None metadata: dict[str, Any] = field(default_factory=dict) def to_dict(self) -> dict[str, Any]: return { "label": self.label.value if isinstance(self.label, RetrievalFeedbackLabel) else self.label, "query": dict(self.query), "result_ref": dict(self.result_ref), "notes": self.notes, "metadata": dict(self.metadata), } @dataclass(frozen=True) class RetrievalFeedbackResult: record: RetrievalFeedbackRecord | None correlation_id: str diagnostics: tuple[Diagnostic, ...] = () success: bool = True def __post_init__(self) -> None: object.__setattr__(self, "diagnostics", tuple(self.diagnostics)) def to_dict(self) -> dict[str, Any]: return { "success": self.success, "correlation_id": self.correlation_id, "record": self.record.to_dict() if self.record else None, "diagnostics": [diagnostic.to_dict() for diagnostic in self.diagnostics], } @dataclass(frozen=True) class RetrievalQualityMetrics: query_count: int zero_result_count: int zero_result_rate: float feedback_count: int useful_count: int unsafe_count: int low_confidence_count: int precision_at_k: float | None citation_precision: float | None permission_filter_observation_count: int average_permission_filter_duration_ms: float | None def to_dict(self) -> dict[str, Any]: return { "query_count": self.query_count, "zero_result_count": self.zero_result_count, "zero_result_rate": self.zero_result_rate, "feedback_count": self.feedback_count, "useful_count": self.useful_count, "unsafe_count": self.unsafe_count, "low_confidence_count": self.low_confidence_count, "precision_at_k": self.precision_at_k, "citation_precision": self.citation_precision, "permission_filter_observation_count": self.permission_filter_observation_count, "average_permission_filter_duration_ms": self.average_permission_filter_duration_ms, } @dataclass(frozen=True) class _LexicalDocument: asset_id: str representation_id: str text: str media_type: str source_ref_id: str | None = None storage_ref: str | None = None provenance: dict[str, Any] = field(default_factory=dict) class AssetRetrievalService: def __init__( self, repository: AssetRegistryRepository, *, policy_gateway: PolicyGateway | None = None, ) -> None: self.repository = repository self.policy_gateway = policy_gateway or AllowAllPolicyGateway() self._lexical_index: tuple[_LexicalDocument, ...] = () self._last_refresh = LexicalIndexRefreshResult(indexed_assets=0, indexed_representations=0) def refresh_index(self) -> LexicalIndexRefreshResult: documents: list[_LexicalDocument] = [] for asset in self.repository.list_assets(): for representation in self.repository.list_representations( asset_id=asset.id, kind=RepresentationKind.NORMALIZED, ): text = representation.metadata.get("search_text") if not isinstance(text, str) or not text: continue documents.append( _LexicalDocument( asset_id=asset.id, representation_id=representation.representation_id, text=text, media_type=representation.media_type, source_ref_id=representation.source_ref_id, storage_ref=representation.storage_ref, provenance=_snippet_provenance(representation), ) ) self._lexical_index = tuple(documents) self._last_refresh = LexicalIndexRefreshResult( indexed_assets=len({document.asset_id for document in documents}), indexed_representations=len(documents), ) return self._last_refresh def query_assets( self, request: AssetQueryRequest, context: OperationContext, ) -> AssetQueryResult: diagnostics, normalized = _validate_request(request) if diagnostics: return AssetQueryResult( request=request, correlation_id=context.correlation_id, total=0, diagnostics=tuple(diagnostics), success=False, ) scope_decision = self._authorize_for_retrieval( context, "retrieval.assets.query", "retrieval:assets", resource_metadata={"query": normalized.to_dict()}, ) if not scope_decision.allowed: self._audit_retrieval( "retrieval.assets.query", "retrieval:assets", AuditOutcome.DENIED, context, scope_decision, details={"query": normalized.to_dict()}, ) return AssetQueryResult( request=normalized, correlation_id=context.correlation_id, total=0, diagnostics=( _permission_denied_diagnostic(scope_decision), ), success=False, ) assets = self.repository.list_assets( lifecycle=normalized.lifecycle, asset_type=normalized.asset_type, sensitivity=normalized.sensitivity, owner=normalized.owner, topic=normalized.topic, review_state=normalized.review_state, metadata_filters=normalized.metadata_filters or None, confirmed_metadata_only=normalized.confirmed_metadata_only, ) assets = [ asset for asset in assets if _source_matches( asset, source_system=normalized.source_system, source_path=normalized.source_path, ) ] assets = [ asset for asset in assets if _collection_matches( asset, self.repository.list_metadata_records(asset.id), normalized.collection, ) and _tags_match( asset, self.repository.list_metadata_records(asset.id), normalized.tags, ) and _timestamp_matches(asset, normalized) ] if normalized.representation_kind is not None: assets = [ asset for asset in assets if self.repository.list_representations( asset_id=asset.id, kind=normalized.representation_kind, ) ] relationship_context_by_asset: dict[str, tuple[CoreRelationship, ...]] = {} context_entities_by_asset: dict[str, tuple[ContextEntity, ...]] = {} if _asset_query_has_graph_filter(normalized): graph_matches: list[KnowledgeAsset] = [] for asset in assets: relationships, entities = self._relationship_context_for_asset(asset.id, normalized, context) if relationships: graph_matches.append(asset) relationship_context_by_asset[asset.id] = relationships context_entities_by_asset[asset.id] = entities assets = graph_matches relevance_by_asset: dict[str, dict[str, Any]] = {} snippets_by_asset: dict[str, tuple[RetrievalSnippet, ...]] = {} if normalized.text: relevance_by_asset, snippets_by_asset = self._lexical_matches( normalized.text, include_snippets=normalized.include_snippets, max_snippets=normalized.max_snippets, snippet_radius=normalized.snippet_radius, ) assets = [asset for asset in assets if asset.id in relevance_by_asset] permission_filter_started = perf_counter() asset_count_before_policy = len(assets) assets = [asset for asset in assets if self._asset_allowed(asset, context)] permission_filtered_count = asset_count_before_policy - len(assets) permission_filter_duration_ms = _elapsed_ms(permission_filter_started) ordered = _sort_assets(assets, normalized.sort_by, normalized.sort_order) total = len(ordered) page = ordered[normalized.offset : normalized.offset + normalized.limit] items = tuple( self._item_for_asset( asset, normalized.representation_kind, relevance=relevance_by_asset.get(asset.id, {}), include_relationships=normalized.include_relationships, relationships=relationship_context_by_asset.get(asset.id), context_entities=context_entities_by_asset.get(asset.id), snippets=snippets_by_asset.get(asset.id, ()), context=context, ) for asset in page ) result = AssetQueryResult( request=normalized, correlation_id=context.correlation_id, total=total, items=items, metadata={ "zero_result": total == 0, "graph_filter": _asset_query_has_graph_filter(normalized), "policy_enforced": True, "lexical_index": self._last_refresh.to_dict(), }, ) self._audit_retrieval( "retrieval.assets.query", "retrieval:assets", AuditOutcome.PARTIAL if permission_filtered_count else AuditOutcome.SUCCESS, context, scope_decision, details={ "query": normalized.to_dict(), "result_count": result.result_count, "total": result.total, "permission_filtered_count": permission_filtered_count, "permission_filter_duration_ms": permission_filter_duration_ms, }, ) return result def query_context_entities( self, request: ContextEntityQueryRequest, context: OperationContext, ) -> ContextEntityQueryResult: diagnostics, normalized = _validate_context_entity_request(request) if diagnostics: return ContextEntityQueryResult( request=request, correlation_id=context.correlation_id, total=0, diagnostics=tuple(diagnostics), success=False, ) scope_decision = self._authorize_for_retrieval( context, "retrieval.context_entities.query", "retrieval:context_entities", resource_metadata={"query": normalized.to_dict()}, ) if not scope_decision.allowed: self._audit_retrieval( "retrieval.context_entities.query", "retrieval:context_entities", AuditOutcome.DENIED, context, scope_decision, details={"query": normalized.to_dict()}, ) return ContextEntityQueryResult( request=normalized, correlation_id=context.correlation_id, total=0, diagnostics=( _permission_denied_diagnostic(scope_decision), ), success=False, ) entities = [ entity for entity in self.repository.list_context_entities() if _context_entity_matches( entity, entity_id=normalized.entity_id, entity_type=normalized.entity_type, name=normalized.name, external_ref=normalized.external_ref, metadata_filters=normalized.metadata_filters, ) ] permission_filter_started = perf_counter() entity_count_before_policy = len(entities) entities = [entity for entity in entities if self._context_entity_allowed(entity, context)] permission_filtered_count = entity_count_before_policy - len(entities) permission_filter_duration_ms = _elapsed_ms(permission_filter_started) ordered = _sort_context_entities(entities, normalized.sort_by, normalized.sort_order) total = len(ordered) page = ordered[normalized.offset : normalized.offset + normalized.limit] items = tuple(self._context_entity_item(entity, context) for entity in page) result = ContextEntityQueryResult( request=normalized, correlation_id=context.correlation_id, total=total, items=items, ) self._audit_retrieval( "retrieval.context_entities.query", "retrieval:context_entities", AuditOutcome.PARTIAL if permission_filtered_count else AuditOutcome.SUCCESS, context, scope_decision, details={ "query": normalized.to_dict(), "result_count": result.result_count, "total": result.total, "permission_filtered_count": permission_filtered_count, "permission_filter_duration_ms": permission_filter_duration_ms, }, ) return result def query_relationships( self, request: RelationshipQueryRequest, context: OperationContext, ) -> RelationshipQueryResult: diagnostics, normalized = _validate_relationship_request(request) if diagnostics: return RelationshipQueryResult( request=request, correlation_id=context.correlation_id, total=0, diagnostics=tuple(diagnostics), success=False, ) scope_decision = self._authorize_for_retrieval( context, "retrieval.relationships.query", "retrieval:relationships", resource_metadata={"query": normalized.to_dict()}, ) if not scope_decision.allowed: self._audit_retrieval( "retrieval.relationships.query", "retrieval:relationships", AuditOutcome.DENIED, context, scope_decision, details={"query": normalized.to_dict()}, ) return RelationshipQueryResult( request=normalized, correlation_id=context.correlation_id, total=0, diagnostics=( _permission_denied_diagnostic(scope_decision), ), success=False, ) relationships = self._relationships_for_request(normalized) permission_filter_started = perf_counter() relationship_count_before_policy = len(relationships) relationships = [ relationship for relationship in relationships if self._relationship_allowed(relationship, context) ] permission_filtered_count = relationship_count_before_policy - len(relationships) permission_filter_duration_ms = _elapsed_ms(permission_filter_started) ordered = _sort_relationships(relationships, normalized.sort_by, normalized.sort_order) total = len(ordered) page = ordered[normalized.offset : normalized.offset + normalized.limit] items = tuple(self._relationship_item(relationship) for relationship in page) result = RelationshipQueryResult( request=normalized, correlation_id=context.correlation_id, total=total, items=items, ) self._audit_retrieval( "retrieval.relationships.query", "retrieval:relationships", AuditOutcome.PARTIAL if permission_filtered_count else AuditOutcome.SUCCESS, context, scope_decision, details={ "query": normalized.to_dict(), "result_count": result.result_count, "total": result.total, "permission_filtered_count": permission_filtered_count, "permission_filter_duration_ms": permission_filter_duration_ms, }, ) return result def record_feedback( self, request: RetrievalFeedbackRequest, context: OperationContext, ) -> RetrievalFeedbackResult: diagnostics: list[Diagnostic] = [] label = _parse_feedback_label(request.label, diagnostics) if diagnostics: return RetrievalFeedbackResult( record=None, correlation_id=context.correlation_id, diagnostics=tuple(diagnostics), success=False, ) decision = self._authorize_for_retrieval( context, "retrieval.feedback.record", "retrieval:feedback", resource_metadata=request.to_dict(), ) if not decision.allowed: self._audit_retrieval( "retrieval.feedback.record", "retrieval:feedback", AuditOutcome.DENIED, context, decision, details=request.to_dict(), ) return RetrievalFeedbackResult( record=None, correlation_id=context.correlation_id, diagnostics=( _permission_denied_diagnostic(decision), ), success=False, ) record = RetrievalFeedbackRecord( label=label or RetrievalFeedbackLabel.LOW_CONFIDENCE, query=dict(request.query), result_ref=dict(request.result_ref), actor_id=context.actor.id, correlation_id=context.correlation_id, notes=request.notes, metadata=dict(request.metadata), ) saved = self.repository.save_retrieval_feedback(record) self._audit_retrieval( "retrieval.feedback.record", f"retrieval_feedback:{saved.feedback_id}", AuditOutcome.SUCCESS, context, decision, details={ "feedback_id": saved.feedback_id, "label": saved.label.value, "result_ref": dict(saved.result_ref), }, ) return RetrievalFeedbackResult(record=saved, correlation_id=context.correlation_id) def list_feedback( self, *, correlation_id: str | None = None, label: RetrievalFeedbackLabel | str | None = None, ) -> tuple[RetrievalFeedbackRecord, ...]: label_value = label.value if isinstance(label, RetrievalFeedbackLabel) else label return tuple(self.repository.list_retrieval_feedback(correlation_id=correlation_id, label=label_value)) def quality_metrics( self, *, query_results: tuple[AssetQueryResult, ...] = (), precision_at_k: int = 5, ) -> RetrievalQualityMetrics: feedback = self.repository.list_retrieval_feedback() query_count = len(query_results) zero_result_count = sum(1 for result in query_results if result.total == 0) ranked_feedback = [ record for record in feedback if _feedback_rank(record) is not None and (_feedback_rank(record) or 0) <= precision_at_k ] useful_ranked = [record for record in ranked_feedback if record.label == RetrievalFeedbackLabel.USEFUL] citation_feedback = [record for record in feedback if _feedback_has_citation_ref(record)] useful_citations = [record for record in citation_feedback if record.label == RetrievalFeedbackLabel.USEFUL] permission_latencies = [ float(event.details["permission_filter_duration_ms"]) for event in self.repository.list_audit_events() if event.operation.startswith("retrieval.") and "permission_filter_duration_ms" in event.details ] return RetrievalQualityMetrics( query_count=query_count, zero_result_count=zero_result_count, zero_result_rate=zero_result_count / query_count if query_count else 0.0, feedback_count=len(feedback), useful_count=sum(1 for record in feedback if record.label == RetrievalFeedbackLabel.USEFUL), unsafe_count=sum(1 for record in feedback if record.label == RetrievalFeedbackLabel.UNSAFE), low_confidence_count=sum(1 for record in feedback if record.label == RetrievalFeedbackLabel.LOW_CONFIDENCE), precision_at_k=len(useful_ranked) / len(ranked_feedback) if ranked_feedback else None, citation_precision=len(useful_citations) / len(citation_feedback) if citation_feedback else None, permission_filter_observation_count=len(permission_latencies), average_permission_filter_duration_ms=sum(permission_latencies) / len(permission_latencies) if permission_latencies else None, ) def _item_for_asset( self, asset: KnowledgeAsset, representation_kind: RepresentationKind | None, *, relevance: dict[str, Any] | None = None, include_relationships: bool = False, relationships: tuple[CoreRelationship, ...] | None = None, context_entities: tuple[ContextEntity, ...] | None = None, snippets: tuple[RetrievalSnippet, ...] = (), context: OperationContext | None = None, ) -> AssetQueryItem: if relationships is None and include_relationships: relationships = self._relationships_for_asset(asset.id, "both") if context is not None: relationships = tuple( relationship for relationship in relationships if self._relationship_allowed(relationship, context) ) if context_entities is None and relationships: context_entities = self._context_entities_for_relationships(relationships) return AssetQueryItem( asset=asset, representations=tuple( self.repository.list_representations( asset_id=asset.id, kind=representation_kind, ) ), metadata_records=tuple(self.repository.list_metadata_records(asset.id)), relationships=tuple(relationships or ()), context_entities=tuple(context_entities or ()), snippets=tuple(snippets), relevance=dict(relevance or {}), ) def _lexical_matches( self, text: str, *, include_snippets: bool, max_snippets: int, snippet_radius: int, ) -> tuple[dict[str, dict[str, Any]], dict[str, tuple[RetrievalSnippet, ...]]]: if not self._lexical_index: self.refresh_index() needle = text.casefold() matches: dict[str, dict[str, Any]] = {} snippets: dict[str, list[RetrievalSnippet]] = {} for document in self._lexical_index: haystack = document.text.casefold() count = haystack.count(needle) if count <= 0: continue current = matches.setdefault( document.asset_id, { "strategy": "lexical_substring", "query": text, "match_count": 0, "representation_ids": [], }, ) current["match_count"] += count current["representation_ids"].append(document.representation_id) if include_snippets: asset_snippets = snippets.setdefault(document.asset_id, []) remaining = max_snippets - len(asset_snippets) if remaining > 0: asset_snippets.extend( _snippets_for_document( document, text, max_snippets=remaining, snippet_radius=snippet_radius, ) ) current["snippet_count"] = len(asset_snippets) return matches, {asset_id: tuple(items) for asset_id, items in snippets.items()} def _context_entity_item(self, entity: ContextEntity, context: OperationContext) -> ContextEntityQueryItem: relationships = [ relationship for relationship in self.repository.list_relationships(target_id=entity.entity_id) if relationship.target_kind == RelationshipTargetKind.CONTEXT_ENTITY and self._relationship_allowed(relationship, context) ] return ContextEntityQueryItem( entity=entity, asset_ids=tuple(sorted({relationship.source_id for relationship in relationships})), relationship_count=len(relationships), ) def _relationship_item(self, relationship: CoreRelationship) -> RelationshipQueryItem: source_asset = self.repository.get_asset(relationship.source_id) target_asset = None target_entity = None if relationship.target_kind == RelationshipTargetKind.ASSET: target_asset = self.repository.get_asset(relationship.target_id) else: target_entity = self.repository.get_context_entity(relationship.target_id) return RelationshipQueryItem( relationship=relationship, source_asset=source_asset, target_asset=target_asset, target_entity=target_entity, ) def _relationship_context_for_asset( self, asset_id: str, request: AssetQueryRequest, context: OperationContext, ) -> tuple[tuple[CoreRelationship, ...], tuple[ContextEntity, ...]]: relationships = self._relationships_for_asset(asset_id, request.relationship_direction) entity_ids = self._context_entity_ids_for_asset_query(request) filtered: list[CoreRelationship] = [] for relationship in relationships: if request.relationship_predicate is not None and relationship.predicate != request.relationship_predicate: continue if request.related_asset_id is not None and not _relationship_connects_asset( relationship, asset_id=asset_id, related_asset_id=request.related_asset_id, ): continue if entity_ids is not None and not ( relationship.target_kind == RelationshipTargetKind.CONTEXT_ENTITY and relationship.target_id in entity_ids ): continue if not self._relationship_allowed(relationship, context): continue filtered.append(relationship) ordered = tuple(_sort_relationships(_unique_relationships(filtered), "source_id", "asc")) return ordered, self._context_entities_for_relationships(ordered) def _relationships_for_request(self, request: RelationshipQueryRequest) -> list[CoreRelationship]: relationships = self._relationships_for_relationship_query_base(request) entity_ids = self._context_entity_ids_for_relationship_query(request) filtered: list[CoreRelationship] = [] for relationship in relationships: if request.target_kind is not None and relationship.target_kind != request.target_kind: continue if request.predicate is not None and relationship.predicate != request.predicate: continue if entity_ids is not None and not ( relationship.target_kind == RelationshipTargetKind.CONTEXT_ENTITY and relationship.target_id in entity_ids ): continue filtered.append(relationship) return _unique_relationships(filtered) def _relationships_for_relationship_query_base( self, request: RelationshipQueryRequest, ) -> list[CoreRelationship]: if request.asset_id is not None: return list(self._relationships_for_asset(request.asset_id, request.direction)) if request.source_id is not None and request.target_id is not None: return [ relationship for relationship in self.repository.list_relationships(source_id=request.source_id) if relationship.target_id == request.target_id ] if request.source_id is not None: return list(self.repository.list_relationships(source_id=request.source_id)) if request.target_id is not None: return list(self.repository.list_relationships(target_id=request.target_id)) return _unique_relationships( [ relationship for asset in self.repository.list_assets() for relationship in self.repository.list_relationships(source_id=asset.id) ] ) def _relationships_for_asset(self, asset_id: str, direction: str) -> tuple[CoreRelationship, ...]: relationships: list[CoreRelationship] = [] if direction in {"outbound", "both"}: relationships.extend(self.repository.list_relationships(source_id=asset_id)) if direction in {"inbound", "both"}: relationships.extend( relationship for relationship in self.repository.list_relationships(target_id=asset_id) if relationship.target_kind == RelationshipTargetKind.ASSET ) return tuple(_unique_relationships(relationships)) def _context_entity_ids_for_asset_query(self, request: AssetQueryRequest) -> set[str] | None: if not _asset_query_has_context_entity_filter(request): return None return self._resolve_context_entity_ids( entity_id=request.context_entity_id, entity_type=request.context_entity_type, name=request.context_entity_name, external_ref=request.context_entity_external_ref, workflow_run_id=request.workflow_run_id, ) def _context_entity_ids_for_relationship_query(self, request: RelationshipQueryRequest) -> set[str] | None: if not _relationship_query_has_context_entity_filter(request): return None return self._resolve_context_entity_ids( entity_id=request.context_entity_id, entity_type=request.context_entity_type, name=request.context_entity_name, external_ref=request.context_entity_external_ref, workflow_run_id=request.workflow_run_id, ) def _resolve_context_entity_ids( self, *, entity_id: str | None, entity_type: ContextEntityType | None, name: str | None, external_ref: str | None, workflow_run_id: str | None, ) -> set[str]: entities = self.repository.list_context_entities() if workflow_run_id is not None: entity_type = ContextEntityType.WORKFLOW_RUN return { entity.entity_id for entity in entities if _context_entity_matches( entity, entity_id=entity_id, entity_type=entity_type, name=name, external_ref=external_ref, metadata_filters={}, ) and _workflow_run_matches(entity, workflow_run_id) } def _context_entities_for_relationships( self, relationships: tuple[CoreRelationship, ...], ) -> tuple[ContextEntity, ...]: entities: list[ContextEntity] = [] for relationship in relationships: if relationship.target_kind != RelationshipTargetKind.CONTEXT_ENTITY: continue entities.append(self.repository.get_context_entity(relationship.target_id)) return tuple(_unique_context_entities(entities)) def _asset_allowed(self, asset: KnowledgeAsset, context: OperationContext) -> bool: decision = self._authorize_for_retrieval( context, "asset.retrieve", f"asset:{asset.id}", resource_metadata={ "asset_id": asset.id, "asset_type": asset.classification.asset_type, "lifecycle": asset.lifecycle.value, "sensitivity": asset.classification.sensitivity.value, "owner": asset.classification.owner, }, ) return decision.allowed def _context_entity_allowed(self, entity: ContextEntity, context: OperationContext) -> bool: decision = self._authorize_for_retrieval( context, "context_entity.retrieve", f"context_entity:{entity.entity_id}", resource_metadata={ "entity_id": entity.entity_id, "entity_type": entity.entity_type.value, "external_ref": entity.external_ref, }, ) return decision.allowed def _relationship_allowed(self, relationship: CoreRelationship, context: OperationContext) -> bool: decision = self._authorize_for_retrieval( context, "asset.relationship.retrieve", f"relationship:{relationship.relationship_id}", resource_metadata={ "relationship_id": relationship.relationship_id, "source_id": relationship.source_id, "target_id": relationship.target_id, "target_kind": relationship.target_kind.value, "predicate": relationship.predicate, }, ) if not decision.allowed: return False source_asset = self.repository.get_asset(relationship.source_id) if not self._asset_allowed(source_asset, context): return False if relationship.target_kind == RelationshipTargetKind.ASSET: target_asset = self.repository.get_asset(relationship.target_id) return self._asset_allowed(target_asset, context) target_entity = self.repository.get_context_entity(relationship.target_id) return self._context_entity_allowed(target_entity, context) def _authorize_for_retrieval( self, context: OperationContext, action: str, resource: str, *, resource_metadata: dict[str, Any] | None = None, ) -> PolicyDecision: self.repository.save_actor(context.actor) try: return self.policy_gateway.authorize( context, action, resource, resource_metadata=resource_metadata, ) except Exception as exc: return PolicyDecision.fail_closed( context.actor.id, action, resource, reason=str(exc) or "Retrieval policy gateway failed", context={ "gateway_error": type(exc).__name__, "resource_metadata": resource_metadata or {}, }, ) def _audit_retrieval( self, operation: str, target: str, outcome: AuditOutcome, context: OperationContext, policy_decision: PolicyDecision, *, details: dict[str, Any] | None = None, ) -> AuditEvent: event = AuditEvent.from_context( operation, target, outcome, context, policy_decision=policy_decision, details=details, ) return self.repository.save_audit_event(event) def _validate_request(request: AssetQueryRequest) -> tuple[list[Diagnostic], AssetQueryRequest]: diagnostics: list[Diagnostic] = [] lifecycle = _parse_lifecycle(request.lifecycle, diagnostics) sensitivity = _parse_sensitivity(request.sensitivity, diagnostics) representation_kind = _parse_representation_kind(request.representation_kind, diagnostics) context_entity_type = _parse_context_entity_type(request.context_entity_type, diagnostics) relationship_direction = _parse_relationship_direction(request.relationship_direction, diagnostics) if request.limit < 1 or request.limit > 500: diagnostics.append( Diagnostic( severity="error", code="retrieval.limit_invalid", message="Query limit must be between 1 and 500", details={"limit": request.limit}, ) ) if request.offset < 0: diagnostics.append( Diagnostic( severity="error", code="retrieval.offset_invalid", message="Query offset must be zero or greater", details={"offset": request.offset}, ) ) if request.max_snippets < 0 or request.max_snippets > 20: diagnostics.append( Diagnostic( severity="error", code="retrieval.max_snippets_invalid", message="Max snippets must be between 0 and 20", details={"max_snippets": request.max_snippets}, ) ) if request.snippet_radius < 0 or request.snippet_radius > 1000: diagnostics.append( Diagnostic( severity="error", code="retrieval.snippet_radius_invalid", message="Snippet radius must be between 0 and 1000", details={"snippet_radius": request.snippet_radius}, ) ) if request.sort_by not in SUPPORTED_SORT_KEYS: diagnostics.append( Diagnostic( severity="error", code="retrieval.sort_invalid", message="Query sort key is not supported", details={"sort_by": request.sort_by, "supported": sorted(SUPPORTED_SORT_KEYS)}, ) ) if request.sort_order not in {"asc", "desc"}: diagnostics.append( Diagnostic( severity="error", code="retrieval.sort_order_invalid", message="Query sort order must be asc or desc", details={"sort_order": request.sort_order}, ) ) return diagnostics, AssetQueryRequest( text=request.text, asset_type=request.asset_type, lifecycle=lifecycle, sensitivity=sensitivity, owner=request.owner, topic=request.topic, tags=tuple(request.tags), collection=request.collection, review_state=request.review_state, metadata_filters=dict(request.metadata_filters), confirmed_metadata_only=request.confirmed_metadata_only, source_system=request.source_system, source_path=request.source_path, context_entity_id=request.context_entity_id, context_entity_type=context_entity_type, context_entity_name=request.context_entity_name, context_entity_external_ref=request.context_entity_external_ref, workflow_run_id=request.workflow_run_id, related_asset_id=request.related_asset_id, relationship_predicate=request.relationship_predicate, relationship_direction=relationship_direction, include_relationships=request.include_relationships, include_snippets=request.include_snippets, max_snippets=request.max_snippets, snippet_radius=request.snippet_radius, created_after=request.created_after, created_before=request.created_before, updated_after=request.updated_after, updated_before=request.updated_before, representation_kind=representation_kind, sort_by=request.sort_by, sort_order=request.sort_order, limit=request.limit, offset=request.offset, ) def _permission_denied_diagnostic(decision: PolicyDecision) -> Diagnostic: return Diagnostic( severity="error", code="retrieval.permission_denied", message="Retrieval query denied by policy", details={"policy_decision": decision.to_dict()}, ) def _parse_feedback_label( value: RetrievalFeedbackLabel | str, diagnostics: list[Diagnostic], ) -> RetrievalFeedbackLabel | None: if isinstance(value, RetrievalFeedbackLabel): return value try: return RetrievalFeedbackLabel(value) except ValueError: diagnostics.append( Diagnostic( severity="error", code="retrieval.feedback_label_invalid", message="Retrieval feedback label is not supported", details={"label": value, "supported": [item.value for item in RetrievalFeedbackLabel]}, ) ) return None def _feedback_rank(record: RetrievalFeedbackRecord) -> int | None: rank = record.result_ref.get("rank") if isinstance(rank, int): return rank if isinstance(rank, str) and rank.isdigit(): return int(rank) return None def _feedback_has_citation_ref(record: RetrievalFeedbackRecord) -> bool: return any( record.result_ref.get(key) for key in ( "snippet_id", "representation_id", "source_ref_id", ) ) or bool(record.metadata.get("citation")) def _elapsed_ms(started_at: float) -> float: return round((perf_counter() - started_at) * 1000, 3) def _parse_lifecycle(value: LifecycleState | str | None, diagnostics: list[Diagnostic]) -> LifecycleState | None: if value is None or isinstance(value, LifecycleState): return value try: return LifecycleState(value) except ValueError: diagnostics.append( Diagnostic( severity="error", code="retrieval.lifecycle_invalid", message="Lifecycle filter is not supported", details={"lifecycle": value, "supported": [item.value for item in LifecycleState]}, ) ) return None def _parse_sensitivity(value: Sensitivity | str | None, diagnostics: list[Diagnostic]) -> Sensitivity | None: if value is None or isinstance(value, Sensitivity): return value try: return Sensitivity(value) except ValueError: diagnostics.append( Diagnostic( severity="error", code="retrieval.sensitivity_invalid", message="Sensitivity filter is not supported", details={"sensitivity": value, "supported": [item.value for item in Sensitivity]}, ) ) return None def _parse_representation_kind( value: RepresentationKind | str | None, diagnostics: list[Diagnostic], ) -> RepresentationKind | None: if value is None or isinstance(value, RepresentationKind): return value try: return RepresentationKind(value) except ValueError: diagnostics.append( Diagnostic( severity="error", code="retrieval.representation_kind_invalid", message="Representation kind filter is not supported", details={"representation_kind": value, "supported": [item.value for item in RepresentationKind]}, ) ) return None def _parse_context_entity_type( value: ContextEntityType | str | None, diagnostics: list[Diagnostic], ) -> ContextEntityType | None: if value is None or isinstance(value, ContextEntityType): return value try: return ContextEntityType(value) except ValueError: diagnostics.append( Diagnostic( severity="error", code="retrieval.context_entity_type_invalid", message="Context entity type filter is not supported", details={"context_entity_type": value, "supported": [item.value for item in ContextEntityType]}, ) ) return None def _parse_relationship_target_kind( value: RelationshipTargetKind | str | None, diagnostics: list[Diagnostic], ) -> RelationshipTargetKind | None: if value is None or isinstance(value, RelationshipTargetKind): return value try: return RelationshipTargetKind(value) except ValueError: diagnostics.append( Diagnostic( severity="error", code="retrieval.relationship_target_kind_invalid", message="Relationship target kind is not supported", details={"target_kind": value, "supported": [item.value for item in RelationshipTargetKind]}, ) ) return None def _parse_relationship_direction(value: str, diagnostics: list[Diagnostic]) -> str: if value in {"outbound", "inbound", "both"}: return value diagnostics.append( Diagnostic( severity="error", code="retrieval.relationship_direction_invalid", message="Relationship direction must be outbound, inbound, or both", details={"direction": value}, ) ) return "both" def _validate_context_entity_request( request: ContextEntityQueryRequest, ) -> tuple[list[Diagnostic], ContextEntityQueryRequest]: diagnostics: list[Diagnostic] = [] entity_type = _parse_context_entity_type(request.entity_type, diagnostics) _validate_window_and_sort( diagnostics, limit=request.limit, offset=request.offset, sort_by=request.sort_by, sort_order=request.sort_order, supported_sort_keys=SUPPORTED_CONTEXT_ENTITY_SORT_KEYS, ) return diagnostics, ContextEntityQueryRequest( entity_id=request.entity_id, entity_type=entity_type, name=request.name, external_ref=request.external_ref, metadata_filters=dict(request.metadata_filters), sort_by=request.sort_by, sort_order=request.sort_order, limit=request.limit, offset=request.offset, ) def _validate_relationship_request( request: RelationshipQueryRequest, ) -> tuple[list[Diagnostic], RelationshipQueryRequest]: diagnostics: list[Diagnostic] = [] context_entity_type = _parse_context_entity_type(request.context_entity_type, diagnostics) target_kind = _parse_relationship_target_kind(request.target_kind, diagnostics) direction = _parse_relationship_direction(request.direction, diagnostics) _validate_window_and_sort( diagnostics, limit=request.limit, offset=request.offset, sort_by=request.sort_by, sort_order=request.sort_order, supported_sort_keys=SUPPORTED_RELATIONSHIP_SORT_KEYS, ) return diagnostics, RelationshipQueryRequest( source_id=request.source_id, target_id=request.target_id, asset_id=request.asset_id, context_entity_id=request.context_entity_id, context_entity_type=context_entity_type, context_entity_name=request.context_entity_name, context_entity_external_ref=request.context_entity_external_ref, workflow_run_id=request.workflow_run_id, target_kind=target_kind, predicate=request.predicate, direction=direction, sort_by=request.sort_by, sort_order=request.sort_order, limit=request.limit, offset=request.offset, ) def _validate_window_and_sort( diagnostics: list[Diagnostic], *, limit: int, offset: int, sort_by: str, sort_order: str, supported_sort_keys: set[str], ) -> None: if limit < 1 or limit > 500: diagnostics.append( Diagnostic( severity="error", code="retrieval.limit_invalid", message="Query limit must be between 1 and 500", details={"limit": limit}, ) ) if offset < 0: diagnostics.append( Diagnostic( severity="error", code="retrieval.offset_invalid", message="Query offset must be zero or greater", details={"offset": offset}, ) ) if sort_by not in supported_sort_keys: diagnostics.append( Diagnostic( severity="error", code="retrieval.sort_invalid", message="Query sort key is not supported", details={"sort_by": sort_by, "supported": sorted(supported_sort_keys)}, ) ) if sort_order not in {"asc", "desc"}: diagnostics.append( Diagnostic( severity="error", code="retrieval.sort_order_invalid", message="Query sort order must be asc or desc", details={"sort_order": sort_order}, ) ) def _source_matches( asset: KnowledgeAsset, *, source_system: str | None, source_path: str | None, ) -> bool: if source_system is None and source_path is None: return True for source_ref in asset.source_refs: if source_system is not None and source_ref.source_system != source_system: continue if source_path is not None and source_ref.path != source_path: continue return True return False def _collection_matches( asset: KnowledgeAsset, metadata_records: list[MetadataRecord], collection: str | None, ) -> bool: if collection is None: return True if asset.metadata.get("collection") == collection: return True if asset.classification.metadata.get("collection") == collection: return True return any(record.key == "collection" and record.value == collection for record in metadata_records) def _tags_match( asset: KnowledgeAsset, metadata_records: list[MetadataRecord], tags: tuple[str, ...], ) -> bool: if not tags: return True values = set(asset.classification.topics) for record in metadata_records: if record.key not in {"tag", "tags"}: continue if isinstance(record.value, list): values.update(str(item) for item in record.value) elif record.value is not None: values.add(str(record.value)) return all(tag in values for tag in tags) def _timestamp_matches(asset: KnowledgeAsset, request: AssetQueryRequest) -> bool: if request.created_after and asset.created_at < request.created_after: return False if request.created_before and asset.created_at > request.created_before: return False if request.updated_after and asset.updated_at < request.updated_after: return False if request.updated_before and asset.updated_at > request.updated_before: return False return True def _asset_query_has_graph_filter(request: AssetQueryRequest) -> bool: return ( _asset_query_has_context_entity_filter(request) or request.related_asset_id is not None or request.relationship_predicate is not None ) def _asset_query_has_context_entity_filter(request: AssetQueryRequest) -> bool: return any( value is not None for value in ( request.context_entity_id, request.context_entity_type, request.context_entity_name, request.context_entity_external_ref, request.workflow_run_id, ) ) def _relationship_query_has_context_entity_filter(request: RelationshipQueryRequest) -> bool: return any( value is not None for value in ( request.context_entity_id, request.context_entity_type, request.context_entity_name, request.context_entity_external_ref, request.workflow_run_id, ) ) def _context_entity_matches( entity: ContextEntity, *, entity_id: str | None, entity_type: ContextEntityType | None, name: str | None, external_ref: str | None, metadata_filters: dict[str, Any], ) -> bool: if entity_id is not None and entity.entity_id != entity_id: return False if entity_type is not None and entity.entity_type != entity_type: return False if name is not None and entity.name != name: return False if external_ref is not None and entity.external_ref != external_ref: return False for key, expected in metadata_filters.items(): if not _metadata_value_matches(entity.metadata.get(key), expected): return False return True def _workflow_run_matches(entity: ContextEntity, workflow_run_id: str | None) -> bool: if workflow_run_id is None: return True if entity.entity_id == workflow_run_id or entity.external_ref == workflow_run_id: return True return entity.metadata.get("workflow_run_id") == workflow_run_id def _snippet_provenance(representation: AssetRepresentation) -> dict[str, Any]: provenance: dict[str, Any] = {} for key in ( "adapter_provenance", "context_span", "extractor", "markitect_selector", "snapshot", "source_span", ): value = representation.metadata.get(key) if value: provenance[key] = value return provenance def _snippets_for_document( document: _LexicalDocument, query_text: str, *, max_snippets: int, snippet_radius: int, ) -> list[RetrievalSnippet]: if max_snippets <= 0: return [] needle = query_text.casefold() haystack = document.text.casefold() snippets: list[RetrievalSnippet] = [] cursor = 0 while len(snippets) < max_snippets: match_start = haystack.find(needle, cursor) if match_start < 0: break match_end = match_start + len(query_text) snippet_start = max(0, match_start - snippet_radius) snippet_end = min(len(document.text), match_end + snippet_radius) snippet_text = document.text[snippet_start:snippet_end].strip() if snippet_start > 0: snippet_text = "..." + snippet_text if snippet_end < len(document.text): snippet_text = snippet_text + "..." snippets.append( RetrievalSnippet( asset_id=document.asset_id, representation_id=document.representation_id, source_ref_id=document.source_ref_id, storage_ref=document.storage_ref, media_type=document.media_type, text=snippet_text, start_offset=snippet_start, end_offset=snippet_end, match_text=document.text[match_start:match_end], provenance=dict(document.provenance), ) ) cursor = match_end return snippets def _relationship_connects_asset( relationship: CoreRelationship, *, asset_id: str, related_asset_id: str, ) -> bool: if relationship.target_kind != RelationshipTargetKind.ASSET: return False return ( relationship.source_id == asset_id and relationship.target_id == related_asset_id or relationship.source_id == related_asset_id and relationship.target_id == asset_id ) def _metadata_value_matches(value: Any, expected: Any) -> bool: if isinstance(value, list): return expected in value return value == expected def _unique_relationships(relationships: list[CoreRelationship]) -> list[CoreRelationship]: seen: set[str] = set() unique: list[CoreRelationship] = [] for relationship in relationships: if relationship.relationship_id in seen: continue seen.add(relationship.relationship_id) unique.append(relationship) return unique def _unique_context_entities(entities: list[ContextEntity]) -> list[ContextEntity]: seen: set[str] = set() unique: list[ContextEntity] = [] for entity in entities: if entity.entity_id in seen: continue seen.add(entity.entity_id) unique.append(entity) return sorted(unique, key=lambda item: (item.entity_type.value, item.name, item.entity_id)) def _sort_assets( assets: list[KnowledgeAsset], sort_by: str, sort_order: str, ) -> list[KnowledgeAsset]: reverse = sort_order == "desc" def sort_key(asset: KnowledgeAsset) -> tuple[Any, str]: if sort_by == "asset_id": primary = asset.id elif sort_by == "asset_type": primary = asset.classification.asset_type elif sort_by == "created_at": primary = asset.created_at elif sort_by == "lifecycle": primary = asset.lifecycle.value elif sort_by == "updated_at": primary = asset.updated_at else: primary = asset.title.casefold() return primary, asset.id return sorted(assets, key=sort_key, reverse=reverse) def _sort_context_entities( entities: list[ContextEntity], sort_by: str, sort_order: str, ) -> list[ContextEntity]: reverse = sort_order == "desc" def sort_key(entity: ContextEntity) -> tuple[Any, str]: if sort_by == "entity_id": primary = entity.entity_id elif sort_by == "entity_type": primary = entity.entity_type.value elif sort_by == "external_ref": primary = entity.external_ref or "" else: primary = entity.name.casefold() return primary, entity.entity_id return sorted(entities, key=sort_key, reverse=reverse) def _sort_relationships( relationships: list[CoreRelationship], sort_by: str, sort_order: str, ) -> list[CoreRelationship]: reverse = sort_order == "desc" def sort_key(relationship: CoreRelationship) -> tuple[Any, str]: if sort_by == "created_at": primary = relationship.created_at elif sort_by == "predicate": primary = relationship.predicate elif sort_by == "relationship_id": primary = relationship.relationship_id elif sort_by == "target_id": primary = relationship.target_id elif sort_by == "target_kind": primary = relationship.target_kind.value else: primary = relationship.source_id return primary, relationship.relationship_id return sorted(relationships, key=sort_key, reverse=reverse)