Phase Contracts¶

class PhaseRunner(ABC, Generic[InputT]):
    """
    Base for every pipeline phase. Each phase accepts a typed input contract
    and returns an artifact collection. Phases may still project to the graph incrementally for crash resilience, but artifact collections are the runtime handoff.

    The Pipeline class composes PhaseRunner instances into a DAG and exposes
    .run(), .run_phase(n), and .run_from_phase(n).

    Dispatch contract (O-09)
    ------------------------
    The orchestrator needs three things from a phase that ``name`` cannot
    safely supply: which config block to fingerprint, whether to persist its
    artifacts, and which artifact view to feed it. Those used to be recovered
    by string-matching ``name`` (``name.startswith("Phase 3b")``), which made a
    display string load-bearing — renaming a phase silently changed its
    behaviour, and the match order mattered ("Phase 3b" had to be tested before
    "Phase 3"). They are now explicit class attributes.
    """

    #: Human-readable phase name. **Display only** — appears in manifests,
    #: reports and logs. Never dispatch on it.
    name: ClassVar[str]

    #: Stable machine key for this phase. Part of the orchestrator contract:
    #: the phase's config block is ``getattr(PipelineConfig, phase_key)`` and
    #: its persistence toggle is
    #: ``getattr(ExecutionConfig, f"persist_{phase_key}_artifacts")``.
    #: ``Pipeline.__init__`` validates both, so a typo fails at composition
    #: time rather than silently selecting the wrong config.
    phase_key: ClassVar[str]

    #: Artifact view class this phase consumes, e.g. ``Phase1ArtifactsView``.
    #: ``None`` means the phase is fed the raw ``PipelineInput`` — Phase 1 only.
    input_view: ClassVar[type | None] = None

    @abstractmethod
    async def run(self, input: InputT, context: ArtifactExecutionContext) -> ArtifactCollection: ...

class PipelineInput(BaseModel):
    """
    Entry point for the full pipeline.

    Parameters
    ----------
    source_paths : list[str]
        File paths to the source documents to ingest and process.
    bib_paths : list[str | Path], optional
        Optional file paths to bibliography references (.bib), by default empty list.
    metadata : dict[str, str], optional
        Arbitrary execution metadata strings, by default empty dict.
    structural_anchor : GlobalStructuralAnchor | None, optional
        Global structural anchor coordinate system (ToC outline, document summary, and/or
        global thesis) for the target document or pipeline run, by default None.
    """

    source_paths: list[str]
    bib_paths: list[str | Path] = Field(default_factory=list)
    metadata: dict[str, str] = Field(default_factory=dict)
    structural_anchor: GlobalStructuralAnchor | None = None

class L1Chunk(BaseModel):
    id: str
    text: str
    source_doc_id: str
    chapter_id: str | None = None
    sequence_index: int
    token_count: int
    embedding: list[float] | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)

class L1Document(BaseModel):
    id: str
    title: str
    source_path: str
    ingested_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    chapter_count: int
    chunk_count: int
    structural_anchor: GlobalStructuralAnchor | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)

class L2Entity(BaseModel):
    """Layer 2 (global/logical) entity representation.

    Attributes
    ----------
    id : str
        The unique identifier for the entity (typically a UUID or stable hash).
    label : str
        The entity category or class label (e.g., 'CONCEPT', 'PERSON', 'WORK').
    name : str
        The canonical name of the entity.
    description : str, optional
        A synthesized textual description of what the entity represents.
    textual_envelope : str, optional
        The original text context where this entity was first discovered.
    is_mature : bool, default False
        Flag indicating if the entity's description has gone through the Stage 2
        maturation protocol. If True, the description is stable, centroid-fused,
        and protected against further epistemic drift.
    source_chunk_ids : list of str, default []
        Identifiers of all the source chunks where this entity has been mentioned.
    """
    id: str
    label: str
    name: str
    description: str | None = None
    textual_envelope: str | None = None
    is_mature: bool = False
    confidence: float | None = None
    source_chunk_ids: list[str] = Field(default_factory=list)

class L2Triple(BaseModel):
    subject_id: str
    predicate: str
    object_id: str
    confidence: float = Field(
        description="Heterogeneous metric: could be LLM probability, reranker tau, or hardcoded 1.0."
    )
    rerank_score: float | None = None
    scope: Literal["local", "global"]
    source_chunk_id: str | None = None

class Phase3bLatentConsolidationRunner(PhaseRunner[Phase3ArtifactsView]):
    """Runner for Phase 3b: Latent Graph Consolidation.

    A fast mathematical sweep over dense vectors to merge duplicate nodes
    caused by parallel processing collisions. Runs without LLM calls.
    """
    name = "Phase 3b: Latent Graph Consolidation"
    phase_key = "phase3b"
    input_view = Phase3ArtifactsView

    def __init__(
        self,
        config: Phase3bConfig,
        *,
        embedding_model,
        graph_store: FusionGraph,
        instance_fusion: InstanceFusion | None = None,
    ) -> None:
        """Initialize the Phase 3b runner.

        Parameters
        ----------
        config : Phase3bConfig
            Configuration settings for Phase 3b.
        embedding_model : object
            The embedding model used to project textual envelopes into vector space.
        graph_store : FusionGraph
            The graph database store to interact with the constructed graph.
        instance_fusion : InstanceFusion, optional
            An optional custom InstanceFusion implementation. Defaults to LatentGraphConsolidation.
        """
        self.config = config
        self.graph_store = graph_store
        # Kept on the runner so Pipeline._method_fingerprints can see it — a
        # phase whose embedding model is invisible to invalidation is reused
        # across a model change (O-15).
        self.embedding_model = embedding_model
        self.instance_fusion = instance_fusion or LatentGraphConsolidation(
            embedding_model=embedding_model,
            dense_similarity_threshold=config.dense_similarity_threshold,
            relation_overlap_threshold=config.relation_overlap_threshold,
        )

    @property
    def event_emitter(self) -> EventEmitter:
        from episteme_pipeline.events.context import get_event_emitter
        return get_event_emitter()

    async def run(self, input: Phase3ArtifactsView, context: ArtifactExecutionContext) -> ArtifactCollection:
        """Run the Latent Graph Consolidation phase.

        Sweeps the graph L2 nodes, finds duplicates using vector similarity and 1-hop relation overlap,
        and generates Canonicalization artifacts representing the duplicate-to-canonical mapping.

        Parameters
        ----------
        input : Phase3ArtifactsView
            View of the artifacts generated up to Phase 3.
        context : ArtifactExecutionContext
            The context of the current pipeline execution.

        Returns
        -------
        ArtifactCollection
            A collection of canonicalization artifacts.
        """
        if not self.config.enabled:
            return ArtifactCollection([])

        progress_task = "Phase 3b: Latent Consolidation"
        self.event_emitter.emit(
            ProgressStarted(task_name=progress_task, total_items=1, description="Latent Consolidation")
        )
        fused_map = await self.instance_fusion.fuse(self.graph_store)
        self.event_emitter.emit(ProgressCompleted(task_name=progress_task))

        artifacts = []
        for original_id, canonical_id in fused_map.items():
            artifacts.append(
                build_canonicalization_artifact(
                    original_id, 
                    canonical_id, 
                    run_id=context.run_id, 
                    phase_name=self.name, 
                    method="phase3b.latent_consolidation"
                )
            )
        return ArtifactCollection(artifacts)

class Phase4EntityMaturationRunner(PhaseRunner[Phase3ArtifactsView]):
    """Runner for Phase 4: Entity Maturation (Batch Epistemic Synthesis).

    This phase addresses the 'epistemic drift' problem in knowledge graph construction.
    Instead of using a naive 'first-mention-wins' approach to name and describe entities,
    this phase runs after Phase 3 to mature and stabilize entities before fusion:
    1. Fetches all mention context envelopes (from `EXTRACTED_FROM` relations).
    2. Encodes envelopes using the Phase 2 Bi-Encoder.
    3. Calculates the geometric centroid of the vectors to find the core meaning.
    4. Ranks envelopes and selects the top-K closest to the centroid.
    5. Calls the LLM to synthesize a canonical description from these top-K envelopes.
    6. Persists the matured entity and marks it as mature (`is_mature = True`).

    Parameters
    ----------
    config : Phase4EntityMaturationConfig
        Configuration options for maturation (e.g. top-K envelopes to select).
    llm : LiteLLM
        LLM instance used for structured description synthesis.
    graph_store : ProcessingGraph
        The underlying Neo4j/graph database client.
    """
    name = "Phase 4: Entity Maturation (Batch Epistemic Synthesis)"
    phase_key = "phase4_maturation"
    input_view = Phase3ArtifactsView

    def __init__(
        self,
        config: Phase4EntityMaturationConfig,
        *,
        llm: Any,
        graph_store: ProcessingGraph,
        embedding_model: EmbeddingModel | None = None,
    ) -> None:
        self.config = config
        self.llm = ensure_structured_llm(llm)
        self.graph_store = graph_store
        self.embedding_model: EmbeddingModel = (
            embedding_model or _default_embedding_model()
        )

    @property
    def event_emitter(self) -> EventEmitter:
        from episteme_pipeline.events.context import get_event_emitter
        return get_event_emitter()

    async def run(self, input: Phase3ArtifactsView, context: ArtifactExecutionContext) -> ArtifactCollection:
        # Note: In a production setting with extremely large inputs, this could be triggered
        # dynamically during Phase 2 when an entity reaches a threshold N of mentions.
        # For pipeline orchestration, it runs post Phase 3.

        logger.info("Starting Entity Maturation batch synthesis...")
        from episteme_pipeline.artifacts.builders import build_linked_entity_artifact

        all_entities = await self.graph_store.get_entities()
        immature_entities = [e for e in all_entities if not getattr(e, "is_mature", False)]

        logger.info(f"Found {len(immature_entities)} immature entities out of {len(all_entities)} total.")

        matured_artifacts = []
        from itertools import batched
        import asyncio

        progress_task = f"Phase 4 Maturation: {len(immature_entities)} entities"
        if immature_entities:
            self.event_emitter.emit(
                ProgressStarted(
                    task_name=progress_task,
                    total_items=len(immature_entities),
                    description="Entity Maturation",
                )
            )

        batch_size = getattr(self.config, "batch_size", 10)

        # Process and checkpoint in batches to ensure crash resilience
        for batch_idx, entity_batch in enumerate(batched(immature_entities, batch_size)):
            logger.info(f"Maturing batch {batch_idx + 1}...")

            # Parallelize the LLM synthesis for this small batch
            tasks = [self._process_entity(entity, context.run_id) for entity in entity_batch]
            matured_results = await asyncio.gather(*tasks, return_exceptions=True)

            entities_to_upsert = []
            for entity, matured_entity in zip(entity_batch, matured_results):
                if isinstance(matured_entity, Exception):
                    logger.error(f"Entity Maturation failed for {entity.name}: {matured_entity}")
                    continue
                if matured_entity:
                    entities_to_upsert.append(matured_entity)
                    matured_artifacts.append(
                        build_linked_entity_artifact(
                            matured_entity,
                            [f"artifact::{cid}" for cid in matured_entity.source_chunk_ids],
                            run_id=context.run_id,
                            phase_name=self.name,
                            method="Phase4EntityMaturationRunner",
                        )
                    )

            # Checkpoint graph for this batch
            if entities_to_upsert:
                # Upsert in sub-batches of 500 (though max is 100 here)
                for sub_batch in batched(entities_to_upsert, 500):
                    await self.graph_store.upsert_entities(list(sub_batch))

            self.event_emitter.emit(
                ProgressAdvanced(task_name=progress_task, advance=len(entity_batch))
            )

        if immature_entities:
            self.event_emitter.emit(ProgressCompleted(task_name=progress_task))


        if matured_artifacts:
            return ArtifactCollection(matured_artifacts)
        return ArtifactCollection([])

    async def _process_entity(self, entity: L2Entity, run_id: str) -> L2Entity | None:
        # 1. Fetch all textual envelopes from EXTRACTED_FROM edges.
        # No capability check: get_entity_envelopes is part of the GraphReader
        # contract (F-11). It used to be duck-typed, which turned a missing
        # backend method into "this entity has nothing to mature".
        envelopes = await self.graph_store.get_entity_envelopes(entity.id)
        if not envelopes:
            logger.debug(f"Entity {entity.name} has no valid textual envelopes to synthesize.")
            # Mark as mature anyway to avoid reprocessing empty entities
            matured_entity = entity.model_copy(update={"is_mature": True})
            return matured_entity

        # 2. Geometric Centroid Calculation
        # E_ctx(T_i) for all envelopes
        embeddings = as_tensor(
            await self.embedding_model.aget_text_embedding_batch(envelopes)
        )  # Shape: [N, hidden_size]

        if len(envelopes) <= self.config.maturation_top_k:
            # Not enough envelopes to filter, just use all of them
            top_k_envelopes = envelopes
        else:
            # Compute geometric centroid C
            centroid = torch.mean(embeddings, dim=0, keepdim=True) # Shape: [1, hidden_size]

            # Calculate distances (we can use negative cosine similarity to rank closest)
            # Higher cosine similarity = closer to centroid
            similarities = torch.nn.functional.cosine_similarity(embeddings, centroid)

            # Get Top-k indices
            top_k_indices = torch.topk(similarities, self.config.maturation_top_k).indices.tolist()
            top_k_envelopes = [envelopes[i] for i in top_k_indices]

        # 3. Generative Fusion via LLM
        formatted_envelopes = "\n---\n".join(top_k_envelopes)

        logger.debug(f"Synthesizing description for {entity.name} using {len(top_k_envelopes)} representative contexts.")

        try:
            raw: EntitySynthesisOutput = await self.llm.predict_structured(
                EntitySynthesisOutput,
                self.config.entity_synthesis_prompts,
                strategy=self.config.entity_synthesis_decoding_strategy,
                entity_name=entity.name,
                envelopes=formatted_envelopes,
            )
            synthesized_desc = raw.description
        except Exception as exc:
            logger.warning(f"Synthesis failed for {entity.name}: {exc}")
            return None

        # 4. Update the entity
        matured_entity = entity.model_copy(update={
            "description": synthesized_desc,
            "is_mature": True
        })

        logger.info(f"Successfully matured entity {entity.name}.")

        # Emit event
        self.event_emitter.emit(
            EntityMaturationSynthesized(
                run_id=run_id,
                phase=self.name,
                entity_id=entity.id,
                entity_name=entity.name,
                envelope_count=len(envelopes),
                top_k_used=len(top_k_envelopes),
                synthesized_description=synthesized_desc,
            )
        )
        return matured_entity

class TheoryAtom(BaseModel):
    id: str
    text: str
    component_type: str
    source_chunk_id: str
    confidence: float | None = None
    plausibility: float | None = None
    entity_ids: list[str] = Field(default_factory=list)
    epistemic_status: str | None = None
    scope_type: str | None = None
    measurements: list[Measurement] = Field(default_factory=list)
    parameters: dict[str, Any] = Field(default_factory=dict)
    tenability: TenabilityResult | None = None

class Phase5ArgumentWebRunner(PhaseRunner[Phase4ArtifactsView]):
    name = "Phase 5: Inter-Document Argument Web"
    phase_key = "phase5"
    input_view = Phase4ArtifactsView

    def __init__(
        self,
        config: Phase5Config,
        *,
        embedding_model,
        graph_store: FusionGraph,
        argument_clustering: ArgumentClustering | None = None,
        theory_fusion: TheoryFusion | None = None,
    ) -> None:
        self.config = config
        self.graph_store = graph_store
        self.embedding_model = embedding_model
        self.argument_clustering = argument_clustering or EmbeddingArgumentClustering(
            embedding_model=embedding_model,
            similarity_threshold=config.fusion_similarity_threshold,
        )
        from episteme_pipeline.phases.phase5_fusion.leiden_clustering import LeidenTheoryClustering
        self.theory_fusion = theory_fusion or LeidenTheoryClustering(cluster_layer=config.cluster_layer)

    @property
    def event_emitter(self) -> EventEmitter:
        from episteme_pipeline.events.context import get_event_emitter
        return get_event_emitter()

    async def run(self, input: Phase4ArtifactsView, context: ArtifactExecutionContext) -> ArtifactCollection:
        progress_task = "Phase 5: Argument Web & Theory Fusion"
        self.event_emitter.emit(
            ProgressStarted(task_name=progress_task, total_items=1, description="Theory Fusion")
        )

        cluster_ids: list[list[str]] = []
        if self.config.argument_clustering_enabled:
            cluster_ids = await self.argument_clustering.cluster(self.graph_store)

        theory_fusion_applied = False
        if self.config.theory_fusion_enabled and self.theory_fusion is not None:
            await self.theory_fusion.fuse(self.graph_store)
            theory_fusion_applied = True

        self.event_emitter.emit(ProgressCompleted(task_name=progress_task))

        artifacts = [
            build_fusion_cluster_artifact(
                cluster,
                theory_fusion_applied,
                run_id=context.run_id,
                phase_name=self.name,
                method="phase5.fusion",
            )
            for cluster in cluster_ids
        ]
        return ArtifactCollection(artifacts)

@dataclass(slots=True)
class Phase1ArtifactsView:
    _allowed_kinds = (ArtifactKind.DOCUMENT, ArtifactKind.CHUNK)

    documents: list[L1Document]
    chunks: list[L1Chunk]

    @classmethod
    def from_collection(cls, collection: ArtifactCollection) -> "Phase1ArtifactsView":
        documents: list[L1Document] = []
        chunks: list[L1Chunk] = []
        for artifact in collection.of_kind(*cls._allowed_kinds):
            payload = artifact.payload
            if artifact.kind == ArtifactKind.DOCUMENT and isinstance(
                payload, DocumentArtifact
            ):
                documents.append(
                    L1Document(
                        id=payload.document_id,
                        title=payload.title,
                        source_path=payload.source_path,
                        ingested_at=payload.ingested_at,
                        chapter_count=int(payload.metadata.get("chapter_count", "0")),
                        chunk_count=int(payload.metadata.get("chunk_count", "0")),
                        structural_anchor=payload.structural_anchor,
                        metadata=payload.metadata,
                    )
                )
            elif artifact.kind == ArtifactKind.CHUNK and isinstance(
                payload, ChunkArtifact
            ):
                chunks.append(
                    L1Chunk(
                        id=payload.chunk_id,
                        text=payload.text,
                        source_doc_id=payload.document_id,
                        chapter_id=payload.chapter_id,
                        sequence_index=payload.sequence_index,
                        token_count=payload.token_count,
                    )
                )
        return cls(documents=documents, chunks=chunks)

@dataclass(slots=True)
class Phase2ArtifactsView:
    _allowed_kinds = (
        ArtifactKind.ENTITY_MENTION,
        ArtifactKind.LINKED_ENTITY,
        ArtifactKind.LOCAL_RELATION,
    )

    entities: list[L2Entity]
    local_triples: list[L2Triple]
    entity_count_by_type: dict[str, int]

    @classmethod
    def from_collection(cls, collection: ArtifactCollection) -> "Phase2ArtifactsView":
        entities: list[L2Entity] = []
        triples: list[L2Triple] = []
        entity_count_by_type: dict[str, int] = {}
        mention_chunk_ids_by_entity: dict[str, list[str]] = {}

        for artifact in collection.of_kind(*cls._allowed_kinds):
            payload = artifact.payload
            if artifact.kind == ArtifactKind.ENTITY_MENTION and isinstance(
                payload, EntityMentionArtifact
            ):
                entity_id, chunk_id = parse_mention_id(payload.mention_id)
                mention_chunk_ids_by_entity.setdefault(entity_id, []).append(
                    chunk_id
                )

        for artifact in collection.of_kind(*cls._allowed_kinds):
            payload = artifact.payload
            if artifact.kind == ArtifactKind.LINKED_ENTITY and isinstance(
                payload, LinkedEntityArtifact
            ):
                source_chunk_ids = mention_chunk_ids_by_entity.get(
                    payload.entity_id, []
                )
                entities.append(
                    L2Entity(
                        id=payload.entity_id,
                        label=payload.entity_type,
                        name=payload.canonical_name,
                        description=payload.description,
                        textual_envelope=payload.textual_envelope,
                        is_mature=payload.is_mature,
                        confidence=payload.confidence,
                        source_chunk_ids=source_chunk_ids,
                    )
                )
                entity_count_by_type[payload.entity_type] = (
                    entity_count_by_type.get(payload.entity_type, 0) + 1
                )
            elif artifact.kind == ArtifactKind.LOCAL_RELATION and isinstance(
                payload, LocalRelationArtifact
            ):
                triples.append(
                    L2Triple(
                        subject_id=payload.subject_entity_id,
                        predicate=payload.predicate,
                        object_id=payload.object_entity_id,
                        confidence=payload.confidence,
                        scope="local",
                        source_chunk_id=payload.source_chunk_id,
                    )
                )
        return cls(
            entities=entities,
            local_triples=triples,
            entity_count_by_type=entity_count_by_type,
        )

@dataclass(slots=True)
class Phase3ArtifactsView:
    global_triples: list[L2Triple]
    relation_type_distribution: dict[str, int]
    chunks: list[L1Chunk]

    @classmethod
    def from_collection(cls, collection: ArtifactCollection) -> "Phase3ArtifactsView":
        triples: list[L2Triple] = []
        distribution: dict[str, int] = {}
        chunks: list[L1Chunk] = []
        for artifact in collection.artifacts:
            payload = artifact.payload
            if artifact.kind == ArtifactKind.CHUNK and isinstance(
                payload, ChunkArtifact
            ):
                chunks.append(
                    L1Chunk(
                        id=payload.chunk_id,
                        text=payload.text,
                        source_doc_id=payload.document_id,
                        chapter_id=payload.chapter_id,
                        sequence_index=payload.sequence_index,
                        token_count=payload.token_count,
                    )
                )
            elif artifact.kind == ArtifactKind.GLOBAL_RELATION and isinstance(
                payload, GlobalRelationArtifact
            ):
                triples.append(
                    L2Triple(
                        subject_id=payload.subject_entity_id,
                        predicate=payload.predicate,
                        object_id=payload.object_entity_id,
                        confidence=payload.confidence,
                        scope="global",
                        source_chunk_id=payload.supporting_chunk_ids[0]
                        if payload.supporting_chunk_ids
                        else "",
                    )
                )
                distribution[payload.predicate] = (
                    distribution.get(payload.predicate, 0) + 1
                )
        return cls(
            global_triples=triples,
            relation_type_distribution=distribution,
            chunks=chunks,
        )

@dataclass(slots=True)
class Phase4ArtifactsView:
    theory_atoms: list[TheoryAtom]
    theory_relations: list[TheoryRelation]

    @classmethod
    def from_collection(cls, collection: ArtifactCollection) -> "Phase4ArtifactsView":
        components: list[TheoryAtom] = []
        relations: list[TheoryRelation] = []
        for artifact in collection.artifacts:
            payload = artifact.payload
            if artifact.kind == ArtifactKind.THEORY_ATOM and isinstance(
                payload, TheoryAtomArtifact
            ):
                components.append(
                    TheoryAtom(
                        id=payload.component_id,
                        text=payload.text,
                        component_type=payload.component_type,
                        source_chunk_id=payload.chunk_id,
                        plausibility=payload.plausibility,
                        epistemic_status=payload.epistemic_status,
                        scope_type=payload.scope_type,
                    )
                )
            elif artifact.kind == ArtifactKind.THEORY_RELATION and isinstance(
                payload, TheoryRelationArtifact
            ):
                relations.append(
                    TheoryRelation(
                        source_id=payload.source_component_id,
                        target_id=payload.target_component_id,
                        relation_type=payload.relation_type,
                        confidence=payload.confidence,
                        scope=payload.scope,
                        weight=payload.weight,
                    )
                )
        return cls(theory_atoms=components, theory_relations=relations)