AI Memory & Knowledge Layer

# digital_twin/knowledge/templates.py

TEMPLATES = {
    "design_decision": """
Design Decision: {title}
Status: {status} | Date: {date} | Domain: {domain}
Context: {context}
Options Considered: {options}
Outcome: {outcome}
Rationale: {rationale}
""",
    "component_rationale": """
Component Selection: {mpn} ({manufacturer})
For: {design_element_name} | Domain: {domain}
Package: {package} | Unit Price: ${unit_price}
Selection Rationale: {rationale}
Key Specifications: {specifications}
Alternatives Considered: {alternatives}
""",
    "failure_mode": """
Failure: {test_name} — {status}
Procedure: {procedure_name} | Category: {category}
Root Cause: {root_cause}
Corrective Action: {action}
Lesson Learned: {lesson}
Affected Components: {components}
""",
    "constraint_violation": """
Constraint Violation: {constraint_name}
Domain: {domain} | Severity: {severity}
Expected: {expected} | Actual: {actual}
Affected Nodes: {affected_nodes}
Resolution: {resolution}
""",
    "session_summary": """
Session: {skill} by {agent_code}
Duration: {duration} | Status: {status}
Input Summary: {input_summary}
Output Summary: {output_summary}
Key Decisions: {decisions}
Nodes Modified: {modified_nodes}
""",
    "datasheet_chunk": """
Datasheet: {part_number} — {manufacturer}
Section: {section_title}
Content: {chunk_text}
""",
    "application_note": """
Application Note: {title}
Source: {filename} | Section: {section_title}
Content: {chunk_text}
Tags: {tags}
""",
}

# digital_twin/knowledge/consumer.py

class KnowledgeConsumer:
    """Kafka consumer that processes events into knowledge embeddings."""

    SUBSCRIPTIONS = [
        "graph.mutations",      # Node created/updated events
        "session.completed",    # Agent session completion
        "knowledge.ingest",     # Manual document ingestion requests
    ]

    async def process_event(self, event: GraphEvent) -> None:
        """Route event to appropriate handler based on type."""
        if event.event_type in (EventType.NODE_CREATED, EventType.NODE_UPDATED):
            await self._handle_node_event(event)
        elif event.event_type == "session.completed":
            await self._handle_session_completed(event)
        elif event.event_type == "knowledge.ingest":
            await self._handle_document_ingestion(event)

    async def _handle_node_event(self, event: GraphEvent) -> None:
        """Index graph node changes as knowledge."""
        handler = self.NODE_HANDLERS.get(event.node_type)
        if handler:
            content = handler.render_template(event)
            embedding = await self.embedding_service.embed(content)
            await self.store.upsert(
                knowledge_type=handler.knowledge_type,
                source_event_id=event.event_id,
                source_node_id=event.node_id,
                domain=event.properties.get("domain"),
                content_text=content,
                embedding=embedding,
                metadata=handler.extract_metadata(event),
            )

    NODE_HANDLERS = {
        "DecisionRecord": DecisionHandler(),
        "BOMItem": ComponentRationaleHandler(),
        "TestExecution": FailureModeHandler(),
        "CAPARecord": FailureModeHandler(),
        "ConstraintViolation": ConstraintViolationHandler(),
    }

# digital_twin/knowledge/embedding_service.py

class EmbeddingService:
    """Unified embedding interface — local or cloud."""

    def __init__(self, provider: Literal["local", "openai"] = "local"):
        self.provider = provider
        if provider == "local":
            from sentence_transformers import SentenceTransformer
            self.model = SentenceTransformer("all-MiniLM-L6-v2")
            self.dimensions = 384
        else:
            import openai
            self.client = openai.OpenAI()
            self.dimensions = 1536

    async def embed(self, text: str) -> list[float]:
        """Generate embedding vector for text."""
        if self.provider == "local":
            return self.model.encode(text).tolist()
        else:
            response = await self.client.embeddings.create(
                model="text-embedding-3-small",
                input=text,
            )
            return response.data[0].embedding

    async def embed_batch(self, texts: list[str]) -> list[list[float]]:
        """Batch embedding for efficiency."""
        if self.provider == "local":
            return self.model.encode(texts).tolist()
        else:
            response = await self.client.embeddings.create(
                model="text-embedding-3-small",
                input=texts,
            )
            return [d.embedding for d in response.data]

# digital_twin/knowledge/chunker.py

class DocumentChunker:
    """Split documents into overlapping chunks for embedding."""

    def __init__(
        self,
        chunk_size: int = 512,      # tokens per chunk
        chunk_overlap: int = 64,     # overlapping tokens between chunks
    ):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def chunk_text(self, text: str, metadata: dict) -> list[ChunkResult]:
        """Split text into overlapping chunks with metadata."""
        tokens = self.tokenize(text)
        chunks = []
        start = 0
        while start < len(tokens):
            end = min(start + self.chunk_size, len(tokens))
            chunk_text = self.detokenize(tokens[start:end])
            chunks.append(ChunkResult(
                text=chunk_text,
                chunk_index=len(chunks),
                start_token=start,
                end_token=end,
                metadata={**metadata, "chunk_index": len(chunks)},
            ))
            start += self.chunk_size - self.chunk_overlap
        return chunks

    def chunk_pdf(self, pdf_path: str, metadata: dict) -> list[ChunkResult]:
        """Extract text from PDF and chunk it."""
        text = self.extract_pdf_text(pdf_path)
        return self.chunk_text(text, metadata)

# domain_agents/shared/skills/retrieve_knowledge/handler.py

class RetrieveKnowledgeInput(BaseModel):
    query: str                                       # Natural language query
    knowledge_types: Optional[list[str]] = None      # Filter by type
    domain: Optional[str] = None                     # Filter by domain
    project_id: Optional[str] = None                 # Filter by project
    top_k: int = 5                                   # Number of results
    similarity_threshold: float = 0.7                # Minimum similarity

class RetrieveKnowledgeOutput(BaseModel):
    results: list[KnowledgeResult]
    query_embedding_time_ms: float
    search_time_ms: float

class KnowledgeResult(BaseModel):
    id: str
    knowledge_type: str
    content_text: str
    similarity: float
    domain: Optional[str]
    metadata: dict

class RetrieveKnowledgeSkill(Skill):
    """Pure skill: semantic search over project knowledge."""

    async def execute(self, input: RetrieveKnowledgeInput) -> RetrieveKnowledgeOutput:
        query_embedding = await self.embedding_service.embed(input.query)
        results = await self.store.search(
            embedding=query_embedding,
            knowledge_types=input.knowledge_types,
            domain=input.domain,
            project_id=input.project_id,
            top_k=input.top_k,
            similarity_threshold=input.similarity_threshold,
        )
        return RetrieveKnowledgeOutput(results=results, ...)

# Example: BOM Agent using RAG
async def select_component(self, requirements: dict) -> BOMItem:
    # Step 1: Retrieve relevant historical knowledge
    knowledge = await self.run_skill("retrieve_knowledge", {
        "query": f"component selection for {requirements['function']} "
                 f"in {requirements['domain']}",
        "knowledge_types": ["component_rationale", "failure_mode"],
        "domain": requirements["domain"],
        "top_k": 5,
    })

    # Step 2: Build prompt with RAG context
    prompt = f"""Select the best component for: {requirements['description']}

## Historical Knowledge
{self._format_knowledge_context(knowledge.results)}

## Current Requirements
{json.dumps(requirements, indent=2)}

Based on the historical knowledge above and current requirements,
select and justify a component.
"""

    # Step 3: Call LLM with enriched context
    response = await self.llm.complete(prompt)
    return self.parse_bom_item(response)

def _format_knowledge_context(self, results: list[KnowledgeResult]) -> str:
    """Format knowledge results for prompt injection."""
    if not results:
        return "_No relevant historical knowledge found._"
    sections = []
    for r in results:
        sections.append(
            f"### {r.knowledge_type} (similarity: {r.similarity:.2f})\n"
            f"{r.content_text}"
        )
    return "\n\n".join(sections)

# digital_twin/knowledge/api.py

# GET /api/v1/knowledge/search
@router.get("/search")
async def search_knowledge(
    query: str,
    knowledge_types: Optional[list[str]] = Query(None),
    domain: Optional[str] = None,
    project_id: Optional[str] = None,
    top_k: int = 10,
    similarity_threshold: float = 0.6,
) -> KnowledgeSearchResponse:
    """Semantic search over project knowledge."""
    embedding = await embedding_service.embed(query)
    results = await store.search(
        embedding=embedding,
        knowledge_types=knowledge_types,
        domain=domain,
        project_id=project_id,
        top_k=top_k,
        similarity_threshold=similarity_threshold,
    )
    return KnowledgeSearchResponse(results=results, total=len(results))

# digital_twin/knowledge/reconciler.py

class KnowledgeReconciler:
    """Daily reconciliation between Neo4j and pgvector."""

    RECONCILE_TYPES = {
        "DecisionRecord": "design_decision",
        "BOMItem": "component_rationale",
        "TestExecution": "failure_mode",
        "CAPARecord": "failure_mode",
    }

    async def reconcile(self) -> ReconciliationReport:
        """Compare Neo4j node counts with pgvector embedding counts."""
        report = ReconciliationReport()
        for node_type, knowledge_type in self.RECONCILE_TYPES.items():
            neo4j_count = await self.neo4j.count_nodes(node_type)
            pgvector_count = await self.store.count_by_type(knowledge_type)
            if neo4j_count != pgvector_count:
                report.add_drift(node_type, neo4j_count, pgvector_count)
                await self._backfill_missing(node_type, knowledge_type)
        return report

async def re_embed_all(self, batch_size: int = 100) -> int:
    """Re-embed all knowledge entries with current model. Run on model change."""
    total = 0
    offset = 0
    while True:
        entries = await self.store.fetch_batch(offset=offset, limit=batch_size)
        if not entries:
            break
        texts = [e.content_text for e in entries]
        embeddings = await self.embedding_service.embed_batch(texts)
        await self.store.update_embeddings(
            ids=[e.id for e in entries],
            embeddings=embeddings,
        )
        total += len(entries)
        offset += batch_size
    return total

# domain_agents/shared/skills/ingest_knowledge/handler.py

class IngestKnowledgeInput(BaseModel):
    source_url: str              # MinIO object path
    knowledge_type: Literal["datasheet_chunk", "application_note"]
    metadata: dict               # part_number, manufacturer, tags, etc.

class IngestKnowledgeOutput(BaseModel):
    fragments_created: int
    source_id: str

class IngestKnowledgeSkill(Skill):
    """Pure skill: chunk and embed a document from MinIO."""

    async def execute(self, input: IngestKnowledgeInput) -> IngestKnowledgeOutput:
        # Download document from MinIO
        document = await self.minio.get_object(input.source_url)

        # Chunk the document
        chunks = self.chunker.chunk_pdf(document, input.metadata)

        # Embed and store all chunks
        texts = [c.text for c in chunks]
        embeddings = await self.embedding_service.embed_batch(texts)
        for chunk, embedding in zip(chunks, embeddings):
            await self.store.insert(
                knowledge_type=input.knowledge_type,
                content_text=chunk.text,
                embedding=embedding,
                metadata=chunk.metadata,
            )

        return IngestKnowledgeOutput(
            fragments_created=len(chunks),
            source_id=input.source_url,
        )

# digital_twin/knowledge/models.py

from pydantic import BaseModel
from datetime import datetime
from typing import Optional

class KnowledgeEntry(BaseModel):
    """A single knowledge embedding in pgvector."""
    id: str
    knowledge_type: str
    source_event_id: Optional[str] = None
    source_node_id: Optional[str] = None
    domain: Optional[str] = None
    phase: Optional[str] = None
    project_id: Optional[str] = None
    content_text: str
    metadata: dict = {}
    created_at: datetime
    updated_at: datetime

class KnowledgeSearchRequest(BaseModel):
    """Search request for the knowledge API."""
    query: str
    knowledge_types: Optional[list[str]] = None
    domain: Optional[str] = None
    project_id: Optional[str] = None
    top_k: int = 10
    similarity_threshold: float = 0.6

class KnowledgeSearchResult(BaseModel):
    """A single search result with similarity score."""
    id: str
    knowledge_type: str
    content_text: str
    similarity: float
    domain: Optional[str] = None
    metadata: dict = {}

class KnowledgeSearchResponse(BaseModel):
    """Response from the knowledge search API."""
    results: list[KnowledgeSearchResult]
    total: int
    query_time_ms: float

class ReconciliationReport(BaseModel):
    """Report from the daily reconciliation job."""
    timestamp: datetime
    drifts: list[DriftEntry] = []
    backfilled: int = 0

class DriftEntry(BaseModel):
    """A single drift between Neo4j and pgvector."""
    node_type: str
    neo4j_count: int
    pgvector_count: int
    difference: int

class ChunkResult(BaseModel):
    """A single chunk from document chunking."""
    text: str
    chunk_index: int
    start_token: int
    end_token: int
    metadata: dict = {}