tracerazor 1.1.0__py3-none-win_amd64.whl

tracerazor/__init__.py

@@ -0,0 +1,143 @@
+"""
+TraceRazor 1.1.0
+
+TraceRazor is a local-first efficiency supervisor for production AI agents.
+Its stable core audits runs, diagnoses structural token waste, and verifies
+before/after improvements. Adaptive sampling remains an independent Labs API.
+
+1. Audit: record your agent's reasoning steps and tool calls, then submit
+   the trace to the TraceRazor analyzer to get a Token Alignment Score (TAS),
+   per-metric breakdowns, fix recommendations, and savings estimates.
+
+2. Sample (Labs): TraceRazor's AdaptiveKNode can sample K candidates per step
+   and select a consensus result. This experimental surface has no general
+   efficacy claim and is not part of the agent-supervisor GA contract.
+
+The stable product is the auditor and agent-efficiency workflow. Sampling is
+kept import-compatible as a separate Labs API.
+
+
+AUDIT QUICKSTART
+----------------
+Record steps manually with Tracer:
+
+    from tracerazor import Tracer
+
+    with Tracer(agent_name="support-agent", framework="openai") as t:
+        response = llm.invoke(prompt)
+        t.reasoning(response.text, tokens=response.usage.total_tokens)
+
+        result = lookup_order(order_id="ORD-123")
+        t.tool("lookup_order", params={"order_id": "ORD-123"},
+               output=str(result), success=True, tokens=80)
+
+    report = t.analyse()
+    print(report.summary())
+    # Prefer a same-workload `tracerazor compare` gate in CI.
+
+
+SAMPLING QUICKSTART
+-------------------
+Drop AdaptiveKNode into a LangGraph graph as a replacement for your ReAct node:
+
+    from tracerazor import AdaptiveKNode, openai_llm
+    from openai import AsyncOpenAI
+
+    llm = openai_llm(AsyncOpenAI(), model="gpt-4.1")
+    node = AdaptiveKNode(llm=llm, tools=my_tools, k_max=5)
+
+    graph = StateGraph(AgentState)
+    graph.add_node("agent", node)
+
+After the graph runs, the consensus report is available at:
+    result["consensus_report"].summary()
+
+
+AUDIT API
+---------
+Tracer             Record steps and submit for analysis.
+TraceRazorClient   Lower-level client for submitting trace dicts directly.
+TraceRazorReport   Parsed audit result with TAS score, metrics, and fixes.
+TraceStep          Data class for a single recorded step.
+
+SAMPLING API (LABS)
+-------------------
+AdaptiveKNode        LangGraph node. Samples K candidates per step and picks
+                     the consensus winner. K adapts dynamically.
+ExactMatchConsensus  Aggregates K branch proposals by exact-match comparison.
+MutationMetadata     Classifies tools as mutating vs read-only. Used by
+                     AdaptiveKNode to reset K after state-changing calls.
+NaiveKEnsemble       Baseline: runs K independent full-task agents and picks
+                     the majority result.
+SelfConsistencyBaseline
+                     Baseline: deterministic tool calls, then re-samples the
+                     final response K times and picks the best candidate.
+
+LLM ADAPTERS
+------------
+openai_llm       Create an async LLM adapter from an AsyncOpenAI client.
+anthropic_llm    Create an async LLM adapter from an AsyncAnthropic client.
+mock_llm         Deterministic mock for tests and offline demos.
+
+DATA TYPES
+----------
+BranchProposal   One candidate response from a single LLM sample.
+ConsensusResult  Output of ExactMatchConsensus.aggregate().
+ConsensusReport  Full trajectory record with consensus rates and token counts.
+Outcome          Enum: FULL_CONSENSUS, MAJORITY, DIVERGENT.
+StepRecord       Per-step detail inside a ConsensusReport.
+TokenCounts      Input, output, cached, and fresh token breakdown.
+NaiveRunResult   Result of one agent run inside NaiveKEnsemble.
+SCResult         Result of a SelfConsistencyBaseline run.
+"""
+
+# Audit
+from ._audit_tracer import Tracer
+from ._audit_client import TraceRazorClient, TraceRazorReport
+from ._audit_trace import TraceStep
+from .errors import AuditError, BelowMinStepsError, BinaryNotFoundError, TraceRazorError, VerificationError
+
+# Sampling
+from ._adaptive_k import AdaptiveKNode
+from ._consensus import BranchProposal, ConsensusResult, ExactMatchConsensus, Outcome
+from ._mutation import MutationMetadata
+from ._naive_ensemble import NaiveKEnsemble, NaiveRunResult
+from ._report import ConsensusReport, StepRecord, TokenCounts
+from ._self_consistency import SCResult, SelfConsistencyBaseline
+
+# LLM adapters
+from ._adapters import anthropic_llm, mock_llm, openai_llm
+
+__version__ = "1.1.0"
+__author__ = "Zulfaqar Hafez"
+
+__all__ = [
+    # Audit
+    "Tracer",
+    "TraceRazorClient",
+    "TraceRazorReport",
+    "TraceStep",
+    "TraceRazorError",
+    "BinaryNotFoundError",
+    "AuditError",
+    "BelowMinStepsError",
+    "VerificationError",
+    # Sampling
+    "AdaptiveKNode",
+    "ExactMatchConsensus",
+    "ConsensusReport",
+    "MutationMetadata",
+    "NaiveKEnsemble",
+    "NaiveRunResult",
+    "SelfConsistencyBaseline",
+    "SCResult",
+    "BranchProposal",
+    "ConsensusResult",
+    "Outcome",
+    "StepRecord",
+    "TokenCounts",
+    # LLM adapters
+    "openai_llm",
+    "anthropic_llm",
+    "mock_llm",
+]

tracerazor/_adapters.py

@@ -0,0 +1,193 @@
+"""LLM adapter factories for OpenAI and Anthropic clients.
+
+Usage
+-----
+from tracerazor._adapters import openai_llm, anthropic_llm
+from openai import AsyncOpenAI
+
+llm = openai_llm(AsyncOpenAI(), model="gpt-4.1")
+node = AdaptiveKNode(llm=llm, tools=my_tools)
+"""
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Optional
+
+
+# ── OpenAI ─────────────────────────────────────────────────────────────────
+
+
+def openai_llm(
+    client: Any,
+    model: str = "gpt-4.1",
+    temperature: float = 1.0,
+    **kwargs: Any,
+) -> Callable:
+    """Return an async LLM callable backed by an AsyncOpenAI client.
+
+    The callable signature is ``(messages, tool_schema) -> dict`` as expected
+    by AdaptiveKNode.
+
+    Parameters
+    ----------
+    client:
+        An ``openai.AsyncOpenAI`` instance.
+    model:
+        Model ID.  Default is gpt-4.1 (tau2-bench reference model).
+    temperature:
+        Sampling temperature.  Use > 0 for branch diversity.
+    """
+
+    async def _call(messages: List[dict], tool_schema: List[dict]) -> dict:
+        tools_param = (
+            [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": t["name"],
+                        "description": t.get("description", ""),
+                        "parameters": t.get("parameters", {"type": "object", "properties": {}}),
+                    },
+                }
+                for t in tool_schema
+            ]
+            if tool_schema
+            else None
+        )
+        response = await client.chat.completions.create(
+            model=model,
+            messages=messages,
+            tools=tools_param,
+            temperature=temperature,
+            **kwargs,
+        )
+        msg = response.choices[0].message
+        usage = response.usage
+
+        result: dict = {}
+        if msg.tool_calls:
+            tc = msg.tool_calls[0]
+            import json as _json
+            result["tool_name"] = tc.function.name
+            try:
+                result["arguments"] = _json.loads(tc.function.arguments)
+            except Exception:
+                result["arguments"] = {"raw": tc.function.arguments}
+        else:
+            result["final_answer"] = msg.content or ""
+
+        result["input_tokens"] = getattr(usage, "prompt_tokens", 0)
+        result["output_tokens"] = getattr(usage, "completion_tokens", 0)
+        # OpenAI prompt_tokens_details contains cached_tokens
+        details = getattr(usage, "prompt_tokens_details", None)
+        result["cached_tokens"] = getattr(details, "cached_tokens", 0) if details else 0
+
+        return result
+
+    return _call
+
+
+# ── Anthropic ──────────────────────────────────────────────────────────────
+
+
+def anthropic_llm(
+    client: Any,
+    model: str = "claude-sonnet-4-6",
+    temperature: float = 1.0,
+    **kwargs: Any,
+) -> Callable:
+    """Return an async LLM callable backed by an AsyncAnthropic client.
+
+    Parameters
+    ----------
+    client:
+        An ``anthropic.AsyncAnthropic`` instance.
+    model:
+        Model ID.  Default is claude-sonnet-4-6.
+    temperature:
+        Sampling temperature.
+    """
+
+    async def _call(messages: List[dict], tool_schema: List[dict]) -> dict:
+        # Separate system messages (Anthropic takes system as a top-level param)
+        system_msgs = [m for m in messages if m.get("role") == "system"]
+        other_msgs = [m for m in messages if m.get("role") != "system"]
+        system_text = " ".join(m.get("content", "") for m in system_msgs) or None
+
+        tools_param = (
+            [
+                {
+                    "name": t["name"],
+                    "description": t.get("description", ""),
+                    "input_schema": t.get("parameters", {"type": "object", "properties": {}}),
+                }
+                for t in tool_schema
+            ]
+            if tool_schema
+            else []
+        )
+
+        create_kwargs: dict = dict(
+            model=model,
+            messages=other_msgs,
+            max_tokens=4096,
+            temperature=temperature,
+            **kwargs,
+        )
+        if system_text:
+            create_kwargs["system"] = system_text
+        if tools_param:
+            create_kwargs["tools"] = tools_param
+
+        response = await client.messages.create(**create_kwargs)
+        usage = response.usage
+
+        result: dict = {}
+        for block in response.content:
+            if block.type == "tool_use":
+                result["tool_name"] = block.name
+                result["arguments"] = block.input or {}
+                break
+        else:
+            texts = [b.text for b in response.content if hasattr(b, "text")]
+            result["final_answer"] = " ".join(texts)
+
+        result["input_tokens"] = getattr(usage, "input_tokens", 0)
+        result["output_tokens"] = getattr(usage, "output_tokens", 0)
+        result["cached_tokens"] = getattr(usage, "cache_read_input_tokens", 0)
+
+        return result
+
+    return _call
+
+
+# ── Generic / test helper ──────────────────────────────────────────────────
+
+
+def mock_llm(responses: list) -> Callable:
+    """Deterministic mock LLM for tests and offline demos.
+
+    Returns the *same* response to all K branches at a given step (branches
+    share identical message histories), then advances to the next response when
+    the message history grows (indicating a new step has started).
+
+    Each item in ``responses`` is a dict as returned by the real adapters::
+
+        llm = mock_llm([
+            {"tool_name": "search", "arguments": {"q": "hello"},
+             "input_tokens": 100, "output_tokens": 20},
+            {"final_answer": "The answer is 42",
+             "input_tokens": 120, "output_tokens": 10},
+        ])
+    """
+    _state = {"idx": 0, "last_msg_len": -1}
+
+    async def _call(messages: List[dict], tool_schema: List[dict]) -> dict:
+        msg_len = len(messages)
+        # New step: message history grew → advance response pointer
+        if msg_len != _state["last_msg_len"]:
+            if _state["last_msg_len"] >= 0:
+                _state["idx"] = (_state["idx"] + 1) % len(responses)
+            _state["last_msg_len"] = msg_len
+        return dict(responses[_state["idx"]])
+
+    return _call

tracerazor/_adaptive_k.py

@@ -0,0 +1,268 @@
+"""AdaptiveKNode: core adaptive-K speculative ensemble.
+
+The node runs the agent loop with K parallel LLM samples per step.
+Consensus collapses K → 1 for tool execution (preventing duplicate side effects
+at mutation boundaries) then re-fans. K shrinks under sustained full consensus
+and resets on divergence or after a mutating call.
+"""
+from __future__ import annotations
+
+import asyncio
+from collections import Counter
+from typing import Any, Callable, Dict, List, Optional
+
+from ._consensus import BranchProposal, ConsensusResult, ExactMatchConsensus, Outcome
+from ._mutation import MutationMetadata
+from ._report import ConsensusReport, StepRecord, TokenCounts
+
+
+class AdaptiveKNode:
+    """Adaptive-K speculative ensemble as a LangGraph-compatible node.
+
+    Drop-in replacement for a ReAct-style agent node. Wrap the same LLM call
+    and tool schema; the node handles K-sampling, consensus, and mutation
+    boundary collapsing transparently.
+
+    Parameters
+    ----------
+    llm:
+        Async callable ``(messages, tool_schema) -> dict`` where the returned
+        dict must contain at least one of:
+          - ``"tool_name"`` + ``"arguments"`` (tool call proposal)
+          - ``"final_answer"`` (terminal step)
+        Optional token fields: ``"input_tokens"``, ``"output_tokens"``,
+        ``"cached_tokens"``.
+    tools:
+        ``{tool_name: async_callable}`` mapping.  Each callable receives the
+        tool arguments as keyword arguments and returns a string observation.
+    k_max:
+        Maximum parallel branches per step.  Default 5 (matches the K=5 naive
+        ensemble baseline in the benchmark).
+    k_min:
+        Floor for K under sustained full consensus.  Default 2.
+    consensus:
+        Aggregator instance.  Defaults to ``ExactMatchConsensus()``.
+    mutation_metadata:
+        ``MutationMetadata`` instance or ``{tool_name: bool}`` override dict.
+        Unknown tools are treated as mutating (safe default).
+    max_steps:
+        Hard ceiling on agent loop iterations.  Default 20.
+    """
+
+    def __init__(
+        self,
+        llm: Callable,
+        tools: Dict[str, Callable],
+        *,
+        k_max: int = 5,
+        k_min: int = 2,
+        consensus: Optional[ExactMatchConsensus] = None,
+        mutation_metadata: Optional[Any] = None,
+        max_steps: int = 20,
+    ) -> None:
+        if k_min < 1:
+            raise ValueError("k_min must be >= 1")
+        # clamp k_min to k_max so callers can set k_max=1 without pain
+        k_min = min(k_min, k_max)
+
+        self.llm = llm
+        self.tools = tools
+        self.k_max = k_max
+        self.k_min = k_min
+        self.consensus = consensus or ExactMatchConsensus()
+        self.max_steps = max_steps
+
+        if isinstance(mutation_metadata, dict):
+            mutation_metadata = MutationMetadata(overrides=mutation_metadata)
+        self.mutation: MutationMetadata = mutation_metadata or MutationMetadata()
+
+        self._last_report: Optional[ConsensusReport] = None
+
+    # ── public ─────────────────────────────────────────────────────────────
+
+    @property
+    def report(self) -> Optional[ConsensusReport]:
+        """ConsensusReport from the most recent ``__call__`` invocation."""
+        return self._last_report
+
+    async def __call__(self, state: Dict[str, Any]) -> Dict[str, Any]:
+        """LangGraph node entry point.
+
+        Reads ``state["messages"]``, runs the adaptive-K agent loop, and
+        returns an updated state dict containing:
+        - ``"messages"``: final message list (shared history)
+        - ``"consensus_report"``: the ``ConsensusReport`` for this run
+        """
+        messages: List[dict] = list(state.get("messages", []))
+        tool_schema = self._build_tool_schema()
+
+        trajectory: List[dict] = []
+        step_records: List[StepRecord] = []
+        consensus_rates: List[float] = []
+        divergences: List[int] = []
+        total_tokens = TokenCounts()
+
+        k = self.k_max
+
+        for step in range(self.max_steps):
+            # ── 1. sample K completions in parallel ───────────────────────
+            proposals = await self._sample(messages, tool_schema, k)
+
+            # ── 2. aggregate consensus ────────────────────────────────────
+            result: ConsensusResult = self.consensus.aggregate(proposals)
+            winning = result.winning_proposal
+
+            # ── 3. token accounting ───────────────────────────────────────
+            step_tokens = TokenCounts(
+                input=sum(p.input_tokens for p in proposals),
+                output=sum(p.output_tokens for p in proposals),
+                cached=sum(p.cached_tokens for p in proposals),
+            )
+            total_tokens += step_tokens
+
+            consensus_rates.append(result.consensus_rate)
+            if result.outcome != Outcome.FULL:
+                divergences.append(step)
+
+            # ── 4. build divergent-alternatives record ────────────────────
+            key_counts: Counter[str] = Counter(p._key for p in proposals)
+            alts = [
+                {"tool": p.tool_name, "args": p.arguments, "answer": p.final_answer}
+                for p in proposals
+                if p._key != result.winning_key
+                and p.branch_id
+                == min(
+                    (q.branch_id for q in proposals if q._key == p._key),
+                    default=p.branch_id,
+                )
+            ]
+
+            # ── 5. adapt K for next step ──────────────────────────────────
+            if result.outcome == Outcome.FULL:
+                k = max(self.k_min, k - 1)
+            elif result.outcome == Outcome.DIVERGENT:
+                k = self.k_max
+
+            # ── 6a. final answer → break ──────────────────────────────────
+            if winning.tool_name is None:
+                answer = winning.final_answer or ""
+                trajectory.append({"type": "final_answer", "content": answer})
+                step_records.append(
+                    StepRecord(
+                        step_number=step,
+                        outcome=result.outcome.value,
+                        consensus_rate=result.consensus_rate,
+                        k_used=result.k,
+                        executed_tool=None,
+                        executed_args=None,
+                        final_answer=answer,
+                        is_mutating=False,
+                        tokens=step_tokens,
+                        divergent_alternatives=alts,
+                    )
+                )
+                messages.append({"role": "assistant", "content": answer})
+                break
+
+            # ── 6b. tool call ─────────────────────────────────────────────
+            tool_name = winning.tool_name
+            tool_args = winning.arguments or {}
+            is_mutating = self.mutation.is_mutating(tool_name)
+
+            observation = await self._execute_tool(tool_name, tool_args)
+
+            trajectory.append(
+                {
+                    "type": "tool_call",
+                    "tool": tool_name,
+                    "args": tool_args,
+                    "observation": observation,
+                    "mutating": is_mutating,
+                    "consensus_rate": result.consensus_rate,
+                }
+            )
+            step_records.append(
+                StepRecord(
+                    step_number=step,
+                    outcome=result.outcome.value,
+                    consensus_rate=result.consensus_rate,
+                    k_used=result.k,
+                    executed_tool=tool_name,
+                    executed_args=tool_args,
+                    final_answer=None,
+                    is_mutating=is_mutating,
+                    tokens=step_tokens,
+                    divergent_alternatives=alts,
+                )
+            )
+
+            # append call + observation to shared history
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [{"name": tool_name, "arguments": tool_args}],
+                }
+            )
+            messages.append({"role": "tool", "name": tool_name, "content": observation})
+
+            # re-fan after mutation boundary
+            if is_mutating:
+                k = self.k_max
+
+        self._last_report = ConsensusReport(
+            trajectory=trajectory,
+            consensus_rate=consensus_rates,
+            divergences=divergences,
+            tokens=total_tokens,
+            steps=step_records,
+        )
+
+        return {**state, "messages": messages, "consensus_report": self._last_report}
+
+    # ── private helpers ─────────────────────────────────────────────────────
+
+    def _build_tool_schema(self) -> List[dict]:
+        return [
+            {
+                "name": name,
+                "description": getattr(fn, "__doc__", "") or "",
+                "parameters": getattr(fn, "_schema", {}),
+            }
+            for name, fn in self.tools.items()
+        ]
+
+    async def _sample(
+        self,
+        messages: List[dict],
+        tool_schema: List[dict],
+        k: int,
+    ) -> List[BranchProposal]:
+        """Call the LLM k times in parallel with identical inputs."""
+        tasks = [self.llm(list(messages), tool_schema) for _ in range(k)]
+        raw: list[dict] = await asyncio.gather(*tasks)
+        return [
+            BranchProposal(
+                branch_id=i,
+                tool_name=r.get("tool_name"),
+                arguments=r.get("arguments"),
+                final_answer=r.get("final_answer"),
+                input_tokens=r.get("input_tokens", 0),
+                output_tokens=r.get("output_tokens", 0),
+                cached_tokens=r.get("cached_tokens", 0),
+            )
+            for i, r in enumerate(raw)
+        ]
+
+    async def _execute_tool(self, tool_name: str, args: dict) -> str:
+        """Execute a registered tool once and return the observation string."""
+        fn = self.tools.get(tool_name)
+        if fn is None:
+            return f"[error: tool '{tool_name}' not registered]"
+        try:
+            result = fn(**args)
+            if asyncio.iscoroutine(result):
+                result = await result
+            return str(result)
+        except Exception as exc:
+            return f"[error: {exc}]"