249 lines
9.9 KiB
Python
249 lines
9.9 KiB
Python
"""Reusable per-unit render/critic retry loops.
|
|
|
|
These loops are designed for map/reduce execution where each content unit
|
|
is processed independently. They deep-copy the incoming unit state, then run
|
|
render -> critic until success or retry exhaustion.
|
|
|
|
Minimal tools contract:
|
|
- A ``ToolBox`` instance is still expected by type, but the loop itself only
|
|
relies on downstream agent calls that resolve an LLM via
|
|
``tools.get_llm_tool(state.budget_tracker)``.
|
|
- No triple-store, chunker, converter, or aggregator capabilities are required
|
|
for these atomic loops.
|
|
"""
|
|
|
|
import logging
|
|
|
|
from ontocast.agent.criticise_facts import criticise_facts
|
|
from ontocast.agent.criticise_ontology import criticise_ontology
|
|
from ontocast.agent.external_evidence import (
|
|
fetch_external_evidence_for_node,
|
|
plan_external_evidence_for_node,
|
|
)
|
|
from ontocast.agent.render_facts import render_facts
|
|
from ontocast.agent.render_ontology import render_ontology
|
|
from ontocast.onto.enum import Status, WorkflowNode
|
|
from ontocast.onto.model import ExternalEvidenceCacheEntry, ExternalEvidenceRequest
|
|
from ontocast.onto.unit_states import UnitFactsState, UnitOntologyState
|
|
from ontocast.tool.atomic import AtomicToolBox
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _resolve_max_visits_limit(state_visits: int, override: int | None) -> int:
|
|
"""Return a safe visit limit while respecting explicit overrides."""
|
|
visits = state_visits if override is None else override
|
|
return max(1, visits)
|
|
|
|
|
|
def _reset_node_evidence_context(
|
|
state: UnitFactsState | UnitOntologyState, node: WorkflowNode
|
|
) -> None:
|
|
"""Start node execution in no-search mode with empty evidence context."""
|
|
state.set_external_evidence_request(node, ExternalEvidenceRequest())
|
|
state.set_external_evidence_cache_entry(node, ExternalEvidenceCacheEntry())
|
|
state.load_external_evidence_for_node(node)
|
|
|
|
|
|
async def facts_loop(
|
|
state: UnitFactsState, tools: AtomicToolBox, max_visits_per_node: int | None = None
|
|
) -> UnitFactsState:
|
|
"""Run facts render/critic loop for one content unit.
|
|
|
|
Ontology is selected once per document in the main workflow; ontology_snapshot
|
|
is always provided by the caller.
|
|
"""
|
|
unit_state = state.model_copy(deep=True)
|
|
max_visits = _resolve_max_visits_limit(
|
|
unit_state.max_visits_per_node, max_visits_per_node
|
|
)
|
|
unit_state.max_visits_per_node = max_visits
|
|
|
|
for render_attempt in range(1, max_visits + 1):
|
|
unit_state.node_visits[WorkflowNode.TEXT_TO_FACTS] += 1
|
|
_reset_node_evidence_context(unit_state, WorkflowNode.TEXT_TO_FACTS)
|
|
unit_state = await render_facts(unit_state, tools)
|
|
if unit_state.status != Status.SUCCESS:
|
|
render_request = unit_state.get_external_evidence_request(
|
|
WorkflowNode.TEXT_TO_FACTS
|
|
)
|
|
if render_request.initiate_search:
|
|
unit_state = await plan_external_evidence_for_node(
|
|
unit_state, tools, WorkflowNode.TEXT_TO_FACTS
|
|
)
|
|
unit_state = await fetch_external_evidence_for_node(
|
|
unit_state, tools, WorkflowNode.TEXT_TO_FACTS
|
|
)
|
|
unit_state = await render_facts(unit_state, tools)
|
|
if unit_state.status == Status.SUCCESS:
|
|
logger.info(
|
|
"Unit facts render recovered with search at attempt %s/%s",
|
|
render_attempt,
|
|
max_visits,
|
|
)
|
|
# Continue to critic tier below.
|
|
else:
|
|
logger.info(
|
|
"Unit facts render failed at attempt %s/%s (with search)",
|
|
render_attempt,
|
|
max_visits,
|
|
)
|
|
continue
|
|
else:
|
|
logger.info(
|
|
"Unit facts render failed at attempt %s/%s (no search request)",
|
|
render_attempt,
|
|
max_visits,
|
|
)
|
|
continue
|
|
|
|
for critic_attempt in range(1, max_visits + 1):
|
|
unit_state.node_visits[WorkflowNode.CRITICISE_FACTS] += 1
|
|
_reset_node_evidence_context(unit_state, WorkflowNode.CRITICISE_FACTS)
|
|
unit_state = await criticise_facts(unit_state, tools)
|
|
if unit_state.status == Status.SUCCESS:
|
|
logger.info(
|
|
"Unit facts loop converged at render %s/%s critic %s/%s",
|
|
render_attempt,
|
|
max_visits,
|
|
critic_attempt,
|
|
max_visits,
|
|
)
|
|
return unit_state
|
|
|
|
critic_request = unit_state.get_external_evidence_request(
|
|
WorkflowNode.CRITICISE_FACTS
|
|
)
|
|
if not critic_request.initiate_search:
|
|
logger.info(
|
|
"Unit facts critic failed at render %s/%s critic %s/%s without search request",
|
|
render_attempt,
|
|
max_visits,
|
|
critic_attempt,
|
|
max_visits,
|
|
)
|
|
break
|
|
|
|
unit_state = await plan_external_evidence_for_node(
|
|
unit_state, tools, WorkflowNode.CRITICISE_FACTS
|
|
)
|
|
unit_state = await fetch_external_evidence_for_node(
|
|
unit_state, tools, WorkflowNode.CRITICISE_FACTS
|
|
)
|
|
unit_state = await criticise_facts(unit_state, tools)
|
|
if unit_state.status == Status.SUCCESS:
|
|
logger.info(
|
|
"Unit facts loop converged with critic search at render %s/%s critic %s/%s",
|
|
render_attempt,
|
|
max_visits,
|
|
critic_attempt,
|
|
max_visits,
|
|
)
|
|
return unit_state
|
|
|
|
continue
|
|
|
|
logger.info("Unit facts loop exhausted retries")
|
|
return unit_state
|
|
|
|
|
|
async def ontology_loop(
|
|
state: UnitOntologyState,
|
|
tools: AtomicToolBox,
|
|
max_visits_per_node: int | None = None,
|
|
) -> UnitOntologyState:
|
|
"""Run ontology render/critic loop for one content unit.
|
|
|
|
Ontology is selected once per document in the main workflow; ontology_snapshot
|
|
is always provided by the caller (may be null for fresh-ontology builds).
|
|
"""
|
|
unit_state = state.model_copy(deep=True)
|
|
max_visits = _resolve_max_visits_limit(
|
|
unit_state.max_visits_per_node, max_visits_per_node
|
|
)
|
|
unit_state.max_visits_per_node = max_visits
|
|
|
|
for render_attempt in range(1, max_visits + 1):
|
|
unit_state.node_visits[WorkflowNode.TEXT_TO_ONTOLOGY] += 1
|
|
_reset_node_evidence_context(unit_state, WorkflowNode.TEXT_TO_ONTOLOGY)
|
|
unit_state = await render_ontology(unit_state, tools)
|
|
if unit_state.status != Status.SUCCESS:
|
|
render_request = unit_state.get_external_evidence_request(
|
|
WorkflowNode.TEXT_TO_ONTOLOGY
|
|
)
|
|
if render_request.initiate_search:
|
|
unit_state = await plan_external_evidence_for_node(
|
|
unit_state, tools, WorkflowNode.TEXT_TO_ONTOLOGY
|
|
)
|
|
unit_state = await fetch_external_evidence_for_node(
|
|
unit_state, tools, WorkflowNode.TEXT_TO_ONTOLOGY
|
|
)
|
|
unit_state = await render_ontology(unit_state, tools)
|
|
if unit_state.status == Status.SUCCESS:
|
|
logger.info(
|
|
"Unit ontology render recovered with search at attempt %s/%s",
|
|
render_attempt,
|
|
max_visits,
|
|
)
|
|
else:
|
|
logger.info(
|
|
"Unit ontology render failed at attempt %s/%s (with search)",
|
|
render_attempt,
|
|
max_visits,
|
|
)
|
|
continue
|
|
else:
|
|
logger.info(
|
|
"Unit ontology render failed at attempt %s/%s (no search request)",
|
|
render_attempt,
|
|
max_visits,
|
|
)
|
|
continue
|
|
|
|
for critic_attempt in range(1, max_visits + 1):
|
|
unit_state.node_visits[WorkflowNode.CRITICISE_ONTOLOGY] += 1
|
|
_reset_node_evidence_context(unit_state, WorkflowNode.CRITICISE_ONTOLOGY)
|
|
unit_state = await criticise_ontology(unit_state, tools)
|
|
if unit_state.status == Status.SUCCESS:
|
|
logger.info(
|
|
"Unit ontology loop converged at render %s/%s critic %s/%s",
|
|
render_attempt,
|
|
max_visits,
|
|
critic_attempt,
|
|
max_visits,
|
|
)
|
|
return unit_state
|
|
|
|
critic_request = unit_state.get_external_evidence_request(
|
|
WorkflowNode.CRITICISE_ONTOLOGY
|
|
)
|
|
if not critic_request.initiate_search:
|
|
logger.info(
|
|
"Unit ontology critic failed at render %s/%s critic %s/%s without search request",
|
|
render_attempt,
|
|
max_visits,
|
|
critic_attempt,
|
|
max_visits,
|
|
)
|
|
break
|
|
|
|
unit_state = await plan_external_evidence_for_node(
|
|
unit_state, tools, WorkflowNode.CRITICISE_ONTOLOGY
|
|
)
|
|
unit_state = await fetch_external_evidence_for_node(
|
|
unit_state, tools, WorkflowNode.CRITICISE_ONTOLOGY
|
|
)
|
|
unit_state = await criticise_ontology(unit_state, tools)
|
|
if unit_state.status == Status.SUCCESS:
|
|
logger.info(
|
|
"Unit ontology loop converged with critic search at render %s/%s critic %s/%s",
|
|
render_attempt,
|
|
max_visits,
|
|
critic_attempt,
|
|
max_visits,
|
|
)
|
|
return unit_state
|
|
|
|
logger.info("Unit ontology loop exhausted retries")
|
|
return unit_state
|