orchestrata_ultimate.knowledge_graph.context_validator_loop

Reference › orchestrata_ultimate › knowledge_graph › context_validator_loop

The validator feedback loop (Spec 008, REQ-008-013). When an LLM response fails graph-context validation because core entities are missing from its citation list, ValidatorLoop re-injects a wider graph context and calls the model again, up to a bounded number of retries. It is constructed with a context_injector ContextInjector, an optional factory that builds a graph_context_validator GraphContextValidator for a given context, and max_retries (default DEFAULT_MAX_RETRIES). run is async: it takes the payload and an awaitable call_llm, injects context for the task type, calls the model and validates the answer. On failure _expansion_for_attempt returns an increasingly broad ContextExpansion for each retry. The result is a ValidatorLoopResult holding the final response, its ValidationResult, the retry count and whether the context was expanded. The volk base class in the agents package is its only importer. The loop reads the operational graph only through the injector.

At a glance

Class diagram

Dependency diagram

API

Classes

ValidatorLoopResult (dataclass)

Decorators: dataclass · Line 33

Outcome of a validator feedback loop.

ValidatorLoop

Decorators: traces(spec='008', story='US-003', reqs=['REQ-008-013']) · Line 43

Retry LLM calls with expanded KG context on missing citations.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/knowledge_graph/context_validator_loop.py

"""Validator feedback loop.

When an LLM response fails graph-context validation because core entities are
missing from the citation list, this loop re-injects an expanded context and
retries the call up to a bounded limit.
"""

from __future__ import annotations

import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.knowledge_graph.graph_context_validator import (
    GraphContextValidator,
    ValidationResult,
)

if TYPE_CHECKING:
    from orchestrata_ultimate.knowledge_graph.context_injector import (
        ContextExpansion,
        ContextInjector,
    )

_logger = logging.getLogger(__name__)

DEFAULT_MAX_RETRIES = 2


@dataclass
class ValidatorLoopResult:
    """Outcome of a validator feedback loop."""

    response: str
    validation: ValidationResult
    retries: int = 0
    expanded: bool = False


@traces(spec="008", story="US-003", reqs=["REQ-008-013"])
class ValidatorLoop:
    """Retry LLM calls with expanded KG context on missing citations.

    Args:
        injector: ``ContextInjector`` used to (re-)inject graph context.
        validator_factory: Callable that builds a ``GraphContextValidator``
            from a graph-context string. Defaults to ``GraphContextValidator``.
        max_retries: Maximum number of expansion retries. The loop always makes
            at least one LLM call.
    """

    def __init__(
        self,
        injector: ContextInjector,
        validator_factory: Callable[[str], GraphContextValidator] | None = None,
        max_retries: int = DEFAULT_MAX_RETRIES,
    ) -> None:
        self._injector = injector
        self._validator_factory = validator_factory or GraphContextValidator
        self._max_retries = max(max_retries, 0)

    def _expansion_for_attempt(self, attempt: int) -> ContextExpansion:
        """Return increasingly broad context expansion per retry."""
        from orchestrata_ultimate.knowledge_graph.context_injector import (
            ContextExpansion,
        )

        return ContextExpansion(
            extra_hops=attempt + 1,
            include_tested_by=True,
            include_requirement_closure=True,
            limit_multiplier=attempt + 2,
        )

    async def run(
        self,
        payload: dict[str, Any],
        call_llm: Callable[..., Awaitable[str]],
        *,
        task_type: Any | None = None,
        system_prompt: str = "",
        user_message: str = "",
        call_kwargs: dict[str, Any] | None = None,
    ) -> ValidatorLoopResult:
        """Call ``call_llm``, validate, and retry with expanded context if needed.

        The first call uses the normal injected context. Each subsequent retry
        applies a broader ``ContextExpansion`` and prepends a retry prompt that
        lists the missing core citations.

        Args:
            payload: CXP packet payload to inject context into.
            call_llm: Async callable for the LLM; receives ``system_prompt``,
                ``user_message``, and merged keyword arguments (including
                ``graph_context`` when available).
            task_type: Optional task-type enum forwarded to the injector.
            system_prompt: System prompt for the LLM.
            user_message: User message for the LLM.
            call_kwargs: Additional kwargs merged into every LLM call.

        Returns:
            ValidatorLoopResult with the final response, validation result, and
            retry count.
        """
        call_kwargs = call_kwargs or {}
        working_payload = dict(payload)

        injected = self._injector.inject(
            working_payload, task_type=task_type, expansion=None
        )
        graph_context = injected.get("graph_context", "")
        validator = self._validator_factory(graph_context)

        last_response = ""
        last_result = ValidationResult(passed=False)

        for attempt in range(self._max_retries + 1):
            kwargs = dict(call_kwargs)
            if graph_context:
                kwargs["graph_context"] = graph_context

            response = await call_llm(system_prompt, user_message, **kwargs)
            last_response = response
            result = validator.validate(response)
            last_result = result

            if result.passed:
                return ValidatorLoopResult(
                    response=response,
                    validation=result,
                    retries=attempt,
                    expanded=attempt > 0,
                )

            if attempt == self._max_retries:
                _logger.warning(
                    "Validator loop exhausted after %d retries; response still missing %s",
                    attempt,
                    result.missing_core_citations,
                )
                break

            expansion = self._expansion_for_attempt(attempt)
            _logger.info(
                "Validator loop retry %d: expanding context (extra_hops=%d, "
                "tested_by=%s, closure=%s, limit_multiplier=%d)",
                attempt + 1,
                expansion.extra_hops,
                expansion.include_tested_by,
                expansion.include_requirement_closure,
                expansion.limit_multiplier,
            )

            working_payload = dict(payload)
            injected = self._injector.inject(
                working_payload,
                task_type=task_type,
                expansion=expansion,
            )
            graph_context = injected.get("graph_context", "")
            validator = self._validator_factory(graph_context)

            # Prepend the validator's retry prompt to the user message so the
            # next call explicitly lists the missing citations.
            retry_prompt = validator.build_retry_prompt(user_message, result)
            user_message = retry_prompt

        return ValidatorLoopResult(
            response=last_response,
            validation=last_result,
            retries=self._max_retries,
            expanded=self._max_retries > 0,
        )