Skip to content

Commit 57b64bf

Browse files
committed
fix(agent): adapt image tool outputs for Chat Completions
Fixes #242.
1 parent ff54396 commit 57b64bf

7 files changed

Lines changed: 330 additions & 26 deletions

File tree

‎openkb/agent/chat.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from prompt_toolkit.styles import Style
2525

2626
from openkb.agent.chat_session import ChatSession
27+
from openkb.agent.model_compat import with_chat_completions_compat
2728
from openkb.agent.query import (
2829
MAX_TURNS,
2930
build_chat_agent,
@@ -354,7 +355,12 @@ async def _run_turn(
354355

355356
new_input = session.history + [{"role": "user", "content": user_input}]
356357

357-
result = Runner.run_streamed(agent, new_input, max_turns=MAX_TURNS)
358+
result = Runner.run_streamed(
359+
agent,
360+
new_input,
361+
max_turns=MAX_TURNS,
362+
run_config=with_chat_completions_compat(None),
363+
)
358364

359365
print()
360366
collected: list[str] = []

‎openkb/agent/model_compat.py‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
"""Model-protocol compatibility helpers for agent runs."""
2+
3+
from __future__ import annotations
4+
5+
import inspect
6+
from dataclasses import replace
7+
from typing import Any
8+
9+
from agents import RunConfig
10+
from agents.run_config import CallModelData, ModelInputData
11+
12+
_IMAGE_OUTPUT_ACK = "Image returned successfully; inspect the following user image."
13+
_IMAGE_INPUT_NOTE = "Images returned by the preceding tool call(s):"
14+
15+
16+
def adapt_image_tool_outputs_for_chat_completions(
17+
data: CallModelData[Any],
18+
) -> ModelInputData:
19+
"""Move image tool outputs into a following user message.
20+
21+
Chat Completions accepts images in user content but restricts tool content
22+
to text. The Agents SDK uses Responses-style ``input_image`` parts for
23+
``ToolOutputImage``, so rewrite those parts immediately before the model
24+
call while preserving a text result for every tool call.
25+
"""
26+
adapted: list[Any] = []
27+
pending_images: list[dict[str, Any]] = []
28+
29+
def flush_images() -> None:
30+
if not pending_images:
31+
return
32+
adapted.append(
33+
{
34+
"role": "user",
35+
"content": [
36+
{"type": "input_text", "text": _IMAGE_INPUT_NOTE},
37+
*pending_images,
38+
],
39+
}
40+
)
41+
pending_images.clear()
42+
43+
for item in data.model_data.input:
44+
if not (isinstance(item, dict) and item.get("type") == "function_call_output"):
45+
flush_images()
46+
adapted.append(item)
47+
continue
48+
49+
output = item.get("output")
50+
if not isinstance(output, list):
51+
adapted.append(item)
52+
continue
53+
54+
images = [
55+
part for part in output if isinstance(part, dict) and part.get("type") == "input_image"
56+
]
57+
if not images:
58+
adapted.append(item)
59+
continue
60+
61+
text_parts = [
62+
part for part in output if isinstance(part, dict) and part.get("type") == "input_text"
63+
]
64+
adapted.append({**item, "output": text_parts or _IMAGE_OUTPUT_ACK})
65+
pending_images.extend(images)
66+
67+
flush_images()
68+
return ModelInputData(input=adapted, instructions=data.model_data.instructions)
69+
70+
71+
CHAT_COMPLETIONS_RUN_CONFIG = RunConfig(
72+
call_model_input_filter=adapt_image_tool_outputs_for_chat_completions
73+
)
74+
75+
76+
def with_chat_completions_compat(run_config: RunConfig | None) -> RunConfig:
77+
"""Return a run config that preserves existing settings and adapts image outputs."""
78+
if run_config is None:
79+
return CHAT_COMPLETIONS_RUN_CONFIG
80+
81+
existing_filter = run_config.call_model_input_filter
82+
if existing_filter is None:
83+
return replace(
84+
run_config,
85+
call_model_input_filter=adapt_image_tool_outputs_for_chat_completions,
86+
)
87+
if existing_filter is adapt_image_tool_outputs_for_chat_completions or getattr(
88+
existing_filter, "_openkb_chat_completions_compat", False
89+
):
90+
return run_config
91+
92+
async def composed_filter(data: CallModelData[Any]) -> ModelInputData:
93+
filtered = existing_filter(data)
94+
if inspect.isawaitable(filtered):
95+
filtered = await filtered
96+
return adapt_image_tool_outputs_for_chat_completions(
97+
CallModelData(
98+
model_data=filtered,
99+
agent=data.agent,
100+
context=data.context,
101+
)
102+
)
103+
104+
setattr(composed_filter, "_openkb_chat_completions_compat", True)
105+
return replace(run_config, call_model_input_filter=composed_filter)

‎openkb/agent/query.py‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
from agents import Agent, Runner, ToolOutputImage, ToolOutputText, function_tool
99

10+
from openkb.agent.model_compat import with_chat_completions_compat
1011
from openkb.agent.tools import (
1112
artifact_event_from_write,
1213
get_wiki_page_content,
@@ -159,10 +160,11 @@ async def iter_agent_response_events(
159160
from agents import RawResponsesStreamEvent, RunItemStreamEvent
160161
from openai.types.responses import ResponseTextDeltaEvent
161162

162-
result = (
163-
Runner.run_streamed(agent, input_data, max_turns=max_turns, run_config=run_config)
164-
if run_config
165-
else Runner.run_streamed(agent, input_data, max_turns=max_turns)
163+
result = Runner.run_streamed(
164+
agent,
165+
input_data,
166+
max_turns=max_turns,
167+
run_config=with_chat_completions_compat(run_config),
166168
)
167169
collected: list[str] = []
168170
pending_calls: dict[str, tuple[str, str]] = {}
@@ -387,10 +389,11 @@ async def run_query(
387389
agent = build_query_agent(wiki_root, model, language=language, bundle=bundle)
388390

389391
if not stream:
390-
result = (
391-
await Runner.run(agent, question, max_turns=MAX_TURNS, run_config=run_config)
392-
if run_config
393-
else await Runner.run(agent, question, max_turns=MAX_TURNS)
392+
result = await Runner.run(
393+
agent,
394+
question,
395+
max_turns=MAX_TURNS,
396+
run_config=with_chat_completions_compat(run_config),
394397
)
395398
return result.final_output or ""
396399

@@ -425,10 +428,11 @@ def _start_live() -> Live | None:
425428
live: Live | None = None
426429
last_was_text = False
427430
need_blank_before_text = False
428-
result = (
429-
Runner.run_streamed(agent, question, max_turns=MAX_TURNS, run_config=run_config)
430-
if run_config
431-
else Runner.run_streamed(agent, question, max_turns=MAX_TURNS)
431+
result = Runner.run_streamed(
432+
agent,
433+
question,
434+
max_turns=MAX_TURNS,
435+
run_config=with_chat_completions_compat(run_config),
432436
)
433437
collected: list[str] = []
434438
segment: list[str] = []

‎openkb/agent/skill_runner.py‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232

3333
from agents import Runner, function_tool
3434

35+
from openkb.agent.model_compat import with_chat_completions_compat
3536
from openkb.agent.query import build_query_agent, build_run_config_from_bundle
3637
from openkb.agent.skills import _parse_frontmatter, scan_local_skills
3738
from openkb.agent.tools import read_kb_file, write_kb_file
@@ -190,15 +191,17 @@ def read_output_or_skill_file(path: str) -> str:
190191

191192
# Per-KB credential isolation: when a bundle is supplied (REST path) the
192193
# RunConfig carries a dedicated LitellmModel with this KB's api_key/base_url,
193-
# overriding the process-global provider for this run only. When bundle is
194-
# None (CLI path) build_run_config_from_bundle returns None and the call is
195-
# byte-identical to the pre-bundle behavior (no run_config kwarg passed).
194+
# overriding the process-global provider for this run only. The compatibility
195+
# helper preserves that config and adapts image tool outputs for Chat
196+
# Completions; without a bundle it creates the minimal compatible config.
196197
run_config = build_run_config_from_bundle(model, bundle)
197198
try:
198-
if run_config:
199-
await Runner.run(agent, user_seed, max_turns=max_turns, run_config=run_config)
200-
else:
201-
await Runner.run(agent, user_seed, max_turns=max_turns)
199+
await Runner.run(
200+
agent,
201+
user_seed,
202+
max_turns=max_turns,
203+
run_config=with_chat_completions_compat(run_config),
204+
)
202205
except MaxTurnsExceeded as exc:
203206
raise RuntimeError(
204207
f"Skill {skill_name!r} hit the {max_turns}-step cap before "

‎openkb/skill/creator.py‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from agents import Agent, Runner, ToolOutputImage, ToolOutputText, function_tool
2121
from agents.model_settings import ModelSettings
2222

23+
from openkb.agent.model_compat import with_chat_completions_compat
2324
from openkb.config import LlmCredentialBundle, resolve_model_settings
2425
from openkb.prompts import load_prompt
2526
from openkb.schema import get_agents_md
@@ -228,17 +229,19 @@ async def run_skill_create(
228229
# Per-KB credential isolation: when a bundle is supplied (REST path) the
229230
# RunConfig carries a dedicated LitellmModel with this KB's api_key/base_url,
230231
# overriding the process-global provider for this run only. When bundle is
231-
# None (CLI path) build_run_config_from_bundle returns None and the call is
232-
# byte-identical to the pre-bundle behavior (no run_config kwarg passed).
232+
# None (CLI path) build_run_config_from_bundle returns None, so the shared
233+
# Chat Completions image compatibility config is used instead.
233234
# Lazy import mirrors this module's convention of keeping query imports local.
234235
from openkb.agent.query import build_run_config_from_bundle
235236

236237
run_config = build_run_config_from_bundle(model, bundle)
237238
try:
238-
if run_config:
239-
await Runner.run(agent, seed, max_turns=MAX_TURNS, run_config=run_config)
240-
else:
241-
await Runner.run(agent, seed, max_turns=MAX_TURNS)
239+
await Runner.run(
240+
agent,
241+
seed,
242+
max_turns=MAX_TURNS,
243+
run_config=with_chat_completions_compat(run_config),
244+
)
242245
except MaxTurnsExceeded as exc:
243246
raise RuntimeError(
244247
f"Skill compilation hit the {MAX_TURNS}-step cap before finishing. "

0 commit comments

Comments
 (0)