237 lines
8.4 KiB
Python
237 lines
8.4 KiB
Python
"""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)
|