mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-27 21:05:34 +00:00
hybrid: capture tool calls + tool results + structured content in trace events
Previous trace events only stored the joined `text` blocks from Anthropic
responses and `message.content` from OpenAI/vLLM — every tool_use block,
server_tool_use block, web_search_tool_result, and OpenAI-style tool_calls
list was being thrown away.
Now each call event includes:
Anthropic events:
- `content_blocks`: full list of every block the model emitted, with all
fields preserved (text / tool_use / server_tool_use /
web_search_tool_result / thinking). Nested content (e.g. search-result
citations) recurses.
- `tool_calls`: filtered subset (tool_use + server_tool_use).
- `tool_results`: filtered subset (web_search_tool_result + tool_result).
- `tools_declared`, `tool_choice`, `output_config`, `stop_reason`.
OpenAI / vLLM events:
- `tool_calls`: serialized `message.tool_calls` (id, function name, args).
- `reasoning_content`: thinking-model intermediate reasoning if present.
- `tools_declared`, `tool_choice`, `finish_reason`.
Confirmed against a real Anthropic call with `tools=[web_search]`,
`tool_choice={"type":"any"}`: the resulting event contains the
`server_tool_use` block with the actual search query, the
`web_search_tool_result` with 10 nested citation blocks, and all
subsequent text blocks the model emitted to synthesize the answer.
This commit is contained in:
@@ -95,6 +95,64 @@ def _close_trace() -> None:
|
||||
delattr(_TRACE_STATE, "events")
|
||||
|
||||
|
||||
def _serialize_block(block: Any) -> Dict[str, Any]:
|
||||
"""Turn an Anthropic content block (text / tool_use / server_tool_use /
|
||||
web_search_tool_result / thinking) into a JSON-safe dict.
|
||||
|
||||
Each block type carries different fields; we extract everything we can
|
||||
so the per-task log file is a complete record of what the model emitted
|
||||
(including every tool call request and tool result body).
|
||||
"""
|
||||
out: Dict[str, Any] = {"type": getattr(block, "type", type(block).__name__)}
|
||||
for attr in (
|
||||
"id", "name", "input", "text", "thinking", "signature",
|
||||
"tool_use_id", "content",
|
||||
):
|
||||
if hasattr(block, attr):
|
||||
val = getattr(block, attr)
|
||||
# Nested content (e.g. web_search_tool_result.content is a list of
|
||||
# citation/result blocks). Recurse for completeness.
|
||||
if attr == "content" and isinstance(val, list):
|
||||
out[attr] = [_serialize_block(b) for b in val]
|
||||
else:
|
||||
out[attr] = _jsonable(val)
|
||||
return out
|
||||
|
||||
|
||||
def _serialize_openai_tool_calls(tool_calls: Any) -> List[Dict[str, Any]]:
|
||||
"""vLLM / OpenAI returns ChatCompletionMessageToolCall objects. Pull out
|
||||
id, type, name, and the (JSON-string) arguments so they're round-trippable."""
|
||||
out: List[Dict[str, Any]] = []
|
||||
if not tool_calls:
|
||||
return out
|
||||
for tc in tool_calls:
|
||||
fn = getattr(tc, "function", None)
|
||||
out.append({
|
||||
"id": getattr(tc, "id", None),
|
||||
"type": getattr(tc, "type", "function"),
|
||||
"function": {
|
||||
"name": getattr(fn, "name", None) if fn else None,
|
||||
"arguments": getattr(fn, "arguments", None) if fn else None,
|
||||
},
|
||||
})
|
||||
return out
|
||||
|
||||
|
||||
def _jsonable(v: Any) -> Any:
|
||||
"""Best-effort JSON-friendly conversion. Pydantic models → .model_dump(),
|
||||
dataclasses untouched (json.dumps handles them via default=str)."""
|
||||
if hasattr(v, "model_dump"):
|
||||
try:
|
||||
return v.model_dump()
|
||||
except Exception:
|
||||
pass
|
||||
if isinstance(v, (list, tuple)):
|
||||
return [_jsonable(x) for x in v]
|
||||
if isinstance(v, dict):
|
||||
return {k: _jsonable(x) for k, x in v.items()}
|
||||
return v
|
||||
|
||||
|
||||
class LocalCloudAgent(BaseAgent):
|
||||
"""Base for paradigm agents that coordinate a local + cloud model pair.
|
||||
|
||||
@@ -196,6 +254,13 @@ class LocalCloudAgent(BaseAgent):
|
||||
text = "".join(b.text for b in msg.content if hasattr(b, "text"))
|
||||
srv = getattr(msg.usage, "server_tool_use", None)
|
||||
n_searches = getattr(srv, "web_search_requests", 0) if srv else 0
|
||||
content_blocks = [_serialize_block(b) for b in msg.content]
|
||||
tool_use_blocks = [b for b in content_blocks if b.get("type") in (
|
||||
"tool_use", "server_tool_use",
|
||||
)]
|
||||
tool_result_blocks = [b for b in content_blocks if b.get("type") in (
|
||||
"web_search_tool_result", "tool_result",
|
||||
)]
|
||||
_record_event({
|
||||
"kind": "anthropic",
|
||||
"role": trace_role,
|
||||
@@ -203,12 +268,16 @@ class LocalCloudAgent(BaseAgent):
|
||||
"system": system,
|
||||
"user": user,
|
||||
"response": text,
|
||||
"content_blocks": content_blocks,
|
||||
"tool_calls": tool_use_blocks,
|
||||
"tool_results": tool_result_blocks,
|
||||
"tokens_in": msg.usage.input_tokens,
|
||||
"tokens_out": msg.usage.output_tokens,
|
||||
"n_web_searches": n_searches,
|
||||
"tools": [t.get("name") for t in (tools or [])],
|
||||
"tools_declared": tools,
|
||||
"tool_choice": tool_choice,
|
||||
"output_config": output_config,
|
||||
"stop_reason": getattr(msg, "stop_reason", None),
|
||||
"latency_s": latency,
|
||||
"ts": time.time(),
|
||||
})
|
||||
@@ -223,10 +292,13 @@ class LocalCloudAgent(BaseAgent):
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.0,
|
||||
response_format: Optional[dict] = None,
|
||||
tools: Optional[list] = None,
|
||||
tool_choice: Optional[Any] = None,
|
||||
timeout: float = 600.0,
|
||||
trace_role: str = "cloud",
|
||||
) -> Tuple[str, int, int]:
|
||||
"""Single OpenAI call. Returns (text, p_tok, c_tok). Trace-captured."""
|
||||
"""Single OpenAI call. Returns (text, p_tok, c_tok). Trace-captured;
|
||||
also records any tool_calls the model emits."""
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(timeout=timeout)
|
||||
@@ -242,10 +314,20 @@ class LocalCloudAgent(BaseAgent):
|
||||
kwargs["temperature"] = temperature
|
||||
if response_format is not None:
|
||||
kwargs["response_format"] = response_format
|
||||
if tools:
|
||||
kwargs["tools"] = tools
|
||||
if tool_choice is not None:
|
||||
kwargs["tool_choice"] = tool_choice
|
||||
t0 = time.time()
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
latency = time.time() - t0
|
||||
text = resp.choices[0].message.content or ""
|
||||
choice = resp.choices[0]
|
||||
message = choice.message
|
||||
text = message.content or ""
|
||||
tool_calls = _serialize_openai_tool_calls(getattr(message, "tool_calls", None))
|
||||
reasoning = getattr(message, "reasoning_content", None) or getattr(
|
||||
message, "reasoning", None
|
||||
)
|
||||
u = resp.usage
|
||||
p = getattr(u, "prompt_tokens", 0) if u else 0
|
||||
c = getattr(u, "completion_tokens", 0) if u else 0
|
||||
@@ -256,9 +338,14 @@ class LocalCloudAgent(BaseAgent):
|
||||
"system": system,
|
||||
"user": user,
|
||||
"response": text,
|
||||
"tool_calls": tool_calls,
|
||||
"reasoning_content": reasoning,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"response_format": response_format,
|
||||
"tools_declared": tools,
|
||||
"tool_choice": tool_choice,
|
||||
"finish_reason": getattr(choice, "finish_reason", None),
|
||||
"latency_s": latency,
|
||||
"ts": time.time(),
|
||||
})
|
||||
@@ -274,11 +361,16 @@ class LocalCloudAgent(BaseAgent):
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.0,
|
||||
enable_thinking: bool = False,
|
||||
tools: Optional[list] = None,
|
||||
tool_choice: Optional[Any] = None,
|
||||
timeout: float = 600.0,
|
||||
trace_role: str = "local",
|
||||
) -> Tuple[str, int, int]:
|
||||
"""Local vLLM (OpenAI-compatible) call. Returns (text, p_tok, c_tok).
|
||||
Trace-captured."""
|
||||
Captures the full response into the trace — including any tool_calls
|
||||
the local model emits (vLLM exposes them in
|
||||
``resp.choices[0].message.tool_calls`` when ``--enable-auto-tool-choice``
|
||||
is on)."""
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=timeout)
|
||||
@@ -286,16 +378,27 @@ class LocalCloudAgent(BaseAgent):
|
||||
if system:
|
||||
messages.append({"role": "system", "content": system})
|
||||
messages.append({"role": "user", "content": user})
|
||||
t0 = time.time()
|
||||
resp = client.chat.completions.create(
|
||||
kwargs: Dict[str, Any] = dict(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": enable_thinking}},
|
||||
)
|
||||
if tools:
|
||||
kwargs["tools"] = tools
|
||||
if tool_choice is not None:
|
||||
kwargs["tool_choice"] = tool_choice
|
||||
t0 = time.time()
|
||||
resp = client.chat.completions.create(**kwargs)
|
||||
latency = time.time() - t0
|
||||
text = resp.choices[0].message.content or ""
|
||||
choice = resp.choices[0]
|
||||
message = choice.message
|
||||
text = message.content or ""
|
||||
tool_calls = _serialize_openai_tool_calls(getattr(message, "tool_calls", None))
|
||||
reasoning = getattr(message, "reasoning_content", None) or getattr(
|
||||
message, "reasoning", None
|
||||
)
|
||||
u = resp.usage
|
||||
p = getattr(u, "prompt_tokens", 0) if u else 0
|
||||
c = getattr(u, "completion_tokens", 0) if u else 0
|
||||
@@ -307,11 +410,16 @@ class LocalCloudAgent(BaseAgent):
|
||||
"system": system,
|
||||
"user": user,
|
||||
"response": text,
|
||||
"tool_calls": tool_calls,
|
||||
"reasoning_content": reasoning,
|
||||
"tokens_in": p,
|
||||
"tokens_out": c,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"enable_thinking": enable_thinking,
|
||||
"tools_declared": tools,
|
||||
"tool_choice": tool_choice,
|
||||
"finish_reason": getattr(choice, "finish_reason", None),
|
||||
"latency_s": latency,
|
||||
"ts": time.time(),
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user