Open
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds WebSocket-mode support for the Responses backend: new persistent WebSocket transport, websocket-aware streaming in ResponseBackend, websocket_mode config/validation and API-parameter exclusion, new WebSocket configs, extensive tests, and small formatter tweak preserving messages with an "id". Changes
Sequence Diagram(s)mermaid Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
📝 Coding Plan
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 |
praneeth999
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
massgen/api_params_handler/_response_api_params_handler.py (1)
74-77: Consider aligningChatCompletionsAPIParamsHandlerwithwebsocket_modehandling.Only
ResponseAPIParamsHandlerimplements thewebsocket_modegate on streaming. TheChatCompletionsAPIParamsHandler(per context snippet 3) unconditionally setsstream: True. If WebSocket mode should apply to both backends, the sibling handler may need similar logic.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/api_params_handler/_response_api_params_handler.py` around lines 74 - 77, The ChatCompletionsAPIParamsHandler unconditionally sets api_params["stream"] = True while ResponseAPIParamsHandler gates streaming with websocket_mode; update ChatCompletionsAPIParamsHandler to read websocket_mode = all_params.get("websocket_mode", False) (same as in ResponseAPIParamsHandler) and only set api_params["stream"] = True when websocket_mode is False so WebSocket mode disables backend streaming consistently; look for the websocket_mode variable usage and the api_params["stream"] assignment in ChatCompletionsAPIParamsHandler to apply the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@massgen/api_params_handler/_response_api_params_handler.py`:
- Around line 74-77: The ChatCompletionsAPIParamsHandler unconditionally sets
api_params["stream"] = True while ResponseAPIParamsHandler gates streaming with
websocket_mode; update ChatCompletionsAPIParamsHandler to read websocket_mode =
all_params.get("websocket_mode", False) (same as in ResponseAPIParamsHandler)
and only set api_params["stream"] = True when websocket_mode is False so
WebSocket mode disables backend streaming consistently; look for the
websocket_mode variable usage and the api_params["stream"] assignment in
ChatCompletionsAPIParamsHandler to apply the change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 03601812-a7a7-4b93-a471-258aafbbe7cd
📒 Files selected for processing (5)
massgen/api_params_handler/_api_params_handler_base.pymassgen/api_params_handler/_response_api_params_handler.pymassgen/backend/base.pymassgen/config_validator.pypyproject.toml
praneeth999
force-pushed
the
Web_socket_mode
branch
from
3d6824c to
09a829c
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (5)
massgen/backend/_websocket_transport.py (3)
96-137: Add return type hint and handle unexpected disconnection.The
send_and_receivemethod is missing a return type annotation. Additionally, if the WebSocket connection drops unexpectedly during iteration (line 116), the exception will propagate unhandled. Consider adding explicit handling for connection loss.♻️ Suggested improvements
+from collections.abc import AsyncGenerator +from websockets.exceptions import ConnectionClosed + async def send_and_receive( self, api_params: dict[str, Any], - ): + ) -> AsyncGenerator[dict[str, Any], None]: """Send a response.create event and yield parsed response events. Args: api_params: The API params dict (same as HTTP body, minus stream/background). Yields: Parsed event dicts with a "type" field matching the HTTP SSE event types (e.g. "response.output_text.delta", "response.completed"). + + Raises: + WebSocketConnectionError: If not connected or connection lost during receive. """ if self._ws is None: raise WebSocketConnectionError("Not connected. Call connect() first.") message = self._build_response_create_event(api_params) await self._ws.send(message) logger.debug("[WebSocket] Sent response.create event") - async for raw_message in self._ws: + try: + async for raw_message in self._ws:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/backend/_websocket_transport.py` around lines 96 - 137, Add an explicit async return type to send_and_receive (e.g., AsyncIterator[dict[str, Any]] or AsyncGenerator[dict[str, Any], None]) and wrap the async iteration over self._ws in a try/except that catches connection-related exceptions (e.g., websockets ConnectionClosed, ConnectionResetError, asyncio.CancelledError or a broad Exception as a fallback) and re-raises them as WebSocketConnectionError with a clear message; update any necessary typing imports (typing.AsyncIterator/AsyncGenerator) and reference send_and_receive, self._ws, and WebSocketConnectionError when making the change.
41-50: Add docstring for__init__method.Per coding guidelines, new functions should include Google-style docstrings.
📝 Suggested docstring
def __init__( self, api_key: str, url: str = DEFAULT_WS_URL, organization: str | None = None, ): + """Initialize the WebSocket transport. + + Args: + api_key: OpenAI API key for authentication. + url: WebSocket endpoint URL. + organization: Optional OpenAI organization ID. + """ self.api_key = api_key🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/backend/_websocket_transport.py` around lines 41 - 50, Add a Google-style docstring to the __init__ method of the WebSocket transport class (the __init__ that accepts api_key, url, and organization and sets self._ws), documenting the purpose of the constructor, each parameter (api_key: str, url: str = DEFAULT_WS_URL, organization: Optional[str]) and the attributes initialized (self.api_key, self.url, self.organization, self._ws), and include types and brief behavior notes (e.g., default URL and that _ws is initialized to None).
83-83: Consider catching more specific exceptions for connection retries.The blind
except Exceptioncatches all exceptions includingKeyboardInterruptandSystemExit. For WebSocket connection attempts, catchingwebsocketslibrary exceptions and standard connection errors would be more appropriate.♻️ Suggested improvement
+from websockets.exceptions import WebSocketException + - except Exception as e: + except (OSError, WebSocketException) as e:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/backend/_websocket_transport.py` at line 83, The catch-all "except Exception as e:" in the connection/retry logic should be replaced with specific exception types so we don't swallow critical signals like KeyboardInterrupt/SystemExit; update the except block in the websocket connection routine (the block containing "except Exception as e:") to catch websocket and network-related errors only (for example websockets.exceptions.ConnectionClosedError, websockets.exceptions.InvalidHandshake/InvalidURI, asyncio.TimeoutError, OSError) and handle/retry those, while letting other exceptions propagate. Ensure the variable "e" is preserved in the new except clauses for logging and that any generic cleanup still runs outside these specific exception handlers.massgen/configs/providers/openai/gpt5_2_websocket.yaml (1)
1-13: Add "What happens" comment explaining execution flow.Per coding guidelines, YAML configs should include comments explaining execution flow to help users understand the configuration behavior.
📝 Suggested improvement
# GPT-5.2 with WebSocket Mode # Single agent using the latest GPT-5.2 model over persistent WebSocket. +# +# What happens: +# 1. A persistent WebSocket connection is established to OpenAI's responses API +# 2. All requests/responses flow over this single connection (lower latency) +# 3. Web search and code interpreter tools are available for the agent agents: - id: "gpt-5-2-ws"As per coding guidelines: "Include 'What happens' comments explaining execution flow"
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/configs/providers/openai/gpt5_2_websocket.yaml` around lines 1 - 13, Add a "What happens" comment at the top of the YAML describing execution flow for the agent with id "gpt-5-2-ws": explain that the backend type "openai" will use model "gpt-5.2" in websocket_mode (persistent connection), note that enable_web_search and enable_code_interpreter toggle browsing and code execution capabilities, and state how UI settings (display_type "textual_terminal" and logging_enabled true) affect runtime output and logs; keep the comment concise and placed above the agents block so users immediately see the behavior.massgen/configs/providers/openai/multi_model_websocket.yaml (1)
1-18: Consider adding orchestrator and system_message for multi-agent setup.This config defines two agents without an
orchestratorsection orsystem_messagefields. Per coding guidelines, multi-agent setups should have identical system messages for all agents, and coordination behavior should be explicit.For a minimal example config this may be acceptable, but users might be confused about how the agents coordinate.
📝 Suggested additions for clarity
# Multi-Model WebSocket Mode # Two agents with different models, both using WebSocket transport. +# +# What happens: +# 1. Both agents establish persistent WebSocket connections +# 2. Default voting coordination is used to reconcile agent outputs agents: - id: "gpt-5-2" + system_message: "You are a helpful AI assistant." backend: type: "openai" model: "gpt-5.2" websocket_mode: true enable_code_interpreter: true - id: "gpt-5-nano" + system_message: "You are a helpful AI assistant." backend: type: "openai" model: "gpt-5-nano" websocket_mode: true enable_code_interpreter: true +orchestrator: + voting_sensitivity: "balanced" ui: display_type: "textual_terminal" logging_enabled: trueAs per coding guidelines: "All agents should have identical system_message for multi-agent setups"
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/configs/providers/openai/multi_model_websocket.yaml` around lines 1 - 18, Add an explicit orchestrator section and identical system_message fields for both agents to clarify coordination: update the config to include an "orchestrator" entry describing coordination mode (e.g., turn-taking or leader-election) and add the same "system_message" text to the agent blocks for "gpt-5-2" and "gpt-5-nano" so both agents share identical system instructions and the orchestration behavior is explicit.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@massgen/backend/_websocket_transport.py`:
- Around line 96-137: Add an explicit async return type to send_and_receive
(e.g., AsyncIterator[dict[str, Any]] or AsyncGenerator[dict[str, Any], None])
and wrap the async iteration over self._ws in a try/except that catches
connection-related exceptions (e.g., websockets ConnectionClosed,
ConnectionResetError, asyncio.CancelledError or a broad Exception as a fallback)
and re-raises them as WebSocketConnectionError with a clear message; update any
necessary typing imports (typing.AsyncIterator/AsyncGenerator) and reference
send_and_receive, self._ws, and WebSocketConnectionError when making the change.
- Around line 41-50: Add a Google-style docstring to the __init__ method of the
WebSocket transport class (the __init__ that accepts api_key, url, and
organization and sets self._ws), documenting the purpose of the constructor,
each parameter (api_key: str, url: str = DEFAULT_WS_URL, organization:
Optional[str]) and the attributes initialized (self.api_key, self.url,
self.organization, self._ws), and include types and brief behavior notes (e.g.,
default URL and that _ws is initialized to None).
- Line 83: The catch-all "except Exception as e:" in the connection/retry logic
should be replaced with specific exception types so we don't swallow critical
signals like KeyboardInterrupt/SystemExit; update the except block in the
websocket connection routine (the block containing "except Exception as e:") to
catch websocket and network-related errors only (for example
websockets.exceptions.ConnectionClosedError,
websockets.exceptions.InvalidHandshake/InvalidURI, asyncio.TimeoutError,
OSError) and handle/retry those, while letting other exceptions propagate.
Ensure the variable "e" is preserved in the new except clauses for logging and
that any generic cleanup still runs outside these specific exception handlers.
In `@massgen/configs/providers/openai/gpt5_2_websocket.yaml`:
- Around line 1-13: Add a "What happens" comment at the top of the YAML
describing execution flow for the agent with id "gpt-5-2-ws": explain that the
backend type "openai" will use model "gpt-5.2" in websocket_mode (persistent
connection), note that enable_web_search and enable_code_interpreter toggle
browsing and code execution capabilities, and state how UI settings
(display_type "textual_terminal" and logging_enabled true) affect runtime output
and logs; keep the comment concise and placed above the agents block so users
immediately see the behavior.
In `@massgen/configs/providers/openai/multi_model_websocket.yaml`:
- Around line 1-18: Add an explicit orchestrator section and identical
system_message fields for both agents to clarify coordination: update the config
to include an "orchestrator" entry describing coordination mode (e.g.,
turn-taking or leader-election) and add the same "system_message" text to the
agent blocks for "gpt-5-2" and "gpt-5-nano" so both agents share identical
system instructions and the orchestration behavior is explicit.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 47e0b97a-50c3-4682-816b-cc47a51e3cfa
📒 Files selected for processing (9)
massgen/api_params_handler/_api_params_handler_base.pymassgen/api_params_handler/_response_api_params_handler.pymassgen/backend/_websocket_transport.pymassgen/backend/base.pymassgen/config_validator.pymassgen/configs/providers/openai/gpt5_2_websocket.yamlmassgen/configs/providers/openai/gpt5_nano_websocket.yamlmassgen/configs/providers/openai/multi_model_websocket.yamlpyproject.toml
🚧 Files skipped from review as they are similar to previous changes (3)
- massgen/api_params_handler/_response_api_params_handler.py
- pyproject.toml
- massgen/backend/base.py
praneeth999
force-pushed
the
Web_socket_mode
branch
from
3e56d15 to
c80c38f
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
massgen/backend/_websocket_transport.py (1)
96-137: Add return type hint for the async generator.The method is missing a return type annotation. For async generators, use
AsyncGeneratorfromtypingorcollections.abc.📝 Suggested improvement
+from collections.abc import AsyncGenerator + ... async def send_and_receive( self, api_params: dict[str, Any], - ): + ) -> AsyncGenerator[dict[str, Any], None]: """Send a response.create event and yield parsed response events.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/backend/_websocket_transport.py` around lines 96 - 137, The async generator method send_and_receive lacks a return type annotation; update its signature to return an AsyncGenerator of event dicts (e.g. AsyncGenerator[dict[str, Any], None]) and add the corresponding import (from typing or collections.abc import AsyncGenerator) at the top of the module; ensure the symbol names to edit are send_and_receive and the module imports so tooling and type checkers recognize the async generator return type.massgen/configs/providers/openai/multi_model_websocket.yaml (1)
1-18: Consider adding a "What happens" comment explaining the execution flow.Per coding guidelines for YAML configs, include comments explaining what happens when this configuration is used. This helps users understand the WebSocket transport behavior.
📝 Suggested improvement
# Multi-Model WebSocket Mode # Two agents with different models, both using WebSocket transport. +# +# What happens: +# - Both agents use persistent WebSocket connections instead of HTTP streaming +# - The transport layer handles connection lifecycle and message streaming +# - Code interpreter is enabled for both agents agents: - id: "gpt-5-2"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/configs/providers/openai/multi_model_websocket.yaml` around lines 1 - 18, Add a short "What happens" comment at the top of this YAML (above "Multi-Model WebSocket Mode") describing the execution flow: that the two agents listed under agents (ids "gpt-5-2" and "gpt-5-nano") will run concurrently using OpenAI backends with websocket_mode: true so their messages are streamed over WebSocket, that enable_code_interpreter: true enables the code-execution extension for each agent, and that ui.display_type: "textual_terminal" with logging_enabled: true will present streamed output in the terminal and record logs; place this explanatory comment near the file header so users immediately see the runtime behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@massgen/api_params_handler/_response_api_params_handler.py`:
- Around line 77-82: The code mutates the input dict by calling
all_params.pop("background", None), which can cause unexpected side effects;
change the logic to operate on a shallow copy of all_params (e.g., work with a
new dict variable before modifying) or explicitly build a new params dict
instead of mutating all_params so websocket_mode, api_params, and the rest of
the flow remain the same; update the code paths that reference websocket_mode,
api_params, and any later use of all_params to use the new copy or constructed
dict, or document the mutation in the function docstring if mutation is
intentional.
---
Nitpick comments:
In `@massgen/backend/_websocket_transport.py`:
- Around line 96-137: The async generator method send_and_receive lacks a return
type annotation; update its signature to return an AsyncGenerator of event dicts
(e.g. AsyncGenerator[dict[str, Any], None]) and add the corresponding import
(from typing or collections.abc import AsyncGenerator) at the top of the module;
ensure the symbol names to edit are send_and_receive and the module imports so
tooling and type checkers recognize the async generator return type.
In `@massgen/configs/providers/openai/multi_model_websocket.yaml`:
- Around line 1-18: Add a short "What happens" comment at the top of this YAML
(above "Multi-Model WebSocket Mode") describing the execution flow: that the two
agents listed under agents (ids "gpt-5-2" and "gpt-5-nano") will run
concurrently using OpenAI backends with websocket_mode: true so their messages
are streamed over WebSocket, that enable_code_interpreter: true enables the
code-execution extension for each agent, and that ui.display_type:
"textual_terminal" with logging_enabled: true will present streamed output in
the terminal and record logs; place this explanatory comment near the file
header so users immediately see the runtime behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: aadbdc46-c52e-4aa2-a8ae-eda50b4f5acd
📒 Files selected for processing (11)
massgen/api_params_handler/_api_params_handler_base.pymassgen/api_params_handler/_response_api_params_handler.pymassgen/backend/_websocket_transport.pymassgen/backend/base.pymassgen/config_validator.pymassgen/configs/providers/openai/gpt5_2_websocket.yamlmassgen/configs/providers/openai/gpt5_nano_websocket.yamlmassgen/configs/providers/openai/multi_model_websocket.yamlmassgen/formatter/_response_formatter.pymassgen/tests/test_websocket_mode.pypyproject.toml
🚧 Files skipped from review as they are similar to previous changes (7)
- massgen/config_validator.py
- massgen/configs/providers/openai/gpt5_2_websocket.yaml
- massgen/backend/base.py
- pyproject.toml
- massgen/configs/providers/openai/gpt5_nano_websocket.yaml
- massgen/tests/test_websocket_mode.py
- massgen/api_params_handler/_api_params_handler_base.py
Comment on lines
+77
to
+82
| websocket_mode = all_params.get("websocket_mode", False) # In WebSocket mode, stream/background are not used (transport handles streaming) | ||
| api_params = {"input": converted_messages} | ||
| if not websocket_mode: | ||
| api_params["stream"] = True | ||
| else: | ||
| all_params.pop("background", None) |
There was a problem hiding this comment.
Mutating all_params may cause unintended side effects.
The all_params.pop("background", None) call modifies the input dictionary in place. If the caller reuses all_params after this call, the background key will be unexpectedly missing. Consider operating on a copy or documenting this mutation.
🛡️ Suggested fix to avoid mutation
websocket_mode = all_params.get("websocket_mode", False) # In WebSocket mode, stream/background are not used (transport handles streaming)
api_params = {"input": converted_messages}
if not websocket_mode:
api_params["stream"] = True
else:
- all_params.pop("background", None)
+ # Don't include 'background' in api_params for WebSocket mode
+ # (handled below via excluded params or explicit skip)Alternatively, if background must be removed to prevent it from being added later in the loop, ensure this mutation is documented in the docstring.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@massgen/api_params_handler/_response_api_params_handler.py` around lines 77 -
82, The code mutates the input dict by calling all_params.pop("background",
None), which can cause unexpected side effects; change the logic to operate on a
shallow copy of all_params (e.g., work with a new dict variable before
modifying) or explicitly build a new params dict instead of mutating all_params
so websocket_mode, api_params, and the rest of the flow remain the same; update
the code paths that reference websocket_mode, api_params, and any later use of
all_params to use the new copy or constructed dict, or document the mutation in
the function docstring if mutation is intentional.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
massgen/backend/response.py (1)
40-81: Add Google-style docstrings to the new websocket helpers.
_WSEvent’s new helper methods plus_create_response_stream(),_ws_event_stream(), and_create_client()still use minimal or non-Google docstrings. Please add Google-styleArgs/Returns/Yieldssections so the new websocket surface matches the repo standard.As per coding guidelines,
**/*.py: For new or changed functions, include Google-style docstringsAlso applies to: 1724-1744
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/backend/response.py` around lines 40 - 81, Add Google-style docstrings to the websocket helper classes and functions so they match the repo standard: for class _WSEvent add a class docstring and per-method docstrings (including __init__, _wrap, __getattr__, __getitem__, get, model_dump, etc.) with Args and Returns sections where relevant; for the functions _create_response_stream, _ws_event_stream, and _create_client add docstrings that describe purpose, include Args, Returns and Yields (for generators/streams) as appropriate and mention any raised exceptions; ensure the text is concise and follows the existing Google-style format used elsewhere in the codebase.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@massgen/backend/response.py`:
- Around line 565-584: The code uses a single current_function_call slot which
gets overwritten when multiple concurrent function calls stream; replace
current_function_call with a dict (e.g., in_flight_calls) keyed by the unique
item id/call id (use getattr(chunk.item, "call_id") or chunk.item.item_id if
present) and update the three handlers: in response.output_item.added create an
entry in in_flight_calls[item_id] with name/arguments/call_id, in
response.function_call_arguments.delta append deltas to
in_flight_calls[item_id]["arguments"] only when the chunk's item_id matches an
existing key, and in response.output_item.done move in_flight_calls[item_id]
into captured_function_calls and delete that key; ensure all references to
current_function_call are replaced accordingly so concurrent streams are tracked
independently.
---
Nitpick comments:
In `@massgen/backend/response.py`:
- Around line 40-81: Add Google-style docstrings to the websocket helper classes
and functions so they match the repo standard: for class _WSEvent add a class
docstring and per-method docstrings (including __init__, _wrap, __getattr__,
__getitem__, get, model_dump, etc.) with Args and Returns sections where
relevant; for the functions _create_response_stream, _ws_event_stream, and
_create_client add docstrings that describe purpose, include Args, Returns and
Yields (for generators/streams) as appropriate and mention any raised
exceptions; ensure the text is concise and follows the existing Google-style
format used elsewhere in the codebase.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 7d86f495-c2d8-4226-a339-483f92b80251
📒 Files selected for processing (1)
massgen/backend/response.py
Comment on lines
+565
to
+584
| # Detect function call start | ||
| if chunk.type == "response.output_item.added" and hasattr(chunk, "item") and chunk.item and getattr(chunk.item, "type", None) == "function_call": | ||
| current_function_call = { | ||
| "call_id": getattr(chunk.item, "call_id", ""), | ||
| "name": getattr(chunk.item, "name", ""), | ||
| "arguments": "", | ||
| } | ||
| logger.info( | ||
| f"Function call detected: {current_function_call['name']}", | ||
| ) | ||
|
|
||
| # Handle other streaming events (reasoning, provider tools, etc.) | ||
| else: | ||
| result = self._process_stream_chunk(chunk, agent_id) | ||
| yield result | ||
|
|
||
| # Response completed | ||
| if chunk.type in ["response.completed", "response.incomplete"]: | ||
| response_completed = True | ||
| # Note: Usage tracking is handled in _process_stream_chunk() above | ||
| # Capture response ID and ALL output items for reasoning continuity | ||
| if hasattr(chunk, "response") and chunk.response: | ||
| response_id = getattr(chunk.response, "id", None) | ||
| if response_id: | ||
| logger.debug(f"Captured response ID for reasoning continuity: {response_id}") | ||
| # CRITICAL: Capture ALL output items (reasoning, function_call, message) | ||
| # These must be included in the next request for reasoning models | ||
| output = getattr(chunk.response, "output", []) | ||
| if output: | ||
| for item in output: | ||
| # Convert to dict format for the API | ||
| item_dict = self._convert_to_dict(item) if hasattr(item, "model_dump") or hasattr(item, "dict") else item | ||
| if isinstance(item_dict, dict): | ||
| response_output_items.append(item_dict) | ||
| logger.debug(f"Captured {len(response_output_items)} output items for reasoning continuity") | ||
| if captured_function_calls: | ||
| # Execute captured function calls and recurse | ||
| self.end_api_call_timing(success=True) | ||
| break # Exit chunk loop to execute functions | ||
| # Accumulate function arguments | ||
| elif chunk.type == "response.function_call_arguments.delta" and current_function_call is not None: | ||
| delta = getattr(chunk, "delta", "") | ||
| current_function_call["arguments"] += delta | ||
|
|
||
| # Function call completed | ||
| elif chunk.type == "response.output_item.done" and current_function_call is not None: | ||
| captured_function_calls.append(current_function_call) | ||
| current_function_call = None |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n massgen/backend/response.py | head -600 | tail -100Repository: massgen/MassGen
Length of output: 5628
🏁 Script executed:
cat -n massgen/backend/response.py | sed -n '520,600p'Repository: massgen/MassGen
Length of output: 4564
🏁 Script executed:
rg -n "output_index|item_id" massgen/Repository: massgen/MassGen
Length of output: 30636
🏁 Script executed:
rg -n "function_call|function call" massgen/backend/response.py -A 2 -B 2Repository: massgen/MassGen
Length of output: 7726
🏁 Script executed:
sed -n '620,640p' massgen/backend/response.pyRepository: massgen/MassGen
Length of output: 1499
🏁 Script executed:
rg -n "captured_function_calls|response_output_items" massgen/backend/response.py | head -20Repository: massgen/MassGen
Length of output: 1340
🏁 Script executed:
fd -e py -path "*/test*" | xargs rg -l "function_call.*function_call\|multi.*tool\|multiple.*call" 2>/dev/null || echo "No matches found"Repository: massgen/MassGen
Length of output: 261
🏁 Script executed:
sed -n '657,680p' massgen/backend/response.pyRepository: massgen/MassGen
Length of output: 1305
🏁 Script executed:
sed -n '576,585p' massgen/backend/response.pyRepository: massgen/MassGen
Length of output: 642
🏁 Script executed:
sed -n '576,600p' massgen/backend/response.pyRepository: massgen/MassGen
Length of output: 1464
🏁 Script executed:
sed -n '1077,1086p' massgen/backend/docs/Function\ calling\ openai\ responses.mdRepository: massgen/MassGen
Length of output: 1586
🏁 Script executed:
grep -A 10 "response.function_call_arguments.delta" massgen/backend/docs/Function\ calling\ openai\ responses.md | head -20Repository: massgen/MassGen
Length of output: 1932
Track streamed tool calls by item_id.
The Responses API can emit multiple function calls in a single response, with each concurrent call receiving deltas keyed by item_id. The current implementation uses a single current_function_call slot that gets overwritten when a second response.output_item.added arrives before the first call's arguments are finalized. The delta handler at lines 577-579 checks only current_function_call is not None without verifying the item_id matches, causing argument streams for earlier calls to accumulate into later calls' slots and resulting in corrupted or dropped tool calls. Use a dict keyed by item_id to track in-flight calls separately.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@massgen/backend/response.py` around lines 565 - 584, The code uses a single
current_function_call slot which gets overwritten when multiple concurrent
function calls stream; replace current_function_call with a dict (e.g.,
in_flight_calls) keyed by the unique item id/call id (use getattr(chunk.item,
"call_id") or chunk.item.item_id if present) and update the three handlers: in
response.output_item.added create an entry in in_flight_calls[item_id] with
name/arguments/call_id, in response.function_call_arguments.delta append deltas
to in_flight_calls[item_id]["arguments"] only when the chunk's item_id matches
an existing key, and in response.output_item.done move in_flight_calls[item_id]
into captured_function_calls and delete that key; ensure all references to
current_function_call are replaced accordingly so concurrent streams are tracked
independently.
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
1 participant
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.
Description
Add Web socket mode for resposnes api
Type of change
feat:) - Non-breaking change which adds functionalityChecklist
Summary by CodeRabbit
New Features
Configs
Bug Fixes
Tests
Chores