"""ContextEngine: assembles "il fascicolo completo" for an AI operation. Input: an artifact (or project id, or a free-text question) Output: a ``ContextAssembly`` bundling · related requirements · architecture decisions · relevant source files · docs · tests · deployments · previous AI analysis · incidents · runtime info (currently placeholder) For MVP it walks the ArtifactRelationship graph in MongoDB. Semantic retrieval is hidden behind ``SemanticStore`` so we can add embeddings later without changing the engine surface. """ from __future__ import annotations from dataclasses import dataclass, field from typing import Optional, Sequence from uuid import UUID import structlog from asf_domain import Artifact, ArtifactRelationship, RelationshipType from .semantic_store import SemanticStore, InMemorySemanticStoreMock, SemanticHit log = structlog.get_logger(__name__) @dataclass class AssemblyOptions: max_depth: int = 2 max_nodes: int = 200 include_semantic: bool = False semantic_top_k: int = 10 relationship_whitelist: Optional[Sequence[RelationshipType]] = None @dataclass class ContextAssembly: seed_artifact_id: Optional[UUID] project_id: Optional[UUID] tenant_id: UUID # Artifacts grouped by logical concern (all are Artifact objects) requirements: list[Artifact] = field(default_factory=list) architecture_decisions: list[Artifact] = field(default_factory=list) tasks: list[Artifact] = field(default_factory=list) source_files: list[Artifact] = field(default_factory=list) documents: list[Artifact] = field(default_factory=list) tests: list[Artifact] = field(default_factory=list) builds: list[Artifact] = field(default_factory=list) deployments: list[Artifact] = field(default_factory=list) releases: list[Artifact] = field(default_factory=list) incidents: list[Artifact] = field(default_factory=list) previous_ai_analysis: list[Artifact] = field(default_factory=list) security_findings: list[Artifact] = field(default_factory=list) # Raw graph for UI / debugging relationships_traversed: list[ArtifactRelationship] = field(default_factory=list) all_artifacts: list[Artifact] = field(default_factory=list) semantic_hits: list[SemanticHit] = field(default_factory=list) warnings: list[str] = field(default_factory=list) def summary(self) -> dict: return { "seed_artifact_id": str(self.seed_artifact_id) if self.seed_artifact_id else None, "total_artifacts": len(self.all_artifacts), "total_relationships": len(self.relationships_traversed), "semantic_hits": len(self.semantic_hits), "by_type": { "requirements": len(self.requirements), "architecture_decisions": len(self.architecture_decisions), "tasks": len(self.tasks), "source_files": len(self.source_files), "documents": len(self.documents), "tests": len(self.tests), "builds": len(self.builds), "deployments": len(self.deployments), "releases": len(self.releases), "incidents": len(self.incidents), "previous_ai_analysis": len(self.previous_ai_analysis), "security_findings": len(self.security_findings), }, "warnings": self.warnings, } class ContextEngine: """Assembles contextual "fascicolo" given a seed (artifact, project, ...). For the bootstrap phase we only accept explicit service-layer objects (relationship service + artifact service) so we can run the engine inside asf-core or standalone. No Mongo driver coupling here. """ def __init__( self, *, relationship_service, # duck-typed: for_artifact, outgoing, incoming artifact_service, # duck-typed: get, get_many, list semantic: SemanticStore | None = None, ) -> None: self._rels = relationship_service self._arts = artifact_service self._semantic = semantic or InMemorySemanticStoreMock() async def assemble( self, *, tenant_id: UUID, seed_artifact_id: Optional[UUID] = None, project_id: Optional[UUID] = None, user_query: Optional[str] = None, options: Optional[AssemblyOptions] = None, ) -> ContextAssembly: opts = options or AssemblyOptions() result = ContextAssembly( seed_artifact_id=seed_artifact_id, project_id=project_id, tenant_id=tenant_id, ) visited_ids: set[UUID] = set() frontier: set[UUID] = set() if seed_artifact_id is not None: frontier.add(seed_artifact_id) elif project_id is not None: arts = await self._arts.list( project_id=project_id, limit=opts.max_nodes ) frontier = {a.id for a in arts[: opts.max_nodes]} result.all_artifacts.extend(arts) visited_ids.update(a.id for a in arts) else: result.warnings.append( "No seed_artifact_id nor project_id provided; context will be empty." ) return result # BFS walk for _ in range(opts.max_depth): current = list(frontier - visited_ids) if not current: break visited_ids.update(current) # Load artifacts we haven't loaded yet need = [i for i in current if not any(a.id == i for a in result.all_artifacts)] if need: fetched = await self._arts.get_many(need) result.all_artifacts.extend(fetched) for aid in current: if len(result.relationships_traversed) >= opts.max_nodes: break out = await self._rels.outgoing( aid, limit=max(50, opts.max_nodes // 4) ) inc = await self._rels.incoming( aid, limit=max(50, opts.max_nodes // 4) ) for rel in out + inc: result.relationships_traversed.append(rel) frontier.add(rel.source_artifact_id) frontier.add(rel.target_artifact_id) # Ensure all referenced artifacts are loaded referenced: set[UUID] = set() for rel in result.relationships_traversed: referenced.add(rel.source_artifact_id) referenced.add(rel.target_artifact_id) already_have = {a.id for a in result.all_artifacts} missing = list(referenced - already_have) if missing: fetched = await self._arts.get_many(missing) result.all_artifacts.extend(fetched) # Bucket artifacts by type _bucket(result) # Optional semantic layer if opts.include_semantic and user_query: try: result.semantic_hits = list( await self._semantic.search( tenant_id=tenant_id, query=user_query, project_id=project_id, top_k=opts.semantic_top_k, ) ) except Exception: # noqa: BLE001 await log.aexception("semantic_search_failed") result.warnings.append("Semantic search failed; using structural context only.") return result # --------------------------------------------------------------------------- # Internal helpers # --------------------------------------------------------------------------- _TYPE_BUCKETS = { "requirement": "requirements", "architecture_decision": "architecture_decisions", "task": "tasks", "story": "tasks", "epic": "tasks", "source_file": "source_files", "document": "documents", "test": "tests", "build": "builds", "deployment": "deployments", "release": "releases", "incident": "incidents", "runtime_event": "incidents", "ai_analysis": "previous_ai_analysis", "security_finding": "security_findings", } def _bucket(r: ContextAssembly) -> None: for art in r.all_artifacts: t = art.type if isinstance(art.type, str) else art.type.value bucket = _TYPE_BUCKETS.get(t) if bucket is not None: getattr(r, bucket).append(art)