casebench 0.1.0__py3-none-any.whl

casebench/capture/otlp/mapping.py

@@ -0,0 +1,105 @@
+"""Turn spans in the GenAI or OpenInference conventions into steps (decision 0023)."""
+
+import json
+import re
+from typing import Any
+
+from casebench.capture.otlp.decode import RawSpan
+from casebench.capture.step import NO_DIAGNOSTIC, Step, StepKind
+
+LLM_OPERATIONS = {"chat", "text_completion", "generate_content"}
+
+
+def _indexed_ids(attributes: dict[str, Any], prefix: str) -> list[str]:
+    """OpenInference flattens lists: `<prefix>.0.document.id`, `<prefix>.1.document.id`."""
+    pattern = re.compile(rf"^{re.escape(prefix)}\.(\d+)\.document\.id$")
+    found = sorted(
+        (int(m.group(1)), str(v)) for k, v in attributes.items() if (m := pattern.match(k))
+    )
+    return [value for _, value in found]
+
+
+def _documents(value: Any) -> list[str]:
+    """`gen_ai.retrieval.documents`: a list of objects with `id`, possibly sent as JSON text."""
+    if isinstance(value, str):
+        try:
+            value = json.loads(value)
+        except json.JSONDecodeError:
+            return []
+    return [str(item["id"]) for item in value or [] if isinstance(item, dict) and "id" in item]
+
+
+def _kind(attributes: dict[str, Any]) -> StepKind | None:
+    operation = attributes.get("gen_ai.operation.name")
+    span_kind = str(attributes.get("openinference.span.kind", "")).upper()
+    if operation in LLM_OPERATIONS or span_kind == "LLM":
+        return StepKind.LLM
+    if operation == "execute_tool" or span_kind == "TOOL":
+        return StepKind.TOOL
+    if operation == "retrieval" or span_kind == "RETRIEVER":
+        return StepKind.RETRIEVAL
+    if (
+        operation == "rerank"
+        or span_kind == "RERANKER"
+        or "casebench.rerank.output_ids" in attributes
+    ):
+        return StepKind.RERANK
+    return None
+
+
+def _ids(kind: StepKind, a: dict[str, Any]) -> tuple[list[str], list[str]]:
+    if kind is StepKind.RETRIEVAL:
+        return [], _documents(a.get("gen_ai.retrieval.documents")) or _indexed_ids(
+            a, "retrieval.documents"
+        )
+    if kind is StepKind.RERANK:
+        ids_in = [str(i) for i in a.get("casebench.rerank.input_ids") or []]
+        ids_out = [str(i) for i in a.get("casebench.rerank.output_ids") or []]
+        return (
+            ids_in or _indexed_ids(a, "reranker.input_documents"),
+            ids_out or _indexed_ids(a, "reranker.output_documents"),
+        )
+    return [], []
+
+
+def _int(value: Any) -> int | None:
+    return int(value) if value is not None else None
+
+
+def to_step(span: RawSpan) -> Step | None:
+    a = span.attributes
+    kind = _kind(a)
+    if kind is None:
+        return None
+    ids_in, ids_out = _ids(kind, a)
+    searching = kind in (StepKind.RETRIEVAL, StepKind.RERANK)
+    return Step(
+        kind=kind,
+        name=a.get("gen_ai.tool.name") or a.get("tool.name") or span.name,
+        provider=a.get("gen_ai.provider.name") or a.get("gen_ai.system") or a.get("llm.provider"),
+        model=a.get("gen_ai.response.model")
+        or a.get("gen_ai.request.model")
+        or a.get("llm.model_name"),
+        input_tokens=_int(a.get("gen_ai.usage.input_tokens", a.get("llm.token_count.prompt"))),
+        output_tokens=_int(
+            a.get("gen_ai.usage.output_tokens", a.get("llm.token_count.completion"))
+        ),
+        doc_ids_in=ids_in,
+        doc_ids_out=ids_out,
+        started=span.started,
+        ended=span.ended,
+        duration_ms=(span.ended - span.started) * 1000,
+        note=NO_DIAGNOSTIC if searching and not ids_out else None,
+    )
+
+
+def without_duplicates(span_steps: list[Step], proxied: list[Step]) -> list[Step]:
+    """A model call the proxy saw is kept from the proxy; the span describing it is dropped."""
+    calls = [s for s in proxied if s.kind is StepKind.LLM and s.started is not None]
+
+    def seen_by_proxy(step: Step) -> bool:
+        if step.kind is not StepKind.LLM or step.started is None or step.ended is None:
+            return False
+        return any(step.started <= c.started <= step.ended for c in calls if c.started is not None)
+
+    return [step for step in span_steps if not seen_by_proxy(step)]

casebench/capture/parsers/__init__.py

@@ -0,0 +1 @@
+"""Built-in capture parsers, one per upstream API."""

casebench/capture/parsers/anthropic/__init__.py

@@ -0,0 +1,86 @@
+"""Read Anthropic Messages API calls, streamed or not, into model-call steps."""
+
+import json
+from typing import Any, ClassVar
+
+from casebench.capture.sizes import content_size
+from casebench.capture.step import HttpExchange, Step, StepKind, ToolCall, ToolResult
+
+
+def _tool_results(request: Any) -> list[ToolResult]:
+    results: list[ToolResult] = []
+    for message in (request or {}).get("messages", []):
+        if message.get("role") != "user" or not isinstance(message.get("content"), list):
+            continue
+        for block in message["content"]:
+            if isinstance(block, dict) and block.get("type") == "tool_result":
+                size = content_size(block.get("content", ""))
+                results.append(ToolResult(id=block["tool_use_id"], size=size))
+    return results
+
+
+def _from_message(message: dict[str, Any]) -> tuple[str | None, int | None, int | None, list]:
+    usage = message.get("usage") or {}
+    tools = [
+        ToolCall(id=block["id"], name=block["name"], arguments=block.get("input"))
+        for block in message.get("content") or []
+        if block.get("type") == "tool_use"
+    ]
+    return message.get("model"), usage.get("input_tokens"), usage.get("output_tokens"), tools
+
+
+def _from_events(events: list[Any]) -> tuple[str | None, int | None, int | None, list]:
+    model = input_tokens = output_tokens = None
+    blocks: dict[int, dict[str, Any]] = {}
+    for event in events:
+        kind = event.get("type")
+        if kind == "message_start":
+            model = event["message"].get("model")
+            input_tokens = (event["message"].get("usage") or {}).get("input_tokens")
+        elif kind == "content_block_start" and event["content_block"].get("type") == "tool_use":
+            block = event["content_block"]
+            blocks[event["index"]] = {"id": block["id"], "name": block["name"], "json": ""}
+        elif kind == "content_block_delta" and event["delta"].get("type") == "input_json_delta":
+            if event["index"] in blocks:
+                blocks[event["index"]]["json"] += event["delta"].get("partial_json", "")
+        elif kind == "message_delta":
+            output_tokens = (event.get("usage") or {}).get("output_tokens", output_tokens)
+    tools = [
+        ToolCall(id=b["id"], name=b["name"], arguments=json.loads(b["json"]) if b["json"] else {})
+        for _, b in sorted(blocks.items())
+    ]
+    return model, input_tokens, output_tokens, tools
+
+
+class Parser:
+    hosts: ClassVar[list[str]] = ["api.anthropic.com"]
+    packages: ClassVar[list[str]] = ["anthropic", "@anthropic-ai/sdk"]
+    # The client classes whose calls this parser reads, for `casebench init` (decision 0040).
+    clients: ClassVar[list[str]] = ["anthropic.Anthropic", "anthropic.AsyncAnthropic"]
+    default_route: ClassVar[dict[str, str]] = {
+        "name": "anthropic",
+        "upstream": "https://api.anthropic.com",
+        "env": "ANTHROPIC_BASE_URL",
+    }
+
+    def parse(self, exchange: HttpExchange) -> list[Step]:
+        if not exchange.url.rstrip("/").endswith("/v1/messages"):
+            return []
+        if exchange.response_events is not None:
+            model, tokens_in, tokens_out, tools = _from_events(exchange.response_events)
+        else:
+            model, tokens_in, tokens_out, tools = _from_message(exchange.response_body or {})
+        return [
+            Step(
+                kind=StepKind.LLM,
+                provider="anthropic",
+                model=model or (exchange.request_body or {}).get("model"),
+                input_tokens=tokens_in,
+                output_tokens=tokens_out,
+                requested_tools=tools,
+                tool_results=_tool_results(exchange.request_body),
+                started=exchange.started,
+                ended=exchange.ended,
+                duration_ms=exchange.duration_ms,
+            )
+        ]

casebench/capture/parsers/anthropic/casebench-plugin.json

@@ -0,0 +1,5 @@
+{
+  "kind": "parser",
+  "name": "anthropic",
+  "object": "casebench.capture.parsers.anthropic:Parser"
+}

casebench/capture/parsers/anthropic/samples/plain-with-tool-result.json

@@ -0,0 +1,76 @@
+{
+  "exchange": {
+    "method": "POST",
+    "url": "https://api.anthropic.com/v1/messages",
+    "status": 200,
+    "request_body": {
+      "model": "claude-sonnet-5-5",
+      "max_tokens": 1024,
+      "messages": [
+        {
+          "role": "user",
+          "content": "Is the dismissal valid?"
+        },
+        {
+          "role": "assistant",
+          "content": [
+            {
+              "type": "tool_use",
+              "id": "toolu_1",
+              "name": "search",
+              "input": {
+                "query": "dismissal"
+              }
+            }
+          ]
+        },
+        {
+          "role": "user",
+          "content": [
+            {
+              "type": "tool_result",
+              "tool_use_id": "toolu_1",
+              "content": "3 decisions found"
+            }
+          ]
+        }
+      ]
+    },
+    "response_body": {
+      "type": "message",
+      "model": "claude-sonnet-5-5",
+      "content": [
+        {
+          "type": "text",
+          "text": "No."
+        }
+      ],
+      "usage": {
+        "input_tokens": 200,
+        "output_tokens": 40
+      }
+    },
+    "route": "anthropic",
+    "started": 1760000002.5,
+    "ended": 1760000003.5,
+    "duration_ms": 1000.0
+  },
+  "expected": [
+    {
+      "kind": "llm",
+      "provider": "anthropic",
+      "model": "claude-sonnet-5-5",
+      "input_tokens": 200,
+      "output_tokens": 40,
+      "tool_results": [
+        {
+          "id": "toolu_1",
+          "size": 17
+        }
+      ],
+      "started": 1760000002.5,
+      "ended": 1760000003.5,
+      "duration_ms": 1000.0
+    }
+  ]
+}

casebench/capture/parsers/anthropic/samples/streamed-tool-call.json

@@ -0,0 +1,97 @@
+{
+  "exchange": {
+    "method": "POST",
+    "url": "https://api.anthropic.com/v1/messages",
+    "status": 200,
+    "request_body": {
+      "model": "claude-sonnet-5-5",
+      "max_tokens": 1024,
+      "stream": true,
+      "messages": [
+        {
+          "role": "user",
+          "content": "Is the dismissal valid?"
+        }
+      ]
+    },
+    "response_events": [
+      {
+        "type": "message_start",
+        "message": {
+          "model": "claude-sonnet-5-5",
+          "usage": {
+            "input_tokens": 120,
+            "output_tokens": 1
+          }
+        }
+      },
+      {
+        "type": "content_block_start",
+        "index": 0,
+        "content_block": {
+          "type": "tool_use",
+          "id": "toolu_1",
+          "name": "search",
+          "input": {}
+        }
+      },
+      {
+        "type": "content_block_delta",
+        "index": 0,
+        "delta": {
+          "type": "input_json_delta",
+          "partial_json": "{\"query\": "
+        }
+      },
+      {
+        "type": "content_block_delta",
+        "index": 0,
+        "delta": {
+          "type": "input_json_delta",
+          "partial_json": "\"dismissal\"}"
+        }
+      },
+      {
+        "type": "content_block_stop",
+        "index": 0
+      },
+      {
+        "type": "message_delta",
+        "delta": {
+          "stop_reason": "tool_use"
+        },
+        "usage": {
+          "output_tokens": 30
+        }
+      },
+      {
+        "type": "message_stop"
+      }
+    ],
+    "route": "anthropic",
+    "started": 1760000000.0,
+    "ended": 1760000001.5,
+    "duration_ms": 1500.0
+  },
+  "expected": [
+    {
+      "kind": "llm",
+      "provider": "anthropic",
+      "model": "claude-sonnet-5-5",
+      "input_tokens": 120,
+      "output_tokens": 30,
+      "requested_tools": [
+        {
+          "id": "toolu_1",
+          "name": "search",
+          "arguments": {
+            "query": "dismissal"
+          }
+        }
+      ],
+      "started": 1760000000.0,
+      "ended": 1760000001.5,
+      "duration_ms": 1500.0
+    }
+  ]
+}

casebench/capture/parsers/gemini/__init__.py

@@ -0,0 +1,77 @@
+"""Read Gemini API generateContent calls, streamed or not, into model-call steps.
+
+Gemini may omit function-call ids; a call without one is paired with its result by the
+function's name.
+"""
+
+import re
+from typing import Any, ClassVar
+
+from casebench.capture.sizes import content_size
+from casebench.capture.step import HttpExchange, Step, StepKind, ToolCall, ToolResult
+
+METHOD = re.compile(r"/models/([^/:]+):(generateContent|streamGenerateContent)")
+
+
+def _tool_results(request: Any) -> list[ToolResult]:
+    results: list[ToolResult] = []
+    for content in (request or {}).get("contents", []):
+        for part in content.get("parts") or []:
+            answer = part.get("functionResponse")
+            if answer:
+                key = answer.get("id") or answer["name"]
+                results.append(ToolResult(id=key, size=content_size(answer.get("response", {}))))
+    return results
+
+
+def _read(chunks: list[Any]) -> tuple[str | None, int | None, int | None, list[ToolCall]]:
+    model = prompt = completion = None
+    tools: list[ToolCall] = []
+    for chunk in chunks:
+        model = chunk.get("modelVersion") or model
+        usage = chunk.get("usageMetadata") or {}
+        prompt = usage.get("promptTokenCount", prompt)
+        completion = usage.get("candidatesTokenCount", completion)
+        for candidate in chunk.get("candidates") or []:
+            for part in (candidate.get("content") or {}).get("parts") or []:
+                call = part.get("functionCall")
+                if call:
+                    key = call.get("id") or call["name"]
+                    tools.append(ToolCall(id=key, name=call["name"], arguments=call.get("args")))
+    return model, prompt, completion, tools
+
+
+class Parser:
+    hosts: ClassVar[list[str]] = ["generativelanguage.googleapis.com"]
+    packages: ClassVar[list[str]] = ["google-genai", "google-generativeai", "@google/genai"]
+    # The client classes whose calls this parser reads, for `casebench init` (decision 0040).
+    clients: ClassVar[list[str]] = ["google.genai.Client"]
+    default_route: ClassVar[dict[str, str]] = {
+        "name": "gemini",
+        "upstream": "https://generativelanguage.googleapis.com",
+        "env": "GOOGLE_GEMINI_BASE_URL",
+    }
+
+    def parse(self, exchange: HttpExchange) -> list[Step]:
+        method = METHOD.search(exchange.url)
+        if method is None:
+            return []
+        if exchange.response_events is not None:
+            chunks = exchange.response_events
+        else:
+            chunks = [exchange.response_body or {}]
+        model, prompt, completion, tools = _read(chunks)
+        return [
+            Step(
+                kind=StepKind.LLM,
+                provider="gemini",
+                model=model or method.group(1),
+                input_tokens=prompt,
+                output_tokens=completion,
+                requested_tools=tools,
+                tool_results=_tool_results(exchange.request_body),
+                started=exchange.started,
+                ended=exchange.ended,
+                duration_ms=exchange.duration_ms,
+            )
+        ]

casebench/capture/parsers/gemini/casebench-plugin.json

@@ -0,0 +1,5 @@
+{
+  "kind": "parser",
+  "name": "gemini",
+  "object": "casebench.capture.parsers.gemini:Parser"
+}

casebench/capture/parsers/gemini/samples/plain-with-function-response.json

@@ -0,0 +1,86 @@
+{
+  "exchange": {
+    "method": "POST",
+    "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-3-pro:generateContent",
+    "status": 200,
+    "request_body": {
+      "contents": [
+        {
+          "role": "user",
+          "parts": [
+            {
+              "text": "Is the dismissal valid?"
+            }
+          ]
+        },
+        {
+          "role": "model",
+          "parts": [
+            {
+              "functionCall": {
+                "name": "search",
+                "args": {
+                  "query": "dismissal"
+                }
+              }
+            }
+          ]
+        },
+        {
+          "role": "user",
+          "parts": [
+            {
+              "functionResponse": {
+                "name": "search",
+                "response": {
+                  "found": 3
+                }
+              }
+            }
+          ]
+        }
+      ]
+    },
+    "response_body": {
+      "candidates": [
+        {
+          "content": {
+            "role": "model",
+            "parts": [
+              {
+                "text": "No."
+              }
+            ]
+          }
+        }
+      ],
+      "usageMetadata": {
+        "promptTokenCount": 120,
+        "candidatesTokenCount": 4
+      },
+      "modelVersion": "gemini-3-pro"
+    },
+    "route": "gemini",
+    "started": 1760000002.0,
+    "ended": 1760000003.0,
+    "duration_ms": 1000.0
+  },
+  "expected": [
+    {
+      "kind": "llm",
+      "provider": "gemini",
+      "model": "gemini-3-pro",
+      "input_tokens": 120,
+      "output_tokens": 4,
+      "tool_results": [
+        {
+          "id": "search",
+          "size": 12
+        }
+      ],
+      "started": 1760000002.0,
+      "ended": 1760000003.0,
+      "duration_ms": 1000.0
+    }
+  ]
+}

casebench/capture/parsers/gemini/samples/streamed-function-call.json

@@ -0,0 +1,86 @@
+{
+  "exchange": {
+    "method": "POST",
+    "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-3-pro:streamGenerateContent?alt=sse",
+    "status": 200,
+    "request_body": {
+      "contents": [
+        {
+          "role": "user",
+          "parts": [
+            {
+              "text": "Is the dismissal valid?"
+            }
+          ]
+        }
+      ]
+    },
+    "response_events": [
+      {
+        "candidates": [
+          {
+            "content": {
+              "role": "model",
+              "parts": [
+                {
+                  "functionCall": {
+                    "name": "search",
+                    "args": {
+                      "query": "dismissal"
+                    }
+                  }
+                }
+              ]
+            }
+          }
+        ],
+        "modelVersion": "gemini-3-pro"
+      },
+      {
+        "candidates": [
+          {
+            "content": {
+              "role": "model",
+              "parts": [
+                {
+                  "text": ""
+                }
+              ]
+            },
+            "finishReason": "STOP"
+          }
+        ],
+        "usageMetadata": {
+          "promptTokenCount": 80,
+          "candidatesTokenCount": 9
+        },
+        "modelVersion": "gemini-3-pro"
+      }
+    ],
+    "route": "gemini",
+    "started": 1760000000.0,
+    "ended": 1760000001.0,
+    "duration_ms": 1000.0
+  },
+  "expected": [
+    {
+      "kind": "llm",
+      "provider": "gemini",
+      "model": "gemini-3-pro",
+      "input_tokens": 80,
+      "output_tokens": 9,
+      "requested_tools": [
+        {
+          "id": "search",
+          "name": "search",
+          "arguments": {
+            "query": "dismissal"
+          }
+        }
+      ],
+      "started": 1760000000.0,
+      "ended": 1760000001.0,
+      "duration_ms": 1000.0
+    }
+  ]
+}