chorusgraph 1.0.1__py3-none-any.whl

chorusgraph/__init__.py

@@ -0,0 +1,57 @@
+"""ChorusGraph — native Prism execution engine with cache, memory, and Route Ledger."""
+
+__version__ = "1.0.1"
+
+from chorusgraph.adapter import RunnableWithLedger, wrap
+from chorusgraph.cache_gate import Decision, DecisionKind, SidecarStore, gate, seed_cache_entry
+from chorusgraph.checkpoint import PrismCheckpointer, create_checkpointer, sqlite_checkpointer
+from chorusgraph.compose import ChorusStack, RedisCacheBackend
+from chorusgraph.core import END, START, CompiledGraph, Graph, NodeContext, NodeFn
+from chorusgraph.transport import TransportMode, publish_hop
+from chorusgraph.ledger import (
+    LedgerStep,
+    LedgerSink,
+    PostgresLedgerSink,
+    RouteLedger,
+    SqliteLedgerSink,
+    get_run,
+    list_runs,
+)
+from chorusgraph.memory import CortexMemoryService, get_cortex_service
+from chorusgraph.sections import CachePolicy, Section
+from chorusgraph.shadow import run_shadow_measurement
+
+__all__ = [
+    "__version__",
+    "CachePolicy",
+    "CompiledGraph",
+    "CortexMemoryService",
+    "ChorusStack",
+    "Decision",
+    "DecisionKind",
+    "END",
+    "Graph",
+    "NodeContext",
+    "NodeFn",
+    "LedgerSink",
+    "LedgerStep",
+    "TransportMode",
+    "PrismCheckpointer",
+    "RedisCacheBackend",
+    "RouteLedger",
+    "RunnableWithLedger",
+    "Section",
+    "SidecarStore",
+    "SqliteLedgerSink",
+    "START",
+    "create_checkpointer",
+    "gate",
+    "get_cortex_service",
+    "get_run",
+    "list_runs",
+    "publish_hop",
+    "run_shadow_measurement",
+    "seed_cache_entry",
+    "sqlite_checkpointer",
+    "wrap",
+]

chorusgraph/adapter/__init__.py

@@ -0,0 +1,5 @@
+"""LangGraph adapter — observe execution, emit Route Ledger."""
+
+from chorusgraph.adapter.wrap import RunnableWithLedger, wrap
+
+__all__ = ["RunnableWithLedger", "wrap"]

chorusgraph/adapter/wrap.py

@@ -0,0 +1,249 @@
+"""Adapter — native CompiledGraph ledger + legacy LangGraph fallback."""
+
+from __future__ import annotations
+
+import time
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+from chorusgraph.ledger.models import LedgerStep, RouteLedger
+from chorusgraph.ledger.sink import LedgerSink, SqliteLedgerSink
+
+_APPEND_KEYS = frozenset(
+    {
+        "rule_chain",
+        "prism_sequence",
+        "hop_metrics",
+        "vector_hops",
+        "pipeline_trace",
+        "tool_calls",
+        "conversation_history",
+        "agent_trace",
+    }
+)
+
+
+def _merge_state(base: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any]:
+    merged = dict(base)
+    for key, value in update.items():
+        if key in _APPEND_KEYS and key in merged:
+            left = merged[key]
+            if isinstance(left, list) and isinstance(value, list):
+                merged[key] = left + value
+                continue
+        if (
+            key not in _APPEND_KEYS
+            and key in merged
+            and isinstance(merged[key], list)
+            and isinstance(value, list)
+        ):
+            merged[key] = merged[key] + value
+        else:
+            merged[key] = value
+    return merged
+
+
+def _parse_edge_from_triggers(triggers: Optional[List[str]]) -> Optional[str]:
+    if not triggers:
+        return None
+    for trigger in triggers:
+        if trigger.startswith("branch:to:"):
+            return trigger.split(":", 2)[-1]
+    return None
+
+
+def _extract_grounding_score(
+    result: Dict[str, Any],
+    node_name: str,
+    *,
+    merged_state: Optional[Dict[str, Any]] = None,
+) -> Optional[float]:
+    if node_name != "writer":
+        return None
+    if result.get("memory_confidence") is not None:
+        return float(result["memory_confidence"])
+    if merged_state is not None and merged_state.get("memory_confidence") is not None:
+        return float(merged_state["memory_confidence"])
+    return None
+
+
+def _append_agent_trace_steps(
+    ledger: RouteLedger,
+    node: str,
+    result: Dict[str, Any],
+    *,
+    timestamp: datetime,
+) -> None:
+    trace = result.get("agent_trace") or []
+    for step in trace:
+        if isinstance(step, dict):
+            kind = step.get("kind") or "step"
+            content = step.get("content") or ""
+        else:
+            kind = getattr(step, "kind", "step")
+            if hasattr(kind, "value"):
+                kind = kind.value
+            content = getattr(step, "content", str(step))
+        ledger.add_step(
+            LedgerStep(
+                node=f"{node}/{kind}",
+                rule_chain=[content[:500] if content else kind],
+                timestamp=timestamp,
+            )
+        )
+
+
+def _extract_rule_chain(
+    result: Dict[str, Any],
+    node_name: str,
+    *,
+    merged_state: Optional[Dict[str, Any]] = None,
+) -> Optional[List[str]]:
+    chain = result.get("rule_chain")
+    if chain:
+        return list(chain)
+
+    sequence = result.get("prism_sequence") or []
+    if not sequence and merged_state is not None:
+        sequence = merged_state.get("prism_sequence") or []
+    for envelope in reversed(sequence):
+        if envelope.get("agent_id") == node_name and envelope.get("rule_chain"):
+            return list(envelope["rule_chain"])
+    return None
+
+
+class RunnableWithLedger:
+    """Wraps a compiled graph and ensures Route Ledger persistence."""
+
+    def __init__(
+        self,
+        compiled_graph: Any,
+        *,
+        tenant_id: str,
+        graph_id: str,
+        sink: Optional[LedgerSink] = None,
+    ) -> None:
+        self._graph = compiled_graph
+        self._tenant_id = tenant_id
+        self._graph_id = graph_id
+        self._sink = sink or SqliteLedgerSink()
+        self.last_ledger: Optional[RouteLedger] = None
+
+        if getattr(compiled_graph, "_native", False):
+            compiled_graph.attach_ledger(
+                tenant_id=tenant_id,
+                graph_id=graph_id,
+                sink=self._sink,
+            )
+
+    @property
+    def graph(self) -> Any:
+        return self._graph
+
+    def invoke(self, input: Dict[str, Any], /, **kwargs: Any) -> Dict[str, Any]:
+        if getattr(self._graph, "_native", False):
+            result = self._graph.invoke(input, **kwargs)
+            self.last_ledger = self._graph.last_ledger
+            if self.last_ledger and result.get("agent_trace"):
+                ts = datetime.now(timezone.utc)
+                for node_name in ("react_agent", "plan_solve", "validator"):
+                    if any(s.node == node_name for s in self.last_ledger.steps):
+                        _append_agent_trace_steps(self.last_ledger, node_name, result, timestamp=ts)
+                        break
+            return result
+        return self._run_legacy(input, kwargs, async_mode=False)
+
+    async def ainvoke(self, input: Dict[str, Any], /, **kwargs: Any) -> Dict[str, Any]:
+        if getattr(self._graph, "_native", False):
+            result = await self._graph.ainvoke(input, **kwargs)
+            self.last_ledger = self._graph.last_ledger
+            return result
+        return await self._run_legacy_async(input, kwargs)
+
+    def _run_legacy(self, input: Dict[str, Any], kwargs: Dict[str, Any], *, async_mode: bool) -> Dict[str, Any]:
+        run_id = str(uuid4())
+        turn_id = input.get("turn_id") or input.get("request_id")
+        ledger = RouteLedger(
+            run_id=run_id,
+            turn_id=str(turn_id) if turn_id else None,
+            tenant_id=self._tenant_id,
+            graph_id=self._graph_id,
+        )
+
+        state = dict(input)
+        final_state: Dict[str, Any] = state
+        pending: Dict[str, float] = {}
+
+        stream_kwargs = dict(kwargs)
+        stream_kwargs["stream_mode"] = "debug"
+
+        for event in self._graph.stream(state, **stream_kwargs):
+            if not isinstance(event, dict):
+                continue
+            event_type = event.get("type")
+            payload = event.get("payload") or {}
+
+            if event_type == "task":
+                node = payload.get("name")
+                if not node:
+                    continue
+                pending[node] = time.perf_counter()
+                edge_taken = _parse_edge_from_triggers(payload.get("triggers"))
+                if edge_taken and ledger.steps:
+                    ledger.steps[-1].edge_taken = edge_taken
+
+            elif event_type == "task_result":
+                node = payload.get("name")
+                if not node:
+                    continue
+                started = pending.pop(node, time.perf_counter())
+                duration_ms = int((time.perf_counter() - started) * 1000)
+                result = payload.get("result") or {}
+                if isinstance(result, dict):
+                    state = _merge_state(state, result)
+                    final_state = state
+                ts_raw = event.get("timestamp")
+                timestamp = (
+                    datetime.fromisoformat(ts_raw)
+                    if isinstance(ts_raw, str)
+                    else datetime.now(timezone.utc)
+                )
+                ledger.add_step(
+                    LedgerStep(
+                        node=node,
+                        edge_taken=None,
+                        rule_chain=_extract_rule_chain(result, node, merged_state=state),
+                        duration_ms=duration_ms,
+                        timestamp=timestamp,
+                        grounding_score=_extract_grounding_score(result, node, merged_state=state),
+                    )
+                )
+                if isinstance(result, dict) and result.get("agent_trace"):
+                    _append_agent_trace_steps(ledger, node, result, timestamp=timestamp)
+
+        self._sink.write(ledger)
+        self.last_ledger = ledger
+        return final_state
+
+    async def _run_legacy_async(self, input: Dict[str, Any], kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        return self._run_legacy(input, kwargs, async_mode=True)
+
+
+def wrap(
+    compiled_graph: Any,
+    *,
+    tenant_id: str,
+    graph_id: str,
+    sink: Optional[LedgerSink] = None,
+) -> RunnableWithLedger:
+    """Attach Route Ledger tracking to a compiled graph."""
+    return RunnableWithLedger(
+        compiled_graph,
+        tenant_id=tenant_id,
+        graph_id=graph_id,
+        sink=sink,
+    )
+
+
+__all__ = ["RunnableWithLedger", "wrap"]

chorusgraph/agents/__init__.py

@@ -0,0 +1,43 @@
+"""Agent execution patterns — unified Agent + strategies."""
+
+from chorusgraph.agents.agent import Agent
+from chorusgraph.agents.agent_node import AgentNode, agent_result_to_state, promote_to_agent
+from chorusgraph.agents.loop import AgentTraceStep, LoopOutcome, TraceKind, run_agent_loop
+from chorusgraph.agents.plan_solve import PlanSolveResult, run_plan_solve
+from chorusgraph.agents.plan_utils import PlanStep, plan_tasks
+from chorusgraph.agents.policy import (
+    BeliefPolicy,
+    BeliefPolicyNotCalibratedError,
+    PlanPolicy,
+    PlanSolveOpts,
+    ReActOpts,
+    ReflectionOpts,
+)
+from chorusgraph.agents.react import ReActResult, run_react
+from chorusgraph.agents.reflection import ReflectionResult, ValidationVerdict, run_reflection
+
+__all__ = [
+    "Agent",
+    "AgentNode",
+    "AgentTraceStep",
+    "BeliefPolicy",
+    "BeliefPolicyNotCalibratedError",
+    "LoopOutcome",
+    "PlanPolicy",
+    "PlanSolveOpts",
+    "PlanSolveResult",
+    "PlanStep",
+    "ReActOpts",
+    "ReActResult",
+    "ReflectionOpts",
+    "ReflectionResult",
+    "TraceKind",
+    "ValidationVerdict",
+    "agent_result_to_state",
+    "plan_tasks",
+    "promote_to_agent",
+    "run_agent_loop",
+    "run_plan_solve",
+    "run_react",
+    "run_reflection",
+]

chorusgraph/agents/agent.py

@@ -0,0 +1,69 @@
+"""Unified configurable Agent type (DESIGN §7.8)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, Optional
+
+from chorusgraph.agents.policy import (
+    BeliefPolicy,
+    PATTERN_DEFAULTS,
+    PatternOpts,
+    PlanPolicy,
+    PlanSolveOpts,
+    ReActOpts,
+    ReflectionOpts,
+)
+from chorusgraph.agents.strategies import get_strategy
+from chorusgraph.agents.strategies.base import AgentContext, AgentRunResult
+from chorusgraph.nodes.roles import RoleTemplate
+from chorusgraph.nodes.tool import ToolRegistry
+
+
+@dataclass
+class Agent:
+    """
+    One agent type — pattern selected via pluggable strategy config.
+
+    >>> agent = Agent(pattern="react", tools=registry, model=gemini.generate_json)
+    >>> agent.run("Compare USD/EUR and USD/GBP")
+    """
+
+    pattern: str
+    tools: ToolRegistry
+    model: Callable[[str, str], str]
+    role: Optional[RoleTemplate] = None
+    policy: PlanPolicy = field(default_factory=PlanPolicy)
+    pattern_opts: Optional[PatternOpts] = None
+    belief: BeliefPolicy = field(default_factory=BeliefPolicy)
+    llm_text: Optional[Callable[[str, str], str]] = None
+
+    def __post_init__(self) -> None:
+        if self.pattern_opts is None:
+            self.pattern_opts = PATTERN_DEFAULTS[self.pattern]  # type: ignore[assignment]
+
+    def run(
+        self,
+        question: str,
+        /,
+        *,
+        scratchpad_prefix: str = "",
+        initial_draft: str = "",
+        validate: Optional[Callable[[str], Any]] = None,
+        revise: Optional[Callable[[str, Any], str]] = None,
+    ) -> AgentRunResult:
+        strategy = get_strategy(self.pattern)
+        ctx = AgentContext(
+            question=question,
+            registry=self.tools,
+            llm_json=self.model,
+            llm_text=self.llm_text,
+            policy=self.policy,
+            pattern_opts=self.pattern_opts or PATTERN_DEFAULTS[self.pattern],  # type: ignore[arg-type]
+            belief=self.belief,
+            scratchpad_prefix=scratchpad_prefix,
+            initial_draft=initial_draft,
+            validate=validate,
+            revise=revise,
+        )
+        return strategy.run(ctx)

chorusgraph/agents/agent_node.py

@@ -0,0 +1,82 @@
+"""AgentNode — Agent IS-A Node (DESIGN §7.7)."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Optional
+
+from chorusgraph.agents.agent import Agent
+from chorusgraph.agents.strategies.base import AgentRunResult
+from chorusgraph.nodes.roles import Node, RoleTemplate, promote
+
+
+def _trace_dicts(result: AgentRunResult) -> List[Dict[str, Any]]:
+    return [s.to_dict() for s in result.trace]
+
+
+def _rule_chain_from_result(result: AgentRunResult) -> List[str]:
+    return [s.to_rule() for s in result.trace]
+
+
+def agent_result_to_state(result: AgentRunResult, *, pattern: str, prior: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+    prior = prior or {}
+    tool_results = [obs for obs in result.observations if isinstance(obs, dict) and "rate" in obs]
+    if not tool_results and result.observations:
+        tool_results = [
+            (obs.get("data") if isinstance(obs, dict) else obs) for obs in result.observations
+        ]
+    update: Dict[str, Any] = {
+        "tool_calls": list(prior.get("tool_calls") or []) + result.tool_calls,
+        "tool_results": tool_results,
+        "tool_result": tool_results[-1] if tool_results else None,
+        "needs_tool": False,
+        "agent_trace": _trace_dicts(result),
+        "rule_chain": _rule_chain_from_result(result),
+    }
+    if pattern == "react":
+        update["research_plan"] = "ReAct via Agent"
+    if pattern == "plan_solve":
+        update["plan_steps"] = [
+            {"id": s.id, "description": s.description, "tool": s.tool} for s in result.plan
+        ]
+    if pattern == "reflection":
+        update["validation"] = result.validation
+        update["response"] = result.draft
+        update["draft_response"] = result.draft
+        update["reflection_pass"] = result.passes
+    return update
+
+
+def AgentNode(
+    agent: Agent,
+    node_id: str = "agent",
+    *,
+    state_mapper: Optional[Callable[[Dict[str, Any], AgentRunResult], Dict[str, Any]]] = None,
+) -> Node:
+    """Wrap an Agent as a LangGraph-compatible Node."""
+
+    def handler(state: Dict[str, Any]) -> Dict[str, Any]:
+        if state.get("cache_hit") and state.get("tool_result") and agent.pattern == "react":
+            return {"agent_trace": [], "rule_chain": ["agent=skipped_cache_hit"]}
+        message = state.get("message") or ""
+        kwargs: Dict[str, Any] = {}
+        if agent.pattern == "reflection":
+            kwargs["initial_draft"] = state.get("draft_response") or ""
+            if state.get("_reflection_validate"):
+                kwargs["validate"] = state["_reflection_validate"]
+            if state.get("_reflection_revise"):
+                kwargs["revise"] = state["_reflection_revise"]
+        result = agent.run(message, **kwargs)
+        update = agent_result_to_state(result, pattern=agent.pattern, prior=state)
+        if state_mapper:
+            update = state_mapper(state, result) | update
+        return update
+
+    role = agent.role
+    node = Node(node_id=node_id, role=role, handler=handler)
+    return node
+
+
+def promote_to_agent(node: Node, agent: Agent) -> Node:
+    """Promote a plain Node to an Agent-backed node."""
+    bound = AgentNode(agent, node_id=node.node_id)
+    return Node(node_id=node.node_id, role=agent.role or node.role, handler=bound.handler)

chorusgraph/agents/loop.py

@@ -0,0 +1,137 @@
+"""Generic agent-loop substrate — reason ↔ act ↔ route with observable trace."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional, Protocol
+
+from chorusgraph.agents.policy import PlanPolicy
+
+
+class TraceKind(str, Enum):
+    THOUGHT = "thought"
+    ACTION = "action"
+    OBSERVATION = "observation"
+    PLAN_STEP = "plan_step"
+    REVISION = "revision"
+    ROUTER = "router"
+
+
+@dataclass
+class AgentTraceStep:
+    kind: TraceKind
+    content: str
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "kind": self.kind.value,
+            "content": self.content,
+            "metadata": self.metadata,
+        }
+
+    def to_rule(self) -> str:
+        return f"{self.kind.value}={self.content[:240]}"
+
+
+@dataclass
+class LoopOutcome:
+    finished: bool
+    trace: List[AgentTraceStep] = field(default_factory=list)
+    observations: List[Any] = field(default_factory=list)
+    tokens_used: int = 0
+    steps_used: int = 0
+    finish_reason: Optional[str] = None
+
+
+class Reasoner(Protocol):
+    def __call__(self, *, scratchpad: str, step: int) -> Dict[str, Any]: ...
+
+
+class Actor(Protocol):
+    def __call__(self, action: Dict[str, Any]) -> Any: ...
+
+
+class Router(Protocol):
+    def __call__(self, *, reason_result: Dict[str, Any], observations: List[Any]) -> str: ...
+
+
+def run_agent_loop(
+    *,
+    policy: PlanPolicy,
+    reason: Reasoner,
+    act: Optional[Actor],
+    route: Router,
+    initial_scratchpad: str = "",
+) -> LoopOutcome:
+    """
+    Generic cyclic machinery: llm (reason) ↔ tool (act) ↔ router (act-or-finish).
+
+    Each cycle appends Thought / Action / Observation steps to the trace.
+    """
+    trace: List[AgentTraceStep] = []
+    observations: List[Any] = []
+    scratchpad = initial_scratchpad
+    tokens_used = 0
+    steps_used = 0
+    finished = False
+    finish_reason: Optional[str] = None
+
+    while steps_used < policy.max_steps and policy.within_budget(tokens_used):
+        reason_result = reason(scratchpad=scratchpad, step=steps_used)
+        tokens_used += int(reason_result.get("tokens_used") or 0)
+
+        thought = (reason_result.get("thought") or "").strip()
+        if thought:
+            trace.append(AgentTraceStep(TraceKind.THOUGHT, thought))
+
+        decision = route(reason_result=reason_result, observations=observations)
+        trace.append(AgentTraceStep(TraceKind.ROUTER, decision, metadata={"step": steps_used}))
+
+        if decision == "finish":
+            finished = True
+            finish_reason = reason_result.get("finish_reason") or "router_finish"
+            steps_used += 1
+            break
+
+        if decision == "continue":
+            scratchpad += f"\nThought: {thought}\n(Policy: continuing — tool required before finish)\n"
+            steps_used += 1
+            continue
+
+        if decision == "act" and act is not None:
+            action = reason_result.get("action") or {}
+            tool_name = action.get("tool") or action.get("name") or "unknown"
+            trace.append(
+                AgentTraceStep(
+                    TraceKind.ACTION,
+                    f"{tool_name}({action.get('args') or action.get('arguments') or {}})",
+                    metadata={"tool": tool_name, "args": action.get("args") or action.get("arguments") or {}},
+                )
+            )
+            observation = act(action)
+            observations.append(observation)
+            obs_text = observation if isinstance(observation, str) else str(observation)
+            trace.append(AgentTraceStep(TraceKind.OBSERVATION, obs_text[:2000], metadata={"tool": tool_name}))
+            scratchpad += f"\nThought: {thought}\nAction: {tool_name}\nObservation: {obs_text}\n"
+        else:
+            scratchpad += f"\nThought: {thought}\n"
+            if reason_result.get("error"):
+                trace.append(AgentTraceStep(TraceKind.OBSERVATION, str(reason_result["error"])))
+
+        steps_used += 1
+
+    if not finished and steps_used >= policy.max_steps:
+        finish_reason = "max_steps"
+    elif not finished and not policy.within_budget(tokens_used):
+        finish_reason = "token_budget"
+
+    return LoopOutcome(
+        finished=finished or bool(observations),
+        trace=trace,
+        observations=observations,
+        tokens_used=tokens_used,
+        steps_used=steps_used,
+        finish_reason=finish_reason,
+    )

chorusgraph/agents/plan_solve.py

@@ -0,0 +1,50 @@
+"""Plan-Solve Agent shim."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional
+
+from chorusgraph.agents.plan_utils import PlanStep, plan_tasks, _try_compute_cross
+from chorusgraph.agents.policy import PlanPolicy, PlanSolveOpts
+from chorusgraph.nodes.tool import ToolRegistry
+
+__all__ = ["PlanStep", "PlanSolveResult", "plan_tasks", "run_plan_solve", "_try_compute_cross"]
+
+
+@dataclass
+class PlanSolveResult:
+    plan: List[PlanStep] = field(default_factory=list)
+    trace: List[Any] = field(default_factory=list)
+    observations: List[Any] = field(default_factory=list)
+    tool_calls: List[Dict[str, Any]] = field(default_factory=list)
+    failed_step: Optional[int] = None
+
+
+def run_plan_solve(
+    *,
+    question: str,
+    registry: ToolRegistry,
+    llm_json: Callable[[str, str], str],
+    policy: Optional[PlanPolicy] = None,
+    pattern_opts: Optional[PlanSolveOpts] = None,
+) -> PlanSolveResult:
+    from chorusgraph.agents.agent import Agent
+
+    policy = policy or PlanPolicy(max_steps=10)
+    opts = pattern_opts or PlanSolveOpts(max_plan_steps=policy.max_steps, on_step_failure="abort")
+    agent = Agent(
+        pattern="plan_solve",
+        tools=registry,
+        model=llm_json,
+        policy=policy,
+        pattern_opts=opts,
+    )
+    result = agent.run(question)
+    return PlanSolveResult(
+        plan=result.plan,
+        trace=result.trace,
+        observations=result.observations,
+        tool_calls=result.tool_calls,
+        failed_step=result.failed_step,
+    )