change the code according to the wt-data sdk change - #47
Merged
Conversation
📝 WalkthroughWalkthroughCloud telemetry handling now supports provider-specific response formats, JSON-deserialized values, arbitrary message fields, and JSON-string landing records. Gateway aggregation preserves additional response data, while storage and retrieval normalize endpoint-specific payloads. ChangesTelemetry normalization
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Provider
participant GatewayApp
participant GatewayStorage
participant CloudStrategy
Provider->>GatewayApp: Stream response events and message deltas
GatewayApp->>GatewayApp: Merge output and arbitrary message fields
GatewayApp->>GatewayStorage: Persist provider-specific telemetry
GatewayStorage->>GatewayStorage: Extract endpoint inputs and outputs
GatewayStorage->>CloudStrategy: Store normalized JSON values
CloudStrategy->>GatewayStorage: Request deserialized cloud records
GatewayStorage-->>CloudStrategy: Return normalized messages, responses, and metadata
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
gateway/app.py (1)
1469-1485: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueConcatenating extra fields treats snapshots as deltas.
_merge_chat_completion_choicecalls_merge_message_payloadfor bothchoice["message"]andchoice["delta"]. For amessagepayload, the value is a full snapshot, not a fragment. Line 1481 concatenates string values, so a repeated or snapshot-style extra field accumulates duplicated text inextra_message_fields.Pass the payload kind into
_merge_message_payloadand concatenate only for delta payloads.♻️ Proposed fix to separate snapshot and delta semantics
-def _merge_message_payload(state: dict[str, Any], payload: dict[str, Any]) -> None: +def _merge_message_payload(state: dict[str, Any], payload: dict[str, Any], *, is_delta: bool = True) -> None: @@ extra_fields = state.setdefault("extra_message_fields", {}) for key, value in payload.items(): if key in known_fields: continue - if isinstance(value, str) and isinstance(extra_fields.get(key), str): + if is_delta and isinstance(value, str) and isinstance(extra_fields.get(key), str): extra_fields[key] += value else: extra_fields[key] = valueUpdate the call sites in
_merge_chat_completion_choice:message = choice.get("message") if isinstance(message, dict): - _merge_message_payload(state, message) + _merge_message_payload(state, message, is_delta=False) delta = choice.get("delta") if isinstance(delta, dict): - _merge_message_payload(state, delta) + _merge_message_payload(state, delta, is_delta=True)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gateway/app.py` around lines 1469 - 1485, Update _merge_message_payload to accept the payload kind and only concatenate string extra fields when processing delta payloads; for message snapshots, replace the existing value. Update both _merge_message_payload call sites in _merge_chat_completion_choice to pass the appropriate message or delta kind, preserving existing handling for known fields.gateway/storage.py (1)
659-669: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd a fallback for unparsable
responsesrequests.
_openai_request_messagesfalls back torecord.messageswhen the request payload does not parse._responses_request_inputreturns[]instead. Aresponsesrecord with a non-JSON or non-dict request then stores an empty message list, and the trajectory loses the prompt.Accept a
fallbackargument and mirror thechat/completionsbehavior.♻️ Proposed fix to add a fallback
-def _responses_request_input(request_payload: Any) -> Any: +def _responses_request_input( + request_payload: Any, + *, + fallback: list[dict[str, Any]], +) -> Any: if not isinstance(request_payload, dict): - return [] + return [dict(message) for message in fallback] input_value = request_payload.get("input") if isinstance(input_value, list): return input_value if isinstance(input_value, dict): return [input_value] if input_value is None: - return [] + return [dict(message) for message in fallback] return [{"role": "user", "content": input_value}]Update the call site at line 393:
- stored_messages = _responses_request_input(request_payload) + stored_messages = _responses_request_input( + request_payload, + fallback=record.messages, + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@gateway/storage.py` around lines 659 - 669, Update _responses_request_input to accept a fallback argument and return it when request_payload is not a dictionary or input cannot be parsed, matching _openai_request_messages behavior. Update its call site in the responses request handling flow to pass record.messages as the fallback, while preserving existing list, dictionary, and scalar input handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@core/data_manager/strategy/cloud_strategy_impl.py`:
- Around line 100-107: The SDK schema rejection must remain sticky across
repeated loads. Update _load_wt_sdk() to store the compatibility exception
before fallback handling and re-raise it on subsequent calls, without clearing
only LandingRecord; restrict _install_mock_wt_sdk_fallbacks() to explicit test
setup, and add a regression test that calls _load_wt_sdk() twice after schema
rejection and verifies both calls fail with the compatibility error.
- Around line 262-268: Streamed responses with delta.content are not being
extracted, leaving protocol envelopes in trajectory data. At
core/data_manager/strategy/cloud_strategy_impl.py#L262-L268, update the list
comprehension to fall back to extracting choice["delta"]["content"] when
choice.get("message") returns None or empty text, preserving the existing
extract_text flow for messages that have text. At
evaluator/trajectory_reader.py#L141-L166, apply _extract_choice_text(item) to
each dictionary item before checking the standard fields output_text, text, and
content, so that delta-based completions are normalized to the expected format
before field inspection.
In `@gateway/storage.py`:
- Around line 672-710: Update _chat_completion_output and _responses_output to
return the parsed payload as the fallback whenever the expected success shape is
absent, including invalid or non-JSON response bodies according to the existing
_json_loads behavior. Preserve the current extraction of valid messages and
output values, but replace None fallbacks so response storage retains upstream
error, incomplete-stream, and non-conforming bodies.
- Around line 364-409: The step construction around provider_meta must ensure
provider_meta is never included for non-cloud storage. Restrict the
provider_meta assignment to the cloud storage path, while preserving the
existing Anthropic metadata behavior for cloud records and keeping non-cloud
record_steps_batch compatible with its step schema.
---
Nitpick comments:
In `@gateway/app.py`:
- Around line 1469-1485: Update _merge_message_payload to accept the payload
kind and only concatenate string extra fields when processing delta payloads;
for message snapshots, replace the existing value. Update both
_merge_message_payload call sites in _merge_chat_completion_choice to pass the
appropriate message or delta kind, preserving existing handling for known
fields.
In `@gateway/storage.py`:
- Around line 659-669: Update _responses_request_input to accept a fallback
argument and return it when request_payload is not a dictionary or input cannot
be parsed, matching _openai_request_messages behavior. Update its call site in
the responses request handling flow to pass record.messages as the fallback,
while preserving existing list, dictionary, and scalar input handling.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 5e621d15-1713-412d-ae9e-ab9ea9fb7a79
📒 Files selected for processing (4)
core/data_manager/strategy/cloud_strategy_impl.pyevaluator/trajectory_reader.pygateway/app.pygateway/storage.py
Comment on lines
+100
to
+107
| LandingRecord = None | ||
| raise RuntimeError( | ||
| "Installed wt_sdk uses the legacy LandingRecord schema. " | ||
| "Cloud storage requires messages/response/meta_json JSON strings; " | ||
| "reinstall requirements-cloud.txt with --upgrade --force-reinstall. " | ||
| f"Loaded wt_sdk version={getattr(wt_sdk, '__version__', 'unknown')} " | ||
| f"from {getattr(wt_sdk, '__file__', 'unknown')}" | ||
| ) from exc |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Keep the SDK-schema failure sticky.
Line 100 clears only LandingRecord. On the next _load_wt_sdk() call, Lines 59-73 install _Model as LandingRecord because the real client symbols remain set. The call then bypasses the schema check and can send a mock record to the real client.
Store and re-raise the compatibility error before the fallback path. Restrict _install_mock_wt_sdk_fallbacks() to explicit test setup. Add a regression test that invokes _load_wt_sdk() twice after schema rejection.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@core/data_manager/strategy/cloud_strategy_impl.py` around lines 100 - 107,
The SDK schema rejection must remain sticky across repeated loads. Update
_load_wt_sdk() to store the compatibility exception before fallback handling and
re-raise it on subsequent calls, without clearing only LandingRecord; restrict
_install_mock_wt_sdk_fallbacks() to explicit test setup, and add a regression
test that calls _load_wt_sdk() twice after schema rejection and verifies both
calls fail with the compatibility error.
Comment on lines
+262
to
+268
| choices = payload.get("choices") | ||
| if isinstance(choices, list): | ||
| return "".join( | ||
| extract_text(choice.get("message")) | ||
| for choice in choices | ||
| if isinstance(choice, dict) | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Extract choices[].delta.content from streamed response lists.
A payload such as [{"choices": [{"delta": {"content": "text"}}]}] reaches both paths. Both paths return serialized provider JSON instead of completion text. This places protocol envelopes in trajectory data.
core/data_manager/strategy/cloud_strategy_impl.py#L262-L268: Extractchoice["delta"]whenchoice["message"]has no text.evaluator/trajectory_reader.py#L141-L166: Call_extract_choice_text(item)for each dictionary item before checkingoutput_text,text, andcontent.
Proposed fix
# core/data_manager/strategy/cloud_strategy_impl.py
- extract_text(choice.get("message"))
+ extract_text(choice.get("message") or choice.get("delta"))# evaluator/trajectory_reader.py
if not isinstance(item, dict):
continue
+ choice_text = _extract_choice_text(item)
+ if choice_text:
+ text_chunks.append(choice_text)
+ continue
for key in ("output_text", "text", "content"):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| choices = payload.get("choices") | |
| if isinstance(choices, list): | |
| return "".join( | |
| extract_text(choice.get("message")) | |
| for choice in choices | |
| if isinstance(choice, dict) | |
| ) | |
| choices = payload.get("choices") | |
| if isinstance(choices, list): | |
| return "".join( | |
| extract_text(choice.get("message") or choice.get("delta")) | |
| for choice in choices | |
| if isinstance(choice, dict) | |
| ) |
Suggested change
| choices = payload.get("choices") | |
| if isinstance(choices, list): | |
| return "".join( | |
| extract_text(choice.get("message")) | |
| for choice in choices | |
| if isinstance(choice, dict) | |
| ) | |
| if isinstance(parsed, list): | |
| text_chunks: list[str] = [] | |
| for item in parsed: | |
| if isinstance(item, str): | |
| text_chunks.append(item) | |
| continue | |
| if not isinstance(item, dict): | |
| continue | |
| choice_text = _extract_choice_text(item) | |
| if choice_text: | |
| text_chunks.append(choice_text) | |
| continue | |
| for key in ("output_text", "text", "content"): | |
| content = item.get(key) | |
| if isinstance(content, str): | |
| text_chunks.append(content) | |
| break | |
| if isinstance(content, list): | |
| chunks = [ | |
| str(part.get("text") or part.get("content") or "") | |
| for part in content | |
| if isinstance(part, dict) | |
| ] | |
| combined = "".join(chunks) | |
| if combined: | |
| text_chunks.append(combined) | |
| break | |
| if text_chunks: | |
| return "".join(text_chunks) | |
| return json.dumps(parsed, ensure_ascii=False, default=str) |
📍 Affects 2 files
core/data_manager/strategy/cloud_strategy_impl.py#L262-L268(this comment)evaluator/trajectory_reader.py#L141-L166
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@core/data_manager/strategy/cloud_strategy_impl.py` around lines 262 - 268,
Streamed responses with delta.content are not being extracted, leaving protocol
envelopes in trajectory data. At
core/data_manager/strategy/cloud_strategy_impl.py#L262-L268, update the list
comprehension to fall back to extracting choice["delta"]["content"] when
choice.get("message") returns None or empty text, preserving the existing
extract_text flow for messages that have text. At
evaluator/trajectory_reader.py#L141-L166, apply _extract_choice_text(item) to
each dictionary item before checking the standard fields output_text, text, and
content, so that delta-based completions are normalized to the expected format
before field inspection.
Comment on lines
+364
to
+409
| provider_meta: dict[str, Any] | None = None | ||
| if self.cfg.storage_type == "cloud": | ||
| response_body = _json_loads(record.response) | ||
| try: | ||
| content = response_body["choices"][0]["message"]["content"] | ||
| except (TypeError, KeyError, IndexError): | ||
| pass | ||
| else: | ||
| if isinstance(content, str) or content is None: | ||
| stored_response = content or "" | ||
| request_payload = _json_loads(record.request) | ||
| if record.endpoint == "messages": | ||
| # Anthropic payloads remain provider-native and live only | ||
| # in meta_json. Do not fabricate OpenAI messages from them. | ||
| response_payload = _json_loads(record.response) | ||
| stored_messages = None | ||
| stored_response = None | ||
| provider_meta = { | ||
| "provider": "anthropic", | ||
| "request": ( | ||
| request_payload | ||
| if request_payload is not None | ||
| else record.request | ||
| ), | ||
| "response": ( | ||
| response_payload | ||
| if response_payload is not None | ||
| else record.response | ||
| ), | ||
| } | ||
| elif record.endpoint == "chat/completions": | ||
| stored_messages = _openai_request_messages( | ||
| request_payload, | ||
| fallback=record.messages, | ||
| ) | ||
| stored_response = _chat_completion_output(record.response) | ||
| elif record.endpoint == "responses": | ||
| stored_messages = _responses_request_input(request_payload) | ||
| stored_response = _responses_output(record.response) | ||
|
|
||
| step = { | ||
| "session": session, | ||
| "step_id": record.seq_id, | ||
| "messages": _trajectory_messages(record), | ||
| "request": record.request, | ||
| "messages": stored_messages, | ||
| "request": None if provider_meta is not None else record.request, | ||
| "response": stored_response, | ||
| "step_reward": 0.0, | ||
| "env_state": json.dumps(self._metadata(record), ensure_ascii=False, default=str), | ||
| "terminated": False, | ||
| "truncated": record.is_truncated, | ||
| "is_trainable": False, | ||
| } | ||
| if provider_meta is not None: | ||
| step["provider_meta"] = provider_meta |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Verify the provider_meta contract across strategies and readers.
set -uo pipefail
# Locate record_steps_batch definitions and their step-key handling.
fd -e py | xargs rg -n -C 20 'def record_steps_batch' || true
# Check whether provider_meta is recognized anywhere outside gateway/storage.py.
rg -n -C 5 'provider_meta' --glob '*.py' || true
# Check whether the trajectory reader has an Anthropic / provider_meta path.
fd -g 'trajectory_reader.py' --exec rg -n -C 5 'provider_meta|anthropic|meta_json'Repository: AI45Lab/SAfactory
Length of output: 10521
🏁 Script executed:
#!/bin/bash
set -uo pipefail
echo "## candidate file line counts"
wc -l core/data_manager/strategy/cloud_strategy_impl.py core/data_manager/strategy/base_strategy.py evaluator/trajectory_reader.py gateway/storage.py 2>/dev/null || true
echo
echo "## cloud _build_step_record signature and body"
sed -n '520,595p' core/data_manager/strategy/cloud_strategy_impl.py
echo
echo "## base record_step signature/body relevant sections"
sed -n '110,165p' core/data_manager/strategy/base_strategy.py
echo
echo "## cloud strategy bulk path"
sed -n '640,685p' core/data_manager/strategy/cloud_strategy_impl.py
echo
echo "## provider_meta occurrences"
rg -n -C 8 'provider_meta' --glob '*.py' || true
echo
echo "## trajectory_reader key fields"
rg -n -C 5 'def parse_gateway_row|row\.get\(|_extract_response_text|messages|env_state|record' evaluator/trajectory_reader.py || trueRepository: AI45Lab/SAfactory
Length of output: 7056
🏁 Script executed:
#!/bin/bash
set -uo pipefail
echo "## provider_meta occurrences in all tracked files"
python3 - <<'PY'
import subprocess, sys
try:
files = subprocess.check_output(['git', 'ls-files'], text=True).splitlines()
except Exception as exc:
print('git ls-files failed:', exc)
files = subprocess.check_output(['fd', '-e', 'py', '.'], text=True).splitlines()
hits=[]
for path in files:
try:
with open(path, 'r', encoding='utf-8', errors='replace') as f:
txt = f.read()
except Exception:
continue
if 'provider_meta' in txt:
hits.append(path)
for path in sorted(hits):
print(path)
with open(path, 'r', encoding='utf-8', errors='replace') as f:
lines = f.read().splitlines()
for i, line in enumerate(lines, 1):
if 'provider_meta' in line:
lo=max(1,i-4); hi=min(len(lines), i+7)
for j in range(lo, hi+1):
print(f'{j}: {lines[j-1]}')
print('---')
if not hits:
print('No provider_meta matches found in tracked files.')
import io
sub = subprocess.run(['git', 'diff', '--',], input=b'provider_meta', capture_output=True)
print('git diff context exit: %s' % sub.returncode)
for cmd in [['git','diff','--','gateway/storage.py'], ['git','diff','--','core/data_manager/strategy/'], ['git','diff','--','evaluator/trajectory_reader.py']]:
sub=subprocess.run(cmd, text=True, capture_output=True)
if sub.stdout.strip():
print('---', ' '.join(cmd), '---')
lines=sub.stdout.splitlines()
for n,line in enumerate(lines[:220], 1):
if n==220:
print('(truncated)')
break
print(line)
PY
echo
echo "## cloud _build_step_record full function"
sed -n '555,635p' core/data_manager/strategy/cloud_strategy_impl.py
echo
echo "## cloud record provider signature/body"
rg -n -C 25 '_record_record|record_provider|ingest_landing|provider_meta' core/data_manager/strategy/cloud_strategy_impl.py
echo
echo "## evaluate readers provider/meta related keys"
sed -n '1,289p' evaluator/trajectory_reader.py | nl -ba | sed -n '1,289p' | rg -C 8 'provider_meta|meta_json|messages|response|env_state|parse_gateway_row|read' || trueRepository: AI45Lab/SAfactory
Length of output: 17643
🏁 Script executed:
#!/bin/bash
set -uo pipefail
echo "## sqlite strategy record step methods"
rg -n -C 60 'class .*Strategy|def record_step|def record_steps_batch|meta_json|provider_meta' core/data_manager/strategy --glob '*.py' || true
echo
echo "## base record_step kwargs usage"
sed -n '154,165p' core/data_manager/strategy/base_strategy.py
echo
echo "## eval reader relevant sections"
python3 - <<'PY'
from pathlib import Path
p = Path('evaluator/trajectory_reader.py')
txt = p.read_text()
lines = txt.splitlines()
for name in ['parse_gateway_row', '_extract_response_text', '_messages', 'provider_meta', 'meta_json']:
print(f'--- {name} ---')
for i,l in enumerate(lines,1):
if name in l:
print(f'{i}: {l}')
PYRepository: AI45Lab/SAfactory
Length of output: 50374
Do not pass provider_meta to non-cloud storage.
provider_meta is only used by cloud _build_step_record; non-cloud strategies fall back to StorageStrategy.record_steps_batch, which calls record_step(**step), and the non-cloud step schemas do not accept that key, so this raises for storage_type != "cloud". Skip adding it for non-cloud storage, or make all backends accept it.
🧰 Tools
🪛 ast-grep (0.45.0)
[info] 402-402: use jsonify instead of json.dumps for JSON output
Context: json.dumps(self._metadata(record), ensure_ascii=False, default=str)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gateway/storage.py` around lines 364 - 409, The step construction around
provider_meta must ensure provider_meta is never included for non-cloud storage.
Restrict the provider_meta assignment to the cloud storage path, while
preserving the existing Anthropic metadata behavior for cloud records and
keeping non-cloud record_steps_batch compatible with its step schema.
Comment on lines
+672
to
+710
| def _chat_completion_output(response: Any) -> Any: | ||
| payload = _json_loads(response) | ||
| if not isinstance(payload, dict): | ||
| return None | ||
|
|
||
| text_parts: list[str] = [] | ||
| thinking_parts: list[str] = [] | ||
| tool_calls: list[dict[str, Any]] = [] | ||
| for block in content: | ||
| if not isinstance(block, dict): | ||
| continue | ||
| block_type = block.get("type") | ||
| if block_type == "text" and isinstance(block.get("text"), str): | ||
| text_parts.append(block["text"]) | ||
| elif block_type == "thinking" and isinstance(block.get("thinking"), str): | ||
| thinking_parts.append(block["thinking"]) | ||
| elif block_type == "tool_use": | ||
| tool_calls.append( | ||
| { | ||
| "id": str(block.get("id") or ""), | ||
| "type": "function", | ||
| "function": { | ||
| "name": str(block.get("name") or ""), | ||
| "arguments": json.dumps(block.get("input") or {}, ensure_ascii=False), | ||
| }, | ||
| } | ||
| ) | ||
|
|
||
| message = _assistant_message_from_parts("".join(text_parts), "\n\n".join(thinking_parts)) | ||
| if message is None: | ||
| message = {"role": "assistant"} | ||
| if tool_calls: | ||
| message["tool_calls"] = tool_calls | ||
| return message if len(message) > 1 else None | ||
|
|
||
|
|
||
| def _assistant_message_from_chat_message(raw_message: dict[str, Any]) -> dict[str, Any] | None: | ||
| content = raw_message.get("content") | ||
| reasoning = _first_string(raw_message, ("think", "reasoning", "reasoning_content")) | ||
| message = _assistant_message_from_parts(content if isinstance(content, str) else None, reasoning) | ||
|
|
||
| if message is None: | ||
| message = {"role": "assistant"} | ||
| if content is not None and not isinstance(content, str): | ||
| message["content"] = content | ||
|
|
||
| tool_calls = raw_message.get("tool_calls") | ||
| if isinstance(tool_calls, list) and tool_calls: | ||
| message["tool_calls"] = tool_calls | ||
|
|
||
| function_call = raw_message.get("function_call") | ||
| if isinstance(function_call, dict) and function_call: | ||
| message["function_call"] = function_call | ||
|
|
||
| return message if len(message) > 1 else None | ||
|
|
||
|
|
||
| def _assistant_message_from_response_payload(payload: dict[str, Any]) -> dict[str, Any] | None: | ||
| content = _first_string(payload, ("output_text", "text", "content")) | ||
| reasoning = _first_string(payload, ("think", "reasoning_text", "reasoning", "reasoning_content")) | ||
| output_content, output_reasoning = _extract_responses_output(payload.get("output")) | ||
|
|
||
| if output_content: | ||
| content = (content or "") + output_content | ||
| if output_reasoning: | ||
| reasoning = "\n\n".join(part for part in (reasoning, output_reasoning) if part) | ||
|
|
||
| return _assistant_message_from_parts(content, reasoning) | ||
|
|
||
|
|
||
| def _assistant_message_from_stream(stream_text: Any) -> dict[str, Any] | None: | ||
| if isinstance(stream_text, str): | ||
| events = _iter_sse_events(stream_text) | ||
| elif isinstance(stream_text, list): | ||
| events = (event for event in stream_text if isinstance(event, dict)) | ||
| else: | ||
| choices = payload.get("choices") | ||
| if not isinstance(choices, list): | ||
| return None | ||
|
|
||
| content_parts: list[str] = [] | ||
| reasoning_parts: list[str] = [] | ||
| for event in events: | ||
| _collect_event_text(event, content_parts, reasoning_parts) | ||
| return _assistant_message_from_parts("".join(content_parts), "\n\n".join(reasoning_parts)) | ||
|
|
||
|
|
||
| def _collect_event_text( | ||
| event: dict[str, Any], | ||
| content_parts: list[str], | ||
| reasoning_parts: list[str], | ||
| ) -> None: | ||
| choices = event.get("choices") | ||
| if isinstance(choices, list): | ||
| for choice in choices: | ||
| if not isinstance(choice, dict): | ||
| continue | ||
| for key in ("message", "delta"): | ||
| payload = choice.get(key) | ||
| if isinstance(payload, dict): | ||
| _collect_message_text(payload, content_parts, reasoning_parts) | ||
|
|
||
| event_type = event.get("type") | ||
| delta = event.get("delta") | ||
| if event_type == "response.output_text.delta" and isinstance(delta, str): | ||
| content_parts.append(delta) | ||
| elif event_type == "response.reasoning_text.delta" and isinstance(delta, str): | ||
| reasoning_parts.append(delta) | ||
|
|
||
|
|
||
| def _collect_message_text( | ||
| payload: dict[str, Any], | ||
| content_parts: list[str], | ||
| reasoning_parts: list[str], | ||
| ) -> None: | ||
| content = payload.get("content") | ||
| if isinstance(content, str): | ||
| content_parts.append(content) | ||
| reasoning = _first_string(payload, ("think", "reasoning", "reasoning_content")) | ||
| if reasoning: | ||
| reasoning_parts.append(reasoning) | ||
|
|
||
|
|
||
| def _extract_responses_output(output: Any) -> tuple[str, str]: | ||
| if not isinstance(output, list): | ||
| return "", "" | ||
|
|
||
| content_parts: list[str] = [] | ||
| reasoning_parts: list[str] = [] | ||
| for item in output: | ||
| if not isinstance(item, dict): | ||
| continue | ||
| item_type = item.get("type") | ||
| if item_type == "reasoning": | ||
| reasoning = _first_string(item, ("text", "content", "summary")) | ||
| if reasoning: | ||
| reasoning_parts.append(reasoning) | ||
| continue | ||
| if item_type == "message": | ||
| content = item.get("content") | ||
| if isinstance(content, list): | ||
| for part in content: | ||
| if isinstance(part, dict): | ||
| text = _first_string(part, ("text", "content")) | ||
| if text: | ||
| content_parts.append(text) | ||
| else: | ||
| text = _first_string(item, ("text", "content")) | ||
| if text: | ||
| content_parts.append(text) | ||
| return "".join(content_parts), "\n\n".join(reasoning_parts) | ||
|
|
||
|
|
||
| def _assistant_message_from_parts( | ||
| content: str | None = None, | ||
| reasoning: str | None = None, | ||
| ) -> dict[str, Any] | None: | ||
| content = content or "" | ||
| tag_think, clean_content = _split_think_tags(content) | ||
| think = "\n\n".join(part.strip() for part in (reasoning, tag_think) if part and part.strip()) | ||
|
|
||
| message: dict[str, Any] = {"role": "assistant"} | ||
| if clean_content.strip(): | ||
| message["content"] = clean_content.strip() | ||
| elif content or think: | ||
| message["content"] = "" | ||
| if think: | ||
| message["think"] = think | ||
| return message if len(message) > 1 else None | ||
|
|
||
|
|
||
| def _split_think_tags(content: str) -> tuple[str, str]: | ||
| matches = [match.group(1).strip() for match in THINK_TAG_RE.finditer(content)] | ||
| if not matches: | ||
| return "", content | ||
| clean_content = THINK_TAG_RE.sub("", content) | ||
| return "\n\n".join(match for match in matches if match), clean_content | ||
|
|
||
|
|
||
| def _iter_sse_events(stream_text: str): | ||
| for raw_line in stream_text.splitlines(): | ||
| line = raw_line.strip() | ||
| if not line.startswith("data:"): | ||
| continue | ||
| data = line[len("data:") :].strip() | ||
| if not data or data == "[DONE]": | ||
| continue | ||
| parsed = _json_loads(data) | ||
| if isinstance(parsed, dict): | ||
| yield parsed | ||
| messages = [ | ||
| dict(choice["message"]) | ||
| for choice in choices | ||
| if isinstance(choice, dict) and isinstance(choice.get("message"), dict) | ||
| ] | ||
| if len(messages) == 1: | ||
| return messages[0] | ||
| return messages or None | ||
|
|
||
|
|
||
| def _responses_output(response: Any) -> Any: | ||
| payload = _json_loads(response) | ||
| if not isinstance(payload, dict): | ||
| return None | ||
| output = payload.get("output") | ||
| if isinstance(output, (dict, list)): | ||
| return output | ||
| if payload.get("type") in {"message", "function_call", "reasoning"}: | ||
| return payload | ||
|
|
||
| # Streaming summaries may only contain the aggregated model text. Keep the | ||
| # derived message free of HTTP envelope fields. | ||
| output_text = payload.get("output_text") | ||
| reasoning_text = payload.get("reasoning_text") | ||
| if not isinstance(output_text, str) and not isinstance(reasoning_text, str): | ||
| return None | ||
| message: dict[str, Any] = {"role": "assistant", "content": []} | ||
| if isinstance(output_text, str): | ||
| message["content"].append({"type": "output_text", "text": output_text}) | ||
| if isinstance(reasoning_text, str): | ||
| message["reasoning"] = reasoning_text | ||
| return message |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Returning None discards error and non-conforming response bodies.
Both helpers return None whenever the response does not match the expected success shape. Concrete cases:
- An upstream error body such as
{"error": {"message": "..."}}has nochoicesand nooutput, so_chat_completion_outputand_responses_outputboth returnNone. - A non-JSON upstream body (HTML error page, proxy response) fails
_json_loadsand returnsNone. - A stream that terminates before any choice delta produces a summary without
choices, so_chat_completion_outputreturnsNone.
In each case line 401 stores response=None. The previous cloud behavior stored record.response. Failure telemetry now loses the upstream body entirely, and only error_text in env_state remains.
Return the parsed payload as a fallback instead of None.
🐛 Proposed fix to preserve non-conforming bodies
def _chat_completion_output(response: Any) -> Any:
payload = _json_loads(response)
if not isinstance(payload, dict):
- return None
+ return payload if payload is not None else response
choices = payload.get("choices")
if not isinstance(choices, list):
- return None
+ return payload
messages = [
dict(choice["message"])
for choice in choices
if isinstance(choice, dict) and isinstance(choice.get("message"), dict)
]
if len(messages) == 1:
return messages[0]
- return messages or None
+ return messages or payload
def _responses_output(response: Any) -> Any:
payload = _json_loads(response)
if not isinstance(payload, dict):
- return None
+ return payload if payload is not None else response
output = payload.get("output")
if isinstance(output, (dict, list)):
return output
if payload.get("type") in {"message", "function_call", "reasoning"}:
return payload
# Streaming summaries may only contain the aggregated model text. Keep the
# derived message free of HTTP envelope fields.
output_text = payload.get("output_text")
reasoning_text = payload.get("reasoning_text")
if not isinstance(output_text, str) and not isinstance(reasoning_text, str):
- return None
+ return payload📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def _chat_completion_output(response: Any) -> Any: | |
| payload = _json_loads(response) | |
| if not isinstance(payload, dict): | |
| return None | |
| text_parts: list[str] = [] | |
| thinking_parts: list[str] = [] | |
| tool_calls: list[dict[str, Any]] = [] | |
| for block in content: | |
| if not isinstance(block, dict): | |
| continue | |
| block_type = block.get("type") | |
| if block_type == "text" and isinstance(block.get("text"), str): | |
| text_parts.append(block["text"]) | |
| elif block_type == "thinking" and isinstance(block.get("thinking"), str): | |
| thinking_parts.append(block["thinking"]) | |
| elif block_type == "tool_use": | |
| tool_calls.append( | |
| { | |
| "id": str(block.get("id") or ""), | |
| "type": "function", | |
| "function": { | |
| "name": str(block.get("name") or ""), | |
| "arguments": json.dumps(block.get("input") or {}, ensure_ascii=False), | |
| }, | |
| } | |
| ) | |
| message = _assistant_message_from_parts("".join(text_parts), "\n\n".join(thinking_parts)) | |
| if message is None: | |
| message = {"role": "assistant"} | |
| if tool_calls: | |
| message["tool_calls"] = tool_calls | |
| return message if len(message) > 1 else None | |
| def _assistant_message_from_chat_message(raw_message: dict[str, Any]) -> dict[str, Any] | None: | |
| content = raw_message.get("content") | |
| reasoning = _first_string(raw_message, ("think", "reasoning", "reasoning_content")) | |
| message = _assistant_message_from_parts(content if isinstance(content, str) else None, reasoning) | |
| if message is None: | |
| message = {"role": "assistant"} | |
| if content is not None and not isinstance(content, str): | |
| message["content"] = content | |
| tool_calls = raw_message.get("tool_calls") | |
| if isinstance(tool_calls, list) and tool_calls: | |
| message["tool_calls"] = tool_calls | |
| function_call = raw_message.get("function_call") | |
| if isinstance(function_call, dict) and function_call: | |
| message["function_call"] = function_call | |
| return message if len(message) > 1 else None | |
| def _assistant_message_from_response_payload(payload: dict[str, Any]) -> dict[str, Any] | None: | |
| content = _first_string(payload, ("output_text", "text", "content")) | |
| reasoning = _first_string(payload, ("think", "reasoning_text", "reasoning", "reasoning_content")) | |
| output_content, output_reasoning = _extract_responses_output(payload.get("output")) | |
| if output_content: | |
| content = (content or "") + output_content | |
| if output_reasoning: | |
| reasoning = "\n\n".join(part for part in (reasoning, output_reasoning) if part) | |
| return _assistant_message_from_parts(content, reasoning) | |
| def _assistant_message_from_stream(stream_text: Any) -> dict[str, Any] | None: | |
| if isinstance(stream_text, str): | |
| events = _iter_sse_events(stream_text) | |
| elif isinstance(stream_text, list): | |
| events = (event for event in stream_text if isinstance(event, dict)) | |
| else: | |
| choices = payload.get("choices") | |
| if not isinstance(choices, list): | |
| return None | |
| content_parts: list[str] = [] | |
| reasoning_parts: list[str] = [] | |
| for event in events: | |
| _collect_event_text(event, content_parts, reasoning_parts) | |
| return _assistant_message_from_parts("".join(content_parts), "\n\n".join(reasoning_parts)) | |
| def _collect_event_text( | |
| event: dict[str, Any], | |
| content_parts: list[str], | |
| reasoning_parts: list[str], | |
| ) -> None: | |
| choices = event.get("choices") | |
| if isinstance(choices, list): | |
| for choice in choices: | |
| if not isinstance(choice, dict): | |
| continue | |
| for key in ("message", "delta"): | |
| payload = choice.get(key) | |
| if isinstance(payload, dict): | |
| _collect_message_text(payload, content_parts, reasoning_parts) | |
| event_type = event.get("type") | |
| delta = event.get("delta") | |
| if event_type == "response.output_text.delta" and isinstance(delta, str): | |
| content_parts.append(delta) | |
| elif event_type == "response.reasoning_text.delta" and isinstance(delta, str): | |
| reasoning_parts.append(delta) | |
| def _collect_message_text( | |
| payload: dict[str, Any], | |
| content_parts: list[str], | |
| reasoning_parts: list[str], | |
| ) -> None: | |
| content = payload.get("content") | |
| if isinstance(content, str): | |
| content_parts.append(content) | |
| reasoning = _first_string(payload, ("think", "reasoning", "reasoning_content")) | |
| if reasoning: | |
| reasoning_parts.append(reasoning) | |
| def _extract_responses_output(output: Any) -> tuple[str, str]: | |
| if not isinstance(output, list): | |
| return "", "" | |
| content_parts: list[str] = [] | |
| reasoning_parts: list[str] = [] | |
| for item in output: | |
| if not isinstance(item, dict): | |
| continue | |
| item_type = item.get("type") | |
| if item_type == "reasoning": | |
| reasoning = _first_string(item, ("text", "content", "summary")) | |
| if reasoning: | |
| reasoning_parts.append(reasoning) | |
| continue | |
| if item_type == "message": | |
| content = item.get("content") | |
| if isinstance(content, list): | |
| for part in content: | |
| if isinstance(part, dict): | |
| text = _first_string(part, ("text", "content")) | |
| if text: | |
| content_parts.append(text) | |
| else: | |
| text = _first_string(item, ("text", "content")) | |
| if text: | |
| content_parts.append(text) | |
| return "".join(content_parts), "\n\n".join(reasoning_parts) | |
| def _assistant_message_from_parts( | |
| content: str | None = None, | |
| reasoning: str | None = None, | |
| ) -> dict[str, Any] | None: | |
| content = content or "" | |
| tag_think, clean_content = _split_think_tags(content) | |
| think = "\n\n".join(part.strip() for part in (reasoning, tag_think) if part and part.strip()) | |
| message: dict[str, Any] = {"role": "assistant"} | |
| if clean_content.strip(): | |
| message["content"] = clean_content.strip() | |
| elif content or think: | |
| message["content"] = "" | |
| if think: | |
| message["think"] = think | |
| return message if len(message) > 1 else None | |
| def _split_think_tags(content: str) -> tuple[str, str]: | |
| matches = [match.group(1).strip() for match in THINK_TAG_RE.finditer(content)] | |
| if not matches: | |
| return "", content | |
| clean_content = THINK_TAG_RE.sub("", content) | |
| return "\n\n".join(match for match in matches if match), clean_content | |
| def _iter_sse_events(stream_text: str): | |
| for raw_line in stream_text.splitlines(): | |
| line = raw_line.strip() | |
| if not line.startswith("data:"): | |
| continue | |
| data = line[len("data:") :].strip() | |
| if not data or data == "[DONE]": | |
| continue | |
| parsed = _json_loads(data) | |
| if isinstance(parsed, dict): | |
| yield parsed | |
| messages = [ | |
| dict(choice["message"]) | |
| for choice in choices | |
| if isinstance(choice, dict) and isinstance(choice.get("message"), dict) | |
| ] | |
| if len(messages) == 1: | |
| return messages[0] | |
| return messages or None | |
| def _responses_output(response: Any) -> Any: | |
| payload = _json_loads(response) | |
| if not isinstance(payload, dict): | |
| return None | |
| output = payload.get("output") | |
| if isinstance(output, (dict, list)): | |
| return output | |
| if payload.get("type") in {"message", "function_call", "reasoning"}: | |
| return payload | |
| # Streaming summaries may only contain the aggregated model text. Keep the | |
| # derived message free of HTTP envelope fields. | |
| output_text = payload.get("output_text") | |
| reasoning_text = payload.get("reasoning_text") | |
| if not isinstance(output_text, str) and not isinstance(reasoning_text, str): | |
| return None | |
| message: dict[str, Any] = {"role": "assistant", "content": []} | |
| if isinstance(output_text, str): | |
| message["content"].append({"type": "output_text", "text": output_text}) | |
| if isinstance(reasoning_text, str): | |
| message["reasoning"] = reasoning_text | |
| return message | |
| def _chat_completion_output(response: Any) -> Any: | |
| payload = _json_loads(response) | |
| if not isinstance(payload, dict): | |
| return payload if payload is not None else response | |
| choices = payload.get("choices") | |
| if not isinstance(choices, list): | |
| return payload | |
| messages = [ | |
| dict(choice["message"]) | |
| for choice in choices | |
| if isinstance(choice, dict) and isinstance(choice.get("message"), dict) | |
| ] | |
| if len(messages) == 1: | |
| return messages[0] | |
| return messages or payload | |
| def _responses_output(response: Any) -> Any: | |
| payload = _json_loads(response) | |
| if not isinstance(payload, dict): | |
| return payload if payload is not None else response | |
| output = payload.get("output") | |
| if isinstance(output, (dict, list)): | |
| return output | |
| if payload.get("type") in {"message", "function_call", "reasoning"}: | |
| return payload | |
| # Streaming summaries may only contain the aggregated model text. Keep the | |
| # derived message free of HTTP envelope fields. | |
| output_text = payload.get("output_text") | |
| reasoning_text = payload.get("reasoning_text") | |
| if not isinstance(output_text, str) and not isinstance(reasoning_text, str): | |
| return payload | |
| message: dict[str, Any] = {"role": "assistant", "content": []} | |
| if isinstance(output_text, str): | |
| message["content"].append({"type": "output_text", "text": output_text}) | |
| if isinstance(reasoning_text, str): | |
| message["reasoning"] = reasoning_text | |
| return message |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@gateway/storage.py` around lines 672 - 710, Update _chat_completion_output
and _responses_output to return the parsed payload as the fallback whenever the
expected success shape is absent, including invalid or non-JSON response bodies
according to the existing _json_loads behavior. Preserve the current extraction
of valid messages and output values, but replace None fallbacks so response
storage retains upstream error, incomplete-stream, and non-conforming bodies.
hsballoon
approved these changes
Aug 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Bug Fixes