Skip to content

ontocast.stategraph

Agent module for OntoCast.

This module provides a collection of agents that handle various aspects of ontology processing, including document conversion, text chunking, fact aggregation, and ontology management. Each agent is designed to perform a specific task in the ontology processing pipeline.

build_agent_graph(tools)

Build the document-level agent graph without compiling it.

Use this when you need to attach a checkpointer or store yourself, inspect the topology, or splice extra nodes in before compiling. Most callers want :func:create_agent_graph, which compiles for you.

CONVERT -> CHUNK (prepare: segment, tag, filter, size) ->

(conditional extraction)

Per-unit ontology context is assembled inside ontology_loop (not at a document-level select node). For ONTOLOGY_AND_FACTS, the full ontology block completes before the facts map runs; facts use the merged document ontology from AgentState.

Summarization has no node of its own: a unit's summary depends only on that unit, so it runs inside the extraction fan-outs. As a stage it was a barrier that made every unit wait for the slowest summary before any extraction could begin.

Parameters:

Name Type Description Default
tools ToolBox

The dependency container bound into every node.

required

Returns:

Type Description
StateGraph

The uncompiled :class:~langgraph.graph.StateGraph.

Source code in ontocast/stategraph/create.py
def build_agent_graph(tools: ToolBox) -> StateGraph:
    """Build the document-level agent graph without compiling it.

    Use this when you need to attach a checkpointer or store yourself, inspect
    the topology, or splice extra nodes in before compiling. Most callers want
    :func:`create_agent_graph`, which compiles for you.

    Flow: CONVERT -> CHUNK (prepare: segment, tag, filter, size) ->
          (conditional extraction)

    Per-unit ontology context is assembled inside ``ontology_loop`` (not at a
    document-level select node). For ``ONTOLOGY_AND_FACTS``, the full ontology
    block completes before the facts map runs; facts use the merged document
    ontology from ``AgentState``.

    Summarization has no node of its own: a unit's summary depends only on that
    unit, so it runs inside the extraction fan-outs. As a stage it was a barrier
    that made every unit wait for the slowest summary before any extraction
    could begin.

    Args:
        tools: The dependency container bound into every node.

    Returns:
        The uncompiled :class:`~langgraph.graph.StateGraph`.
    """
    workflow = StateGraph(AgentState)

    convert_document_node = partial(convert_document, tools=tools)
    chunk_text_node = partial(chunk_text, tools=tools)
    serialize_node = partial(serialize, tools=tools)

    render_ontology_node = make_render_ontology_node(tools)
    normalize_ontology_node = make_normalize_ontology_node(tools)
    consolidate_ontology_node = make_consolidate_ontology_node(tools)
    render_facts_node = make_render_facts_node(tools)
    merge_facts_node = make_merge_facts_node(tools)
    validate_facts_node = make_validate_facts_node(tools)
    structural_check_node = make_structural_check_node(tools)
    consistency_critic_node = make_consistency_critic_node(tools)

    node_callables: dict[WorkflowNode, Callable[..., Any]] = {
        WorkflowNode.CONVERT_TO_TEXT: convert_document_node,
        WorkflowNode.CHUNK: chunk_text_node,
        WorkflowNode.RENDER_ONTOLOGY_UPDATE: render_ontology_node,
        WorkflowNode.NORMALIZE_ONTOLOGY_UPDATES: normalize_ontology_node,
        WorkflowNode.CONSOLIDATE_ONTOLOGY: consolidate_ontology_node,
        WorkflowNode.RENDER_FACTS: render_facts_node,
        WorkflowNode.MERGE_FACTS: merge_facts_node,
        WorkflowNode.VALIDATE_FACTS: validate_facts_node,
        WorkflowNode.STRUCTURAL_CHECK: structural_check_node,
        WorkflowNode.CONSISTENCY_CRITIC: consistency_critic_node,
        WorkflowNode.SERIALIZE: serialize_node,
    }
    for node, callable_ in node_callables.items():
        workflow.add_node(node, _timed(str(node), callable_))
    workflow.add_edge(START, WorkflowNode.CONVERT_TO_TEXT)
    workflow.add_edge(WorkflowNode.CONVERT_TO_TEXT, WorkflowNode.CHUNK)
    workflow.add_conditional_edges(
        WorkflowNode.CHUNK,
        route_after_tag_or_chunk,
        {
            WorkflowNode.RENDER_ONTOLOGY_UPDATE: WorkflowNode.RENDER_ONTOLOGY_UPDATE,
            WorkflowNode.RENDER_FACTS: WorkflowNode.RENDER_FACTS,
        },
    )
    workflow.add_edge(
        WorkflowNode.RENDER_ONTOLOGY_UPDATE, WorkflowNode.NORMALIZE_ONTOLOGY_UPDATES
    )
    workflow.add_edge(
        WorkflowNode.NORMALIZE_ONTOLOGY_UPDATES, WorkflowNode.CONSOLIDATE_ONTOLOGY
    )
    workflow.add_edge(WorkflowNode.CONSOLIDATE_ONTOLOGY, WorkflowNode.STRUCTURAL_CHECK)
    workflow.add_edge(WorkflowNode.RENDER_FACTS, WorkflowNode.MERGE_FACTS)

    workflow.add_edge(WorkflowNode.STRUCTURAL_CHECK, WorkflowNode.CONSISTENCY_CRITIC)
    workflow.add_edge(WorkflowNode.MERGE_FACTS, WorkflowNode.VALIDATE_FACTS)
    workflow.add_edge(WorkflowNode.VALIDATE_FACTS, WorkflowNode.SERIALIZE)

    def route_after_consistency_critic(state: AgentState) -> str:
        if state.render_facts:
            return WorkflowNode.RENDER_FACTS
        return WorkflowNode.SERIALIZE

    workflow.add_conditional_edges(
        WorkflowNode.CONSISTENCY_CRITIC,
        route_after_consistency_critic,
        {
            WorkflowNode.RENDER_FACTS: WorkflowNode.RENDER_FACTS,
            WorkflowNode.SERIALIZE: WorkflowNode.SERIALIZE,
        },
    )
    workflow.add_edge(WorkflowNode.SERIALIZE, END)

    return workflow

create_agent_graph(tools, *, checkpointer=None, store=None, name=None)

Create and compile the parallel map/reduce agent graph.

Parameters:

Name Type Description Default
tools ToolBox

The dependency container bound into every node.

required
checkpointer BaseCheckpointSaver | None

Optional LangGraph checkpointer for durable execution.

None
store BaseStore | None

Optional LangGraph store for cross-thread memory.

None
name str | None

Optional graph name. Set this when embedding the graph as a node in a parent graph -- LangGraph shows unnamed subgraphs as LangGraph in traces.

None

Returns:

Type Description
CompiledStateGraph

The compiled graph, ready for ainvoke or astream.

Source code in ontocast/stategraph/create.py
def create_agent_graph(
    tools: ToolBox,
    *,
    checkpointer: BaseCheckpointSaver | None = None,
    store: BaseStore | None = None,
    name: str | None = None,
) -> CompiledStateGraph:
    """Create and compile the parallel map/reduce agent graph.

    Args:
        tools: The dependency container bound into every node.
        checkpointer: Optional LangGraph checkpointer for durable execution.
        store: Optional LangGraph store for cross-thread memory.
        name: Optional graph name. Set this when embedding the graph as a node
            in a parent graph -- LangGraph shows unnamed subgraphs as
            ``LangGraph`` in traces.

    Returns:
        The compiled graph, ready for ``ainvoke`` or ``astream``.
    """
    return build_agent_graph(tools).compile(
        checkpointer=checkpointer, store=store, name=name
    )

run_unit_pipeline(agent_state, tools) async

Run conversion, ontology, and facts loops for a single content unit.

Source code in ontocast/stategraph/unit_pipeline.py
async def run_unit_pipeline(
    agent_state: AgentState,
    tools: ToolBox,
) -> tuple[UnitOntologyState | None, UnitFactsState | None]:
    """Run conversion, ontology, and facts loops for a single content unit."""
    convert_document(agent_state, tools)
    if agent_state.failure_stage is not None or agent_state.status == Status.FAILED:
        raise DocumentConversionError(
            agent_state.failure_reason or "Document conversion failed",
            stage=str(agent_state.failure_stage),
        )

    full_text = (
        agent_state.docling_doc.export_to_markdown()
        if agent_state.docling_doc is not None
        else ""
    )
    unit = ContentUnit(
        text=full_text,
        index=0,
        doc_iri=agent_state.doc_iri,
    )
    agent_state.content_units = [unit]

    onto_result: UnitOntologyState | None = None
    facts_result: UnitFactsState | None = None

    max_visits = agent_state.max_visits

    if agent_state.render_ontology:
        ontology_state = UnitOntologyState(
            content_unit=unit,
            ontology_snapshot=_empty_snapshot(),
            ontology_patch_sources=[],
            ontology_user_instruction=agent_state.ontology_user_instruction,
            budget_tracker=deepcopy(agent_state.budget_tracker),
            max_visits_per_node=max_visits,
            max_critic_visits_per_node=(tools.config.server.max_critic_visits_per_node),
            current_domain=agent_state.current_domain,
            ontology_max_triples=tools.config.server.ontology_max_triples,
            llm_graph_format=agent_state.llm_graph_format,
            ontology_context_max_triples=tools.config.server.ontology_context_max_triples,
        )
        logger.info("run_unit_pipeline: starting ontology loop")
        ontology_context = UnitLoopContext.from_agent_state(agent_state)
        onto_result = await ontology_loop(ontology_state, tools, ontology_context)
        logger.info(
            "run_unit_pipeline: ontology loop finished (status=%s)", onto_result.status
        )
        agent_state.retrieval_metrics.update(ontology_context.retrieval_metrics)
        agent_state.budget_tracker = onto_result.budget_tracker
        if (
            onto_result.fresh_ontology is not None
            and not onto_result.fresh_ontology.is_null()
        ):
            agent_state.reduced_ontology_artifacts = [onto_result.fresh_ontology]

    facts_pre_resolved_context = (
        _facts_context_from_ontology_result(onto_result)
        if onto_result is not None
        else None
    )

    if agent_state.render_facts:
        facts_state = UnitFactsState(
            content_unit=unit,
            ontology_snapshot=_empty_snapshot(),
            ontology_patch_sources=[],
            facts_user_instruction=agent_state.facts_user_instruction,
            budget_tracker=deepcopy(agent_state.budget_tracker),
            max_visits_per_node=max_visits,
            max_critic_visits_per_node=(tools.config.server.max_critic_visits_per_node),
            llm_graph_format=agent_state.llm_graph_format,
            ontology_context_max_triples=tools.config.server.ontology_context_max_triples,
        )
        logger.info("run_unit_pipeline: starting facts loop")
        facts_context = UnitLoopContext.from_agent_state(agent_state)
        facts_result = await facts_loop(
            facts_state,
            tools,
            facts_context,
            pre_resolved_context=facts_pre_resolved_context,
        )
        logger.info(
            "run_unit_pipeline: facts loop finished (status=%s)", facts_result.status
        )
        agent_state.retrieval_metrics.update(facts_context.retrieval_metrics)
        agent_state.budget_tracker = facts_result.budget_tracker

    return onto_result, facts_result