123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
"""Programmatic query and retrieval API."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from .artifacts import Artifact, Relationship
|
|
from .errors import Diagnostic, KontextualError
|
|
from .storage import KnowledgeRepository
|
|
|
|
|
|
@dataclass
|
|
class QueryResult:
|
|
query: str
|
|
results: list[dict[str, Any]]
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
diagnostics: list[Diagnostic] = field(default_factory=list)
|
|
success: bool = True
|
|
|
|
@property
|
|
def result_count(self) -> int:
|
|
return len(self.results)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"query": self.query,
|
|
"result_count": self.result_count,
|
|
"results": self.results,
|
|
"metadata": dict(self.metadata),
|
|
"diagnostics": [d.to_dict() for d in self.diagnostics],
|
|
"success": self.success,
|
|
}
|
|
|
|
|
|
class QueryEngine:
|
|
"""Small deterministic query facade over repository state."""
|
|
|
|
def __init__(self, repository: KnowledgeRepository) -> None:
|
|
self.repository = repository
|
|
|
|
def artifacts(
|
|
self,
|
|
*,
|
|
collection_id: str | None = None,
|
|
name: str | None = None,
|
|
digest: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
text_contains: str | None = None,
|
|
) -> QueryResult:
|
|
query_parts = ["artifacts"]
|
|
matches = self.repository.list_artifacts(collection_id=collection_id)
|
|
if collection_id:
|
|
query_parts.append(f"collection_id={collection_id}")
|
|
if name is not None:
|
|
matches = [a for a in matches if a.name == name]
|
|
query_parts.append(f"name={name}")
|
|
if digest is not None:
|
|
matches = [a for a in matches if a.content_digest == digest]
|
|
query_parts.append(f"digest={digest}")
|
|
if metadata:
|
|
for key, value in metadata.items():
|
|
matches = [a for a in matches if _metadata_value(a, key) == value]
|
|
query_parts.append(f"metadata.{key}={value}")
|
|
if text_contains:
|
|
needle = text_contains.lower()
|
|
matches = [a for a in matches if needle in a.content.lower()]
|
|
query_parts.append(f"text~={text_contains}")
|
|
return QueryResult(
|
|
query=" ".join(query_parts),
|
|
results=[a.to_dict() for a in matches],
|
|
metadata={"type": "artifact"},
|
|
)
|
|
|
|
def relationships(self, *, artifact_id: str | None = None) -> QueryResult:
|
|
relationships = self.repository.list_relationships(artifact_id=artifact_id)
|
|
query = "relationships" + (f" artifact_id={artifact_id}" if artifact_id else "")
|
|
return QueryResult(
|
|
query=query,
|
|
results=[r.to_dict() for r in relationships],
|
|
metadata={"type": "relationship"},
|
|
)
|
|
|
|
def related_artifacts(self, artifact_id: str) -> QueryResult:
|
|
try:
|
|
relationships = self.repository.list_relationships(artifact_id=artifact_id)
|
|
related: list[Artifact] = []
|
|
seen: set[str] = set()
|
|
for relationship in relationships:
|
|
other_id = _other_endpoint(relationship, artifact_id)
|
|
if other_id and other_id not in seen:
|
|
related.append(self.repository.get_artifact(other_id))
|
|
seen.add(other_id)
|
|
return QueryResult(
|
|
query=f"related artifact_id={artifact_id}",
|
|
results=[a.to_dict() for a in related],
|
|
metadata={"type": "related_artifacts", "relationship_count": len(relationships)},
|
|
)
|
|
except KontextualError as exc:
|
|
return QueryResult(
|
|
query=f"related artifact_id={artifact_id}",
|
|
results=[],
|
|
diagnostics=[exc.diagnostic()],
|
|
success=False,
|
|
)
|
|
|
|
|
|
def _metadata_value(artifact: Artifact, path: str) -> Any:
|
|
data: Any = artifact.metadata.to_dict()
|
|
for part in path.split("."):
|
|
if not isinstance(data, dict):
|
|
return None
|
|
data = data.get(part)
|
|
return data
|
|
|
|
|
|
def _other_endpoint(relationship: Relationship, artifact_id: str) -> str | None:
|
|
if relationship.source_artifact_id == artifact_id:
|
|
return relationship.target_artifact_id
|
|
if relationship.target_artifact_id == artifact_id:
|
|
return relationship.source_artifact_id
|
|
return None
|
|
|