From 20b6a98d2ba544aa7c0b0727075978676b2b36d7 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Fri, 19 Sep 2025 12:14:20 -0700 Subject: [PATCH 01/19] WIP responses client Signed-off-by: George Armstrong --- nemo_skills/inference/model/responses.py | 343 +++++++++++++++++++++++ 1 file changed, 343 insertions(+) create mode 100644 nemo_skills/inference/model/responses.py diff --git a/nemo_skills/inference/model/responses.py b/nemo_skills/inference/model/responses.py new file mode 100644 index 0000000000..e4b82646b8 --- /dev/null +++ b/nemo_skills/inference/model/responses.py @@ -0,0 +1,343 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging + +from openai import AsyncOpenAI, OpenAI + +from nemo_skills.utils import get_logger_name + +from .base import BaseModel +from .context_retry import with_context_retry + +LOG = logging.getLogger(get_logger_name(__file__)) + + +class ResponsesModel(BaseModel): + """Model implementation using OpenAI Responses API via LiteLLM. + + This model uses the responses API endpoint instead of chat completions, + which is useful for models like gpt-oss that support the responses format. + """ + + MODEL_PROVIDER = "openai" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + # Initialize OpenAI clients for responses API + self.client = OpenAI(base_url=self.base_url, api_key=self.litellm_kwargs["api_key"]) + self.async_client = AsyncOpenAI(base_url=self.base_url, api_key=self.litellm_kwargs["api_key"]) + + def _build_completion_request_params( + self, + prompt: str, + tokens_to_generate: int = 512, + temperature: float = 0.0, + top_p: float = 0.95, + top_k: int = -1, + min_p: float = 0.0, + repetition_penalty: float = 1.0, + random_seed: int = None, + top_logprobs: int | None = None, + timeout: int | None = None, + stop_phrases: list[str] | None = None, + stream: bool = False, + reasoning_effort: str | None = None, + extra_body: dict = None, + tools: list[dict] | None = None, + ) -> dict: + """Build request parameters for responses API with string input.""" + request = { + "input": prompt, + "max_output_tokens": tokens_to_generate, + "temperature": temperature, + "top_p": top_p, + "stream": stream, + "tools": tools, + "extra_body": { + "seed": random_seed, + "reasoning_effort": reasoning_effort, + "timeout": timeout, + "stop": stop_phrases, + "top_logprobs": top_logprobs, + "top_k": top_k, + "min_p": min_p, + "repetition_penalty": repetition_penalty, + **(extra_body or {}), + }, + } + + return request + + def _build_chat_request_params( + self, + messages: list[dict], + stream: bool, + tokens_to_generate: int = 512, + temperature: float = 0.0, + top_p: float = 0.95, + top_k: int = -1, + min_p: float = 0.0, + repetition_penalty: float = 1.0, + random_seed: int = 0, + stop_phrases: list[str] | None = None, + timeout: int | None = None, + top_logprobs: int | None = None, + reasoning_effort: str | None = None, + tools: list[dict] | None = None, + extra_body: dict = None, + ) -> dict: + """Build request parameters for responses API with messages input.""" + request = { + "input": messages, + "max_output_tokens": tokens_to_generate, + "temperature": temperature, + "top_p": top_p, + "stream": stream, + "tools": tools, + "extra_body": { + "seed": random_seed, + "reasoning_effort": reasoning_effort, + "timeout": timeout, + "stop": stop_phrases, + "top_logprobs": top_logprobs, + "top_k": top_k, + "min_p": min_p, + "repetition_penalty": repetition_penalty, + **(extra_body or {}), + }, + } + + return request + + def _parse_responses_response(self, response, include_response: bool = False, **kwargs) -> dict: + """Parse responses API response into standard format.""" + + result = {"generation": "", "num_generated_tokens": 0} + + # Get token usage + if hasattr(response, "usage"): + result["num_generated_tokens"] = getattr(response.usage, "output_tokens", 0) + + # Check for tool calls in the output array + tool_calls = [] + reasoning_content = "" + generation_text = "" + + if hasattr(response, "output") and response.output: + for output_item in response.output: + # Handle reasoning content + if hasattr(output_item, "type") and output_item.type == "reasoning": + if hasattr(output_item, "content") and output_item.content: + for content_item in output_item.content: + if hasattr(content_item, "text"): + reasoning_content += content_item.text + "\n" + + # Handle function calls + elif hasattr(output_item, "type") and output_item.type == "function_call": + tool_calls.append(output_item) + + # Handle message content + elif hasattr(output_item, "type") and output_item.type == "message": + if hasattr(output_item, "content") and output_item.content: + for content_item in output_item.content: + if hasattr(content_item, "text"): + generation_text += content_item.text + + # Set the appropriate response fields + if tool_calls: + result["tool_calls"] = tool_calls + result["generation"] = "" # No text generation when there are tool calls + else: + result["generation"] = generation_text + + # Add reasoning content if available + if reasoning_content: + result["reasoning_content"] = reasoning_content.strip() + + # Add finish reason if available + if hasattr(response, "status"): + result["finish_reason"] = response.status + + if include_response: + result["response"] = response + + return result + + @with_context_retry + async def generate_async( + self, + prompt: str | list[dict], + tokens_to_generate: int | None = None, + temperature: float = 0.0, + top_p: float = 0.95, + top_k: int = -1, + min_p: float = 0.0, + repetition_penalty: float = 1.0, + random_seed: int = None, + stop_phrases: list[str] | None = None, + top_logprobs: int | None = None, + timeout: float | int | None = 14400, + remove_stop_phrases: bool = True, + stream: bool = False, + reasoning_effort: str | None = None, + tools: list[dict] | None = None, + include_response: bool = False, + extra_body: dict = None, + ) -> dict: + """Generate response using responses API.""" + + # Check tool calls are a list of dict + if tools is not None: + for tool in tools: + if not isinstance(tool, dict): + raise ValueError(f"Tool must be a dictionary, got {type(tool)}") + + kwargs = { + "tokens_to_generate": tokens_to_generate, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "repetition_penalty": repetition_penalty, + "random_seed": random_seed, + "stop_phrases": stop_phrases, + "top_logprobs": top_logprobs, + "timeout": timeout, + "reasoning_effort": reasoning_effort, + "tools": tools, + "extra_body": extra_body, + } + + request_params = self._build_request_params(prompt=prompt, stream=stream, **kwargs) + + # Use OpenAI client for responses API + LOG.info(f"Making responses API call with params: {request_params}") + LOG.info(f"Full litellm_kwargs: {self.litellm_kwargs}") + LOG.info(f"Model name from litellm_kwargs: {self.litellm_kwargs['model']}") + LOG.info(f"Base URL: {self.base_url}") + + # Use the full model name as vLLM expects it + model_name = self.litellm_kwargs["model"] # Should be openai/gpt-oss-20b + LOG.info(f"About to call responses.create with model: {model_name}") + + response = await self.async_client.responses.create(model=model_name, **request_params) + + LOG.info(f"Response from server: {response}") + LOG.info(f"Response type: {type(response)}") + + if stream: + # Handle streaming responses + return self._stream_responses_chunks_async(response) + else: + result = self._parse_responses_response(response, include_response=include_response, **kwargs) + self._maybe_apply_stop_phrase_removal(result, remove_stop_phrases, stop_phrases) + return result + + @with_context_retry + def generate_sync( + self, + prompt: str | list[dict], + tokens_to_generate: int | None = None, + temperature: float = 0.0, + top_p: float = 0.95, + top_k: int = -1, + min_p: float = 0.0, + repetition_penalty: float = 1.0, + random_seed: int = None, + stop_phrases: list[str] | None = None, + top_logprobs: int | None = None, + timeout: float | int | None = 14400, + remove_stop_phrases: bool = True, + stream: bool = False, + reasoning_effort: str | None = None, + tools: list[dict] | None = None, + include_response: bool = False, + extra_body: dict = None, + ) -> dict: + """Synchronous version of generate using responses API.""" + + # Check tool calls are a list of dict + if tools is not None: + for tool in tools: + if not isinstance(tool, dict): + raise ValueError(f"Tool must be a dictionary, got {type(tool)}") + + kwargs = { + "tokens_to_generate": tokens_to_generate, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "repetition_penalty": repetition_penalty, + "random_seed": random_seed, + "stop_phrases": stop_phrases, + "top_logprobs": top_logprobs, + "timeout": timeout, + "reasoning_effort": reasoning_effort, + "tools": tools, + "extra_body": extra_body, + } + + request_params = self._build_request_params(prompt=prompt, stream=stream, **kwargs) + + # Use OpenAI client for responses API + LOG.info(f"Making sync responses API call with params: {request_params}") + LOG.info(f"Model name: {self.litellm_kwargs['model']}") + + # Use the full model name as vLLM expects it + model_name = self.litellm_kwargs["model"] # Keep openai/gpt-oss-20b + + response = self.client.responses.create(model=model_name, **request_params) + + LOG.info(f"Response from server: {response}") + LOG.info(f"Response type: {type(response)}") + + if stream: + # Handle streaming responses + return self._stream_responses_chunks_sync(response) + else: + result = self._parse_responses_response(response, include_response=include_response, **kwargs) + self._maybe_apply_stop_phrase_removal(result, remove_stop_phrases, stop_phrases) + return result + + def _stream_responses_chunks_sync(self, response): + """Synchronous version of stream responses chunks.""" + for chunk in response: + result = self._process_responses_chunk(chunk) + if result: + yield result + + async def _stream_responses_chunks_async(self, response): + """Async version of stream responses chunks.""" + async for chunk in response: + result = self._process_responses_chunk(chunk) + if result: + yield result + + def _process_responses_chunk(self, chunk): + """Process a single responses API chunk.""" + # This will depend on the actual streaming format from responses API + # For now, implement basic text streaming + if hasattr(chunk, "output_text"): + return {"generation": chunk.output_text or ""} + elif hasattr(chunk, "output") and chunk.output: + if isinstance(chunk.output, list) and len(chunk.output) > 0: + first_output = chunk.output[0] + if hasattr(first_output, "text"): + return {"generation": first_output.text or ""} + + # Fallback + return {"generation": ""} From 6c9760a61eb4a6be2a0e258f45a12588bb72bc59 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Fri, 19 Sep 2025 13:21:26 -0700 Subject: [PATCH 02/19] WIP make python tool http runnable Signed-off-by: George Armstrong --- nemo_skills/mcp/servers/python_tool.py | 97 +++++++++++++++++++++++--- 1 file changed, 89 insertions(+), 8 deletions(-) diff --git a/nemo_skills/mcp/servers/python_tool.py b/nemo_skills/mcp/servers/python_tool.py index 140b7bd206..853e3c3b6f 100644 --- a/nemo_skills/mcp/servers/python_tool.py +++ b/nemo_skills/mcp/servers/python_tool.py @@ -16,13 +16,23 @@ import logging from collections import defaultdict from dataclasses import dataclass -from typing import Annotated, Any, Dict +from typing import Any, Dict + +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated from httpx import RemoteProtocolError from mcp.server.fastmcp import FastMCP from omegaconf import OmegaConf from pydantic import Field +try: + import uvicorn +except ImportError: + uvicorn = None + from nemo_skills.code_execution.sandbox import get_sandbox from nemo_skills.mcp.tool_providers import MCPClientTool from nemo_skills.mcp.utils import add_config_args, load_mcp_config @@ -69,6 +79,17 @@ async def stateful_python_code_exec( def main(): parser = argparse.ArgumentParser(description="MCP server for executing Python code in a sandbox") add_config_args(parser) + parser.add_argument( + "--transport", + type=str, + choices=["stdio", "streamable-http"], + default="stdio", + help="Transport mode for the MCP server (stdio or streamable-http)", + ) + parser.add_argument( + "--host", type=str, default="127.0.0.1", help="Host to bind to for HTTP transport (default: 127.0.0.1)" + ) + parser.add_argument("--port", type=int, default=8000, help="Port to bind to for HTTP transport (default: 8000)") args = parser.parse_args() try: @@ -81,11 +102,44 @@ def main(): logger.warning(f"{e} Falling back to default local sandbox config.") cfg = OmegaConf.create({"sandbox": {"sandbox_type": "local"}}) + # Start with transport config from file + transport_config = getattr(cfg, "transport", {}) + + # Apply command-line overrides only for non-default values or when config specifies HTTP + if args.transport != "stdio": + # Explicit command-line transport override + transport_config.update({"type": args.transport, "host": args.host, "port": args.port}) + elif transport_config.get("type") == "streamable-http": + # Config file specifies HTTP, only override host/port if provided + if args.host != "127.0.0.1": + transport_config["host"] = args.host + if args.port != 8000: + transport_config["port"] = args.port + else: + # Default to stdio + transport_config.setdefault("type", "stdio") + global sandbox sandbox_cfg = OmegaConf.to_container(cfg.sandbox, resolve=True) sandbox = get_sandbox(**sandbox_cfg) - # Initialize and run the server - mcp.run(transport="stdio") + + # Initialize and run the server based on transport configuration + transport_type = transport_config.get("type", "stdio") + + if transport_type == "streamable-http": + host = transport_config.get("host", "127.0.0.1") + port = transport_config.get("port", 8000) + logger.info(f"Starting Python tool server in streamable HTTP mode on {host}:{port}") + + if uvicorn is None: + raise ImportError("uvicorn not available. Please install uvicorn to use HTTP transport.") + + # Get the Starlette app from FastMCP and run it with uvicorn + app = mcp.streamable_http_app() + uvicorn.run(app, host=host, port=port) + else: + logger.info("Starting Python tool server in stdio mode") + mcp.run(transport="stdio") # ============================== @@ -94,11 +148,24 @@ def main(): class PythonTool(MCPClientTool): - def __init__(self) -> None: + def __init__(self, transport_type: str = "stdio", host: str = "127.0.0.1", port: int = 8000) -> None: super().__init__() - # Defaults for stdio Python MCP using explicit client class - self.apply_config_updates( - { + + if transport_type == "streamable-http": + # Configuration for streamable HTTP client + base_config = { + "client": "nemo_skills.mcp.clients.MCPStreamableHttpClient", + "client_params": { + "base_url": f"http://{host}:{port}", + }, + # hide args from schemas and sanitize at runtime + "hide_args": {"stateful_python_code_exec": ["session_id", "timeout"]}, + # execution-specific default + "exec_timeout_s": 10, + } + else: + # Default stdio configuration + base_config = { "client": "nemo_skills.mcp.clients.MCPStdioClient", "client_params": { "command": "python", @@ -111,9 +178,23 @@ def __init__(self) -> None: # execution-specific default "exec_timeout_s": 10, } - ) + + self.apply_config_updates(base_config) self.requests_to_sessions = defaultdict(lambda: None) + @classmethod + def create_from_config(cls, config: Dict[str, Any] | None = None): + """Factory method to create PythonTool from configuration.""" + if not config: + return cls() + + transport_config = config.get("transport", {}) + transport_type = transport_config.get("type", "stdio") + host = transport_config.get("host", "127.0.0.1") + port = transport_config.get("port", 8000) + + return cls(transport_type=transport_type, host=host, port=port) + async def execute(self, tool_name: str, arguments: Dict[str, Any], extra_args: Dict[str, Any] | None = None): # Ensure timeout is sent via extra_args (post-sanitize), not in main arguments arguments = dict(arguments) From 06dd03c81b2f059f9dd682f0526961d47f3a59b6 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Mon, 22 Sep 2025 15:53:54 -0700 Subject: [PATCH 03/19] ENH add see protocol to python tool Signed-off-by: George Armstrong --- nemo_skills/mcp/servers/python_tool.py | 37 +++++++++++++++++++++----- 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/nemo_skills/mcp/servers/python_tool.py b/nemo_skills/mcp/servers/python_tool.py index 853e3c3b6f..5cd7e1aac0 100644 --- a/nemo_skills/mcp/servers/python_tool.py +++ b/nemo_skills/mcp/servers/python_tool.py @@ -82,9 +82,9 @@ def main(): parser.add_argument( "--transport", type=str, - choices=["stdio", "streamable-http"], + choices=["stdio", "streamable-http", "sse"], default="stdio", - help="Transport mode for the MCP server (stdio or streamable-http)", + help="Transport mode for the MCP server (stdio, streamable-http, or sse)", ) parser.add_argument( "--host", type=str, default="127.0.0.1", help="Host to bind to for HTTP transport (default: 127.0.0.1)" @@ -105,12 +105,12 @@ def main(): # Start with transport config from file transport_config = getattr(cfg, "transport", {}) - # Apply command-line overrides only for non-default values or when config specifies HTTP + # Apply command-line overrides only for non-default values or when config specifies HTTP/SSE if args.transport != "stdio": # Explicit command-line transport override transport_config.update({"type": args.transport, "host": args.host, "port": args.port}) - elif transport_config.get("type") == "streamable-http": - # Config file specifies HTTP, only override host/port if provided + elif transport_config.get("type") in ["streamable-http", "sse"]: + # Config file specifies HTTP or SSE, only override host/port if provided if args.host != "127.0.0.1": transport_config["host"] = args.host if args.port != 8000: @@ -137,6 +137,17 @@ def main(): # Get the Starlette app from FastMCP and run it with uvicorn app = mcp.streamable_http_app() uvicorn.run(app, host=host, port=port) + elif transport_type == "sse": + host = transport_config.get("host", "127.0.0.1") + port = transport_config.get("port", 8000) + logger.info(f"Starting Python tool server in SSE mode on {host}:{port}") + + if uvicorn is None: + raise ImportError("uvicorn not available. Please install uvicorn to use SSE transport.") + + # Get the SSE app from FastMCP and run it with uvicorn + app = mcp.sse_app() + uvicorn.run(app, host=host, port=port) else: logger.info("Starting Python tool server in stdio mode") mcp.run(transport="stdio") @@ -163,7 +174,19 @@ def __init__(self, transport_type: str = "stdio", host: str = "127.0.0.1", port: # execution-specific default "exec_timeout_s": 10, } - else: + # elif transport_type == "sse": + # # Configuration for SSE client (using HTTP client for SSE endpoints) + # base_config = { + # "client": "nemo_skills.mcp.clients.MCPHttpClient", + # "client_params": { + # "base_url": f"http://{host}:{port}", + # }, + # # hide args from schemas and sanitize at runtime + # "hide_args": {"stateful_python_code_exec": ["session_id", "timeout"]}, + # # execution-specific default + # "exec_timeout_s": 10, + # } + elif transport_type == "stdio": # Default stdio configuration base_config = { "client": "nemo_skills.mcp.clients.MCPStdioClient", @@ -178,6 +201,8 @@ def __init__(self, transport_type: str = "stdio", host: str = "127.0.0.1", port: # execution-specific default "exec_timeout_s": 10, } + else: + raise ValueError(f"Invalid transport type: {transport_type}") self.apply_config_updates(base_config) self.requests_to_sessions = defaultdict(lambda: None) From 04d15ef2139ab2fceaf83ee3be4cc11b01ed59be Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Mon, 22 Sep 2025 17:35:19 -0700 Subject: [PATCH 04/19] ENH enable response API for toolwrapper Signed-off-by: George Armstrong --- nemo_skills/inference/model/responses.py | 44 ++++++- nemo_skills/inference/model/tool_call.py | 71 +++++++++-- nemo_skills/mcp/adapters.py | 150 ++++++++++++++++++++++- 3 files changed, 245 insertions(+), 20 deletions(-) diff --git a/nemo_skills/inference/model/responses.py b/nemo_skills/inference/model/responses.py index e4b82646b8..93350e125e 100644 --- a/nemo_skills/inference/model/responses.py +++ b/nemo_skills/inference/model/responses.py @@ -65,7 +65,6 @@ def _build_completion_request_params( "temperature": temperature, "top_p": top_p, "stream": stream, - "tools": tools, "extra_body": { "seed": random_seed, "reasoning_effort": reasoning_effort, @@ -79,6 +78,10 @@ def _build_completion_request_params( }, } + # Only include tools if they are provided + if tools is not None: + request["tools"] = tools + return request def _build_chat_request_params( @@ -106,7 +109,6 @@ def _build_chat_request_params( "temperature": temperature, "top_p": top_p, "stream": stream, - "tools": tools, "extra_body": { "seed": random_seed, "reasoning_effort": reasoning_effort, @@ -120,6 +122,10 @@ def _build_chat_request_params( }, } + # Only include tools if they are provided + if tools is not None: + request["tools"] = tools + return request def _parse_responses_response(self, response, include_response: bool = False, **kwargs) -> dict: @@ -171,11 +177,37 @@ def _parse_responses_response(self, response, include_response: bool = False, ** if hasattr(response, "status"): result["finish_reason"] = response.status + # Add serialized output for conversation history + result["serialized_output"] = self._serialize_response_output(response) + if include_response: result["response"] = response return result + def _serialize_response_output(self, response) -> list[dict]: + """Serialize response output objects using model_dump() for conversation history.""" + serialized_output = [] + + if hasattr(response, "output") and response.output: + for output_item in response.output: + try: + # Try to use model_dump() method if available (Pydantic models) + if hasattr(output_item, "model_dump"): + serialized_output.append(output_item.model_dump()) + # Fallback to dict conversion + elif hasattr(output_item, "__dict__"): + serialized_output.append(output_item.__dict__) + # Last resort: convert to string representation + else: + serialized_output.append({"content": str(output_item), "type": "unknown"}) + except Exception as e: + LOG.warning(f"Failed to serialize output item: {e}") + # Fallback serialization + serialized_output.append({"content": str(output_item), "type": "error", "error": str(e)}) + + return serialized_output + @with_context_retry async def generate_async( self, @@ -229,8 +261,8 @@ async def generate_async( LOG.info(f"Model name from litellm_kwargs: {self.litellm_kwargs['model']}") LOG.info(f"Base URL: {self.base_url}") - # Use the full model name as vLLM expects it - model_name = self.litellm_kwargs["model"] # Should be openai/gpt-oss-20b + # Use the original model name (without litellm prefix) for OpenAI client + model_name = self.model_name_or_path # Just gpt-oss-20b LOG.info(f"About to call responses.create with model: {model_name}") response = await self.async_client.responses.create(model=model_name, **request_params) @@ -297,8 +329,8 @@ def generate_sync( LOG.info(f"Making sync responses API call with params: {request_params}") LOG.info(f"Model name: {self.litellm_kwargs['model']}") - # Use the full model name as vLLM expects it - model_name = self.litellm_kwargs["model"] # Keep openai/gpt-oss-20b + # Use the original model name (without litellm prefix) for OpenAI client + model_name = self.model_name_or_path # Just gpt-oss-20b response = self.client.responses.create(model=model_name, **request_params) diff --git a/nemo_skills/inference/model/tool_call.py b/nemo_skills/inference/model/tool_call.py index bba2f64f9c..b11e80793d 100644 --- a/nemo_skills/inference/model/tool_call.py +++ b/nemo_skills/inference/model/tool_call.py @@ -23,7 +23,12 @@ from nemo_skills.mcp.adapters import ( ChatCompletionCallInterpreter, ChatCompletionResponseFormatter, + CompletionConversationManager, OpenAISchemaAdapter, + ResponsesCallInterpreter, + ResponsesConversationManager, + ResponsesResponseFormatter, + ResponsesSchemaAdapter, ) from nemo_skills.mcp.tool_manager import ToolManager from nemo_skills.utils import get_logger_name @@ -59,10 +64,28 @@ def __init__( overrides=tool_overrides or {}, context=additional_config, ) - # Use sensible defaults for adapters in module-based mode - self.schema_adapter = OpenAISchemaAdapter() - self.call_interpreter = ChatCompletionCallInterpreter() - self.response_formatter = ChatCompletionResponseFormatter() + + # Detect model type and set up appropriate adapters + self._setup_adapters() + + def _setup_adapters(self): + """Set up adapters based on model type.""" + # Import here to avoid circular imports + from .responses import ResponsesModel + + # Detect model type and configure adapters + if isinstance(self.model, ResponsesModel): + # Responses API model - uses flatter tool schema format + self.schema_adapter = ResponsesSchemaAdapter() + self.call_interpreter = ResponsesCallInterpreter() + self.response_formatter = ResponsesResponseFormatter() + self.conversation_manager = ResponsesConversationManager() + else: + # Chat completion model (default) - uses nested function format + self.schema_adapter = OpenAISchemaAdapter() + self.call_interpreter = ChatCompletionCallInterpreter() + self.response_formatter = ChatCompletionResponseFormatter() + self.conversation_manager = CompletionConversationManager() async def _execute_tool_call(self, tool_call, request_id: str): ## TODO(sanyamk): The correct key format needs to be cohesive with other formatters. @@ -103,7 +126,7 @@ async def generate_async( tokens_to_generate: int = None, **generation_kwargs, ) -> Dict: - assert isinstance(prompt, list), "Only use ChatCompletion API for now." + assert isinstance(prompt, list), "Prompt must be a list for tool calling." assert tools is None, "Do not pass 'tools'; they are derived from tool_modules." @@ -133,19 +156,41 @@ async def generate_async( if k in generation: result_steps[k].append(generation[k]) - conversation.append({"role": "assistant", "content": generation["generation"]}) - if "reasoning_content" in generation: - conversation[-1]["reasoning_content"] = generation["reasoning_content"] + # Use conversation manager to add assistant response + self.conversation_manager.add_assistant_response(conversation, generation) + # Check for tool calls (simple and direct) tool_calls = generation.get("tool_calls", []) if tool_calls: tool_calls_message = self.call_interpreter.parse(tool_calls) - conversation[-1].update(tool_calls_message) - tool_calls_output_messages = await self._execute_tool_calls( - tool_calls_message["tool_calls"], request_id=request_id - ) - conversation.extend(tool_calls_output_messages) + # Update the last message with tool calls (for completion models) + if ( + hasattr(self.conversation_manager, "__class__") + and "Completion" in self.conversation_manager.__class__.__name__ + ): + if conversation and conversation[-1].get("role") == "assistant": + conversation[-1].update(tool_calls_message) + + tool_results = await self._execute_tool_calls(tool_calls_message["tool_calls"], request_id=request_id) + + # For completion models, the tool results are already formatted messages + # For responses models, we need to extract the actual results + if ( + hasattr(self.conversation_manager, "__class__") + and "Responses" in self.conversation_manager.__class__.__name__ + ): + # Extract raw results for responses conversation manager + raw_results = [] + for result_msg in tool_results: + # ResponsesResponseFormatter creates {"type": "function_call_output", "call_id": "...", "output": "..."} + raw_results.append(result_msg.get("output", "")) + self.conversation_manager.add_tool_results( + conversation, tool_calls_message["tool_calls"], raw_results + ) + else: + # For completion models, add the formatted messages directly + conversation.extend(tool_results) result_steps["num_tool_calls"].append(len(tool_calls)) diff --git a/nemo_skills/mcp/adapters.py b/nemo_skills/mcp/adapters.py index 6e385d9757..3ee282f2b7 100644 --- a/nemo_skills/mcp/adapters.py +++ b/nemo_skills/mcp/adapters.py @@ -15,7 +15,7 @@ import json from abc import ABC, abstractmethod -from typing import List +from typing import Any, Dict, List from litellm.types.utils import ChatCompletionMessageToolCall @@ -43,6 +43,27 @@ def format(self, tool_call: ChatCompletionMessageToolCall, result: dict) -> dict raise NotImplementedError("Subclasses must implement this method.") +class ConversationManager(ABC): + """Manages conversation history building for different model types.""" + + @abstractmethod + def add_user_message(self, conversation: List[Dict[str, Any]], content: str) -> None: + """Add a user message to the conversation.""" + raise NotImplementedError("Subclasses must implement this method.") + + @abstractmethod + def add_assistant_response(self, conversation: List[Dict[str, Any]], response: Dict[str, Any]) -> None: + """Add an assistant response to the conversation.""" + raise NotImplementedError("Subclasses must implement this method.") + + @abstractmethod + def add_tool_results( + self, conversation: List[Dict[str, Any]], tool_calls: List[Dict[str, Any]], results: List[Dict[str, Any]] + ) -> None: + """Add tool call results to the conversation.""" + raise NotImplementedError("Subclasses must implement this method.") + + # ============================== # ADAPTER IMPLEMENTATIONS # ============================== @@ -116,3 +137,130 @@ def format(self, tool_call, result): "tool_call_id": tool_call["id"], "content": json.dumps(result) if not isinstance(result, str) else result, } + + +# ============================== +# CONVERSATION MANAGERS +# ============================== + + +class CompletionConversationManager(ConversationManager): + """Manages conversation history for chat completion models.""" + + def add_user_message(self, conversation: List[Dict[str, Any]], content: str) -> None: + """Add a user message to the conversation.""" + conversation.append({"role": "user", "content": content}) + + def add_assistant_response(self, conversation: List[Dict[str, Any]], response: Dict[str, Any]) -> None: + """Add an assistant response to the conversation.""" + message = {"role": "assistant", "content": response["generation"]} + + # Add reasoning content if available + if "reasoning_content" in response: + message["reasoning_content"] = response["reasoning_content"] + + conversation.append(message) + + def add_tool_results( + self, conversation: List[Dict[str, Any]], tool_calls: List[Dict[str, Any]], results: List[Dict[str, Any]] + ) -> None: + """Add tool call results to the conversation.""" + # Update the last assistant message with tool calls + if conversation and conversation[-1]["role"] == "assistant": + conversation[-1]["tool_calls"] = tool_calls + + # Add tool result messages + for tool_call, result in zip(tool_calls, results): + conversation.append( + { + "role": "tool", + "name": tool_call["function"]["name"], + "tool_call_id": tool_call["id"], + "content": json.dumps(result) if not isinstance(result, str) else result, + } + ) + + +class ResponsesConversationManager(ConversationManager): + """Manages conversation history for responses API models.""" + + def add_user_message(self, conversation: List[Dict[str, Any]], content: str) -> None: + """Add a user message to the conversation.""" + conversation.append({"role": "user", "content": content}) + + def add_assistant_response(self, conversation: List[Dict[str, Any]], response: Dict[str, Any]) -> None: + """Add an assistant response to the conversation using serialized output.""" + # Use the serialized output from the responses API + if "serialized_output" in response: + conversation.extend(response["serialized_output"]) + else: + # Fallback to basic message format + conversation.append({"role": "assistant", "content": response["generation"]}) + + def add_tool_results( + self, conversation: List[Dict[str, Any]], tool_calls: List[Dict[str, Any]], results: List[Dict[str, Any]] + ) -> None: + """Add tool call results to the conversation.""" + # For responses API, add tool results as function_call_output items + for tool_call, result in zip(tool_calls, results): + conversation.append( + { + "type": "function_call_output", + "call_id": tool_call.get("id", "unknown"), + "output": json.dumps(result) if not isinstance(result, str) else result, + } + ) + + +# ============================== +# RESPONSES API ADAPTERS +# ============================== + + +class ResponsesCallInterpreter(ToolCallInterpreter): + """Convert responses API tool calls to a standardized format.""" + + def parse(self, tool_calls: List[Any]) -> Dict[str, Any]: + """Parse tool calls from responses API format.""" + parsed_calls = [] + + for tool_call in tool_calls: + parsed_call = { + "type": "function", + "id": getattr(tool_call, "call_id", getattr(tool_call, "id", "unknown")), + "function": { + "name": getattr(tool_call, "name", "unknown"), + "arguments": getattr(tool_call, "arguments", "{}"), + }, + } + parsed_calls.append(parsed_call) + + return {"role": "assistant", "tool_calls": parsed_calls} + + +class ResponsesResponseFormatter(ToolResponseFormatter): + """Format tool call results for responses API conversation history.""" + + def format(self, tool_call: Dict[str, Any], result: Dict[str, Any]) -> Dict[str, Any]: + """Format the response from a tool call for responses API.""" + return { + "type": "function_call_output", + "call_id": tool_call["id"], + "output": json.dumps(result) if not isinstance(result, str) else result, + } + + +class ResponsesSchemaAdapter(ToolSchemaAdapter): + """Convert MCP tool definitions to responses API format (flatter structure).""" + + def convert(self, tools): + """Convert tools to responses API format without nested 'function' object.""" + return [ + { + "type": "function", + "name": t["name"], + "description": t["description"], + "parameters": t["input_schema"], + } + for t in tools + ] From 2ef268330dffe7940f856cfe369224c108271108 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Wed, 24 Sep 2025 11:22:29 -0700 Subject: [PATCH 05/19] ENH support responses model type Signed-off-by: George Armstrong --- nemo_skills/inference/model/__init__.py | 2 ++ nemo_skills/pipeline/utils/server.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/nemo_skills/inference/model/__init__.py b/nemo_skills/inference/model/__init__.py index 6ecfaa6e25..732978e338 100644 --- a/nemo_skills/inference/model/__init__.py +++ b/nemo_skills/inference/model/__init__.py @@ -28,6 +28,7 @@ from .megatron import MegatronModel from .openai import OpenAIModel from .parallel_thinking import ParallelThinkingConfig, ParallelThinkingTask +from .responses import ResponsesModel # Tool Calling from .tool_call import ToolCallingWrapper @@ -47,6 +48,7 @@ "gemini": GeminiModel, "vllm": VLLMModel, "sglang": VLLMModel, + "responses": ResponsesModel, } diff --git a/nemo_skills/pipeline/utils/server.py b/nemo_skills/pipeline/utils/server.py index f3b9b3ea55..0946b0bd98 100644 --- a/nemo_skills/pipeline/utils/server.py +++ b/nemo_skills/pipeline/utils/server.py @@ -123,7 +123,7 @@ def get_server_command( # check if the model path is mounted if not vllm, sglang, or trtllm; # vllm, sglang, and trtllm can also pass model name as "model_path" so we need special processing - if server_type not in ["vllm", "sglang", "trtllm"]: + if server_type not in ["vllm", "sglang", "trtllm", "responses"]: check_if_mounted(cluster_config, model_path) # the model path will be mounted, so generally it will start with / @@ -152,7 +152,7 @@ def get_server_command( f" --micro-batch-size 1 " # that's a training argument, ignored here, but required to specify.. f" {server_args} " ) - elif server_type == "vllm": + elif server_type in ["vllm", "responses"]: server_entrypoint = server_entrypoint or "-m nemo_skills.inference.server.serve_vllm" start_vllm_cmd = ( f"python3 {server_entrypoint} " From 70aac38dcb19067b18a9a3190915b02089fee26b Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Wed, 24 Sep 2025 11:54:56 -0700 Subject: [PATCH 06/19] MAINT revert unrelated python tool changes Signed-off-by: George Armstrong --- nemo_skills/mcp/servers/python_tool.py | 122 ++----------------------- 1 file changed, 8 insertions(+), 114 deletions(-) diff --git a/nemo_skills/mcp/servers/python_tool.py b/nemo_skills/mcp/servers/python_tool.py index 5cd7e1aac0..140b7bd206 100644 --- a/nemo_skills/mcp/servers/python_tool.py +++ b/nemo_skills/mcp/servers/python_tool.py @@ -16,23 +16,13 @@ import logging from collections import defaultdict from dataclasses import dataclass -from typing import Any, Dict - -try: - from typing import Annotated -except ImportError: - from typing_extensions import Annotated +from typing import Annotated, Any, Dict from httpx import RemoteProtocolError from mcp.server.fastmcp import FastMCP from omegaconf import OmegaConf from pydantic import Field -try: - import uvicorn -except ImportError: - uvicorn = None - from nemo_skills.code_execution.sandbox import get_sandbox from nemo_skills.mcp.tool_providers import MCPClientTool from nemo_skills.mcp.utils import add_config_args, load_mcp_config @@ -79,17 +69,6 @@ async def stateful_python_code_exec( def main(): parser = argparse.ArgumentParser(description="MCP server for executing Python code in a sandbox") add_config_args(parser) - parser.add_argument( - "--transport", - type=str, - choices=["stdio", "streamable-http", "sse"], - default="stdio", - help="Transport mode for the MCP server (stdio, streamable-http, or sse)", - ) - parser.add_argument( - "--host", type=str, default="127.0.0.1", help="Host to bind to for HTTP transport (default: 127.0.0.1)" - ) - parser.add_argument("--port", type=int, default=8000, help="Port to bind to for HTTP transport (default: 8000)") args = parser.parse_args() try: @@ -102,55 +81,11 @@ def main(): logger.warning(f"{e} Falling back to default local sandbox config.") cfg = OmegaConf.create({"sandbox": {"sandbox_type": "local"}}) - # Start with transport config from file - transport_config = getattr(cfg, "transport", {}) - - # Apply command-line overrides only for non-default values or when config specifies HTTP/SSE - if args.transport != "stdio": - # Explicit command-line transport override - transport_config.update({"type": args.transport, "host": args.host, "port": args.port}) - elif transport_config.get("type") in ["streamable-http", "sse"]: - # Config file specifies HTTP or SSE, only override host/port if provided - if args.host != "127.0.0.1": - transport_config["host"] = args.host - if args.port != 8000: - transport_config["port"] = args.port - else: - # Default to stdio - transport_config.setdefault("type", "stdio") - global sandbox sandbox_cfg = OmegaConf.to_container(cfg.sandbox, resolve=True) sandbox = get_sandbox(**sandbox_cfg) - - # Initialize and run the server based on transport configuration - transport_type = transport_config.get("type", "stdio") - - if transport_type == "streamable-http": - host = transport_config.get("host", "127.0.0.1") - port = transport_config.get("port", 8000) - logger.info(f"Starting Python tool server in streamable HTTP mode on {host}:{port}") - - if uvicorn is None: - raise ImportError("uvicorn not available. Please install uvicorn to use HTTP transport.") - - # Get the Starlette app from FastMCP and run it with uvicorn - app = mcp.streamable_http_app() - uvicorn.run(app, host=host, port=port) - elif transport_type == "sse": - host = transport_config.get("host", "127.0.0.1") - port = transport_config.get("port", 8000) - logger.info(f"Starting Python tool server in SSE mode on {host}:{port}") - - if uvicorn is None: - raise ImportError("uvicorn not available. Please install uvicorn to use SSE transport.") - - # Get the SSE app from FastMCP and run it with uvicorn - app = mcp.sse_app() - uvicorn.run(app, host=host, port=port) - else: - logger.info("Starting Python tool server in stdio mode") - mcp.run(transport="stdio") + # Initialize and run the server + mcp.run(transport="stdio") # ============================== @@ -159,36 +94,11 @@ def main(): class PythonTool(MCPClientTool): - def __init__(self, transport_type: str = "stdio", host: str = "127.0.0.1", port: int = 8000) -> None: + def __init__(self) -> None: super().__init__() - - if transport_type == "streamable-http": - # Configuration for streamable HTTP client - base_config = { - "client": "nemo_skills.mcp.clients.MCPStreamableHttpClient", - "client_params": { - "base_url": f"http://{host}:{port}", - }, - # hide args from schemas and sanitize at runtime - "hide_args": {"stateful_python_code_exec": ["session_id", "timeout"]}, - # execution-specific default - "exec_timeout_s": 10, - } - # elif transport_type == "sse": - # # Configuration for SSE client (using HTTP client for SSE endpoints) - # base_config = { - # "client": "nemo_skills.mcp.clients.MCPHttpClient", - # "client_params": { - # "base_url": f"http://{host}:{port}", - # }, - # # hide args from schemas and sanitize at runtime - # "hide_args": {"stateful_python_code_exec": ["session_id", "timeout"]}, - # # execution-specific default - # "exec_timeout_s": 10, - # } - elif transport_type == "stdio": - # Default stdio configuration - base_config = { + # Defaults for stdio Python MCP using explicit client class + self.apply_config_updates( + { "client": "nemo_skills.mcp.clients.MCPStdioClient", "client_params": { "command": "python", @@ -201,25 +111,9 @@ def __init__(self, transport_type: str = "stdio", host: str = "127.0.0.1", port: # execution-specific default "exec_timeout_s": 10, } - else: - raise ValueError(f"Invalid transport type: {transport_type}") - - self.apply_config_updates(base_config) + ) self.requests_to_sessions = defaultdict(lambda: None) - @classmethod - def create_from_config(cls, config: Dict[str, Any] | None = None): - """Factory method to create PythonTool from configuration.""" - if not config: - return cls() - - transport_config = config.get("transport", {}) - transport_type = transport_config.get("type", "stdio") - host = transport_config.get("host", "127.0.0.1") - port = transport_config.get("port", 8000) - - return cls(transport_type=transport_type, host=host, port=port) - async def execute(self, tool_name: str, arguments: Dict[str, Any], extra_args: Dict[str, Any] | None = None): # Ensure timeout is sent via extra_args (post-sanitize), not in main arguments arguments = dict(arguments) From 2eb538f7f7cb381a072b48db6ff825938a9205ea Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Wed, 24 Sep 2025 14:02:53 -0700 Subject: [PATCH 07/19] MAINT cleanup adapters Signed-off-by: George Armstrong --- nemo_skills/inference/model/tool_call.py | 27 ++------------ nemo_skills/mcp/adapters.py | 45 +++++++++++------------- 2 files changed, 23 insertions(+), 49 deletions(-) diff --git a/nemo_skills/inference/model/tool_call.py b/nemo_skills/inference/model/tool_call.py index b11e80793d..a2135d05d0 100644 --- a/nemo_skills/inference/model/tool_call.py +++ b/nemo_skills/inference/model/tool_call.py @@ -164,33 +164,10 @@ async def generate_async( if tool_calls: tool_calls_message = self.call_interpreter.parse(tool_calls) - # Update the last message with tool calls (for completion models) - if ( - hasattr(self.conversation_manager, "__class__") - and "Completion" in self.conversation_manager.__class__.__name__ - ): - if conversation and conversation[-1].get("role") == "assistant": - conversation[-1].update(tool_calls_message) - tool_results = await self._execute_tool_calls(tool_calls_message["tool_calls"], request_id=request_id) - # For completion models, the tool results are already formatted messages - # For responses models, we need to extract the actual results - if ( - hasattr(self.conversation_manager, "__class__") - and "Responses" in self.conversation_manager.__class__.__name__ - ): - # Extract raw results for responses conversation manager - raw_results = [] - for result_msg in tool_results: - # ResponsesResponseFormatter creates {"type": "function_call_output", "call_id": "...", "output": "..."} - raw_results.append(result_msg.get("output", "")) - self.conversation_manager.add_tool_results( - conversation, tool_calls_message["tool_calls"], raw_results - ) - else: - # For completion models, add the formatted messages directly - conversation.extend(tool_results) + # Use conversation manager for both model types - it handles all the details + self.conversation_manager.add_tool_results(conversation, tool_calls_message, tool_results) result_steps["num_tool_calls"].append(len(tool_calls)) diff --git a/nemo_skills/mcp/adapters.py b/nemo_skills/mcp/adapters.py index 3ee282f2b7..2b26dfc46e 100644 --- a/nemo_skills/mcp/adapters.py +++ b/nemo_skills/mcp/adapters.py @@ -58,7 +58,10 @@ def add_assistant_response(self, conversation: List[Dict[str, Any]], response: D @abstractmethod def add_tool_results( - self, conversation: List[Dict[str, Any]], tool_calls: List[Dict[str, Any]], results: List[Dict[str, Any]] + self, + conversation: List[Dict[str, Any]], + tool_calls_message: Dict[str, Any], + formatted_results: List[Dict[str, Any]], ) -> None: """Add tool call results to the conversation.""" raise NotImplementedError("Subclasses must implement this method.") @@ -162,23 +165,18 @@ def add_assistant_response(self, conversation: List[Dict[str, Any]], response: D conversation.append(message) def add_tool_results( - self, conversation: List[Dict[str, Any]], tool_calls: List[Dict[str, Any]], results: List[Dict[str, Any]] + self, + conversation: List[Dict[str, Any]], + tool_calls_message: Dict[str, Any], + formatted_results: List[Dict[str, Any]], ) -> None: """Add tool call results to the conversation.""" - # Update the last assistant message with tool calls + # Update the last assistant message with tool calls (completion models need this) if conversation and conversation[-1]["role"] == "assistant": - conversation[-1]["tool_calls"] = tool_calls + conversation[-1].update(tool_calls_message) - # Add tool result messages - for tool_call, result in zip(tool_calls, results): - conversation.append( - { - "role": "tool", - "name": tool_call["function"]["name"], - "tool_call_id": tool_call["id"], - "content": json.dumps(result) if not isinstance(result, str) else result, - } - ) + # Add the formatted tool result messages directly - they're already in the right format + conversation.extend(formatted_results) class ResponsesConversationManager(ConversationManager): @@ -198,18 +196,17 @@ def add_assistant_response(self, conversation: List[Dict[str, Any]], response: D conversation.append({"role": "assistant", "content": response["generation"]}) def add_tool_results( - self, conversation: List[Dict[str, Any]], tool_calls: List[Dict[str, Any]], results: List[Dict[str, Any]] + self, + conversation: List[Dict[str, Any]], + tool_calls_message: Dict[str, Any], + formatted_results: List[Dict[str, Any]], ) -> None: """Add tool call results to the conversation.""" - # For responses API, add tool results as function_call_output items - for tool_call, result in zip(tool_calls, results): - conversation.append( - { - "type": "function_call_output", - "call_id": tool_call.get("id", "unknown"), - "output": json.dumps(result) if not isinstance(result, str) else result, - } - ) + # For responses models, we don't need to update the last assistant message + # The tool calls are already part of the serialized output + + # Add the formatted tool result messages directly - they're already in the right format + conversation.extend(formatted_results) # ============================== From 97b4f2f0d71da2cc59f5fc6a756d2046a9c2dbf5 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Wed, 24 Sep 2025 14:09:52 -0700 Subject: [PATCH 08/19] MAINT remove unused method Signed-off-by: George Armstrong --- nemo_skills/mcp/adapters.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/nemo_skills/mcp/adapters.py b/nemo_skills/mcp/adapters.py index 2b26dfc46e..f5f398e163 100644 --- a/nemo_skills/mcp/adapters.py +++ b/nemo_skills/mcp/adapters.py @@ -46,11 +46,6 @@ def format(self, tool_call: ChatCompletionMessageToolCall, result: dict) -> dict class ConversationManager(ABC): """Manages conversation history building for different model types.""" - @abstractmethod - def add_user_message(self, conversation: List[Dict[str, Any]], content: str) -> None: - """Add a user message to the conversation.""" - raise NotImplementedError("Subclasses must implement this method.") - @abstractmethod def add_assistant_response(self, conversation: List[Dict[str, Any]], response: Dict[str, Any]) -> None: """Add an assistant response to the conversation.""" @@ -150,10 +145,6 @@ def format(self, tool_call, result): class CompletionConversationManager(ConversationManager): """Manages conversation history for chat completion models.""" - def add_user_message(self, conversation: List[Dict[str, Any]], content: str) -> None: - """Add a user message to the conversation.""" - conversation.append({"role": "user", "content": content}) - def add_assistant_response(self, conversation: List[Dict[str, Any]], response: Dict[str, Any]) -> None: """Add an assistant response to the conversation.""" message = {"role": "assistant", "content": response["generation"]} @@ -182,10 +173,6 @@ def add_tool_results( class ResponsesConversationManager(ConversationManager): """Manages conversation history for responses API models.""" - def add_user_message(self, conversation: List[Dict[str, Any]], content: str) -> None: - """Add a user message to the conversation.""" - conversation.append({"role": "user", "content": content}) - def add_assistant_response(self, conversation: List[Dict[str, Any]], response: Dict[str, Any]) -> None: """Add an assistant response to the conversation using serialized output.""" # Use the serialized output from the responses API From 92e61b06aa9d2c9b491cb4ed70c3535e914a74e9 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Wed, 24 Sep 2025 15:36:16 -0700 Subject: [PATCH 09/19] MAINT move import to top Signed-off-by: George Armstrong --- nemo_skills/inference/model/tool_call.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/nemo_skills/inference/model/tool_call.py b/nemo_skills/inference/model/tool_call.py index a2135d05d0..dfbdfbb077 100644 --- a/nemo_skills/inference/model/tool_call.py +++ b/nemo_skills/inference/model/tool_call.py @@ -34,6 +34,7 @@ from nemo_skills.utils import get_logger_name from .base import BaseModel +from .responses import ResponsesModel LOG = logging.getLogger(get_logger_name(__file__)) @@ -70,9 +71,6 @@ def __init__( def _setup_adapters(self): """Set up adapters based on model type.""" - # Import here to avoid circular imports - from .responses import ResponsesModel - # Detect model type and configure adapters if isinstance(self.model, ResponsesModel): # Responses API model - uses flatter tool schema format From 7a1f01949e62fdbc46f59e0349da545c31c7a948 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Thu, 25 Sep 2025 11:00:51 -0700 Subject: [PATCH 10/19] Draft add client handler Signed-off-by: George Armstrong --- nemo_skills/inference/model/__init__.py | 15 +- nemo_skills/inference/model/base.py | 521 ++++++++++++++++++----- nemo_skills/inference/model/defaults.py | 90 ++++ nemo_skills/inference/model/openai.py | 45 ++ nemo_skills/inference/model/tool_call.py | 7 +- 5 files changed, 565 insertions(+), 113 deletions(-) create mode 100644 nemo_skills/inference/model/defaults.py diff --git a/nemo_skills/inference/model/__init__.py b/nemo_skills/inference/model/__init__.py index 732978e338..25c7f43a26 100644 --- a/nemo_skills/inference/model/__init__.py +++ b/nemo_skills/inference/model/__init__.py @@ -28,7 +28,6 @@ from .megatron import MegatronModel from .openai import OpenAIModel from .parallel_thinking import ParallelThinkingConfig, ParallelThinkingTask -from .responses import ResponsesModel # Tool Calling from .tool_call import ToolCallingWrapper @@ -48,22 +47,23 @@ "gemini": GeminiModel, "vllm": VLLMModel, "sglang": VLLMModel, - "responses": ResponsesModel, } -def get_model(server_type, tokenizer=None, **kwargs): +def get_model(server_type, client_type="chat_completion", tokenizer=None, **kwargs): """A helper function to make it easier to set server through cmd.""" model_class = models[server_type.lower()] if server_type == "trtllm" and kwargs.get("enable_soft_fail", False): if kwargs.get("context_limit_retry_strategy", None) is not None: raise ValueError("context_limit_retry_strategy is not supported for trtllm") - return model_class(tokenizer=tokenizer, **kwargs) + return model_class(client_type=client_type, tokenizer=tokenizer, **kwargs) -def get_code_execution_model(server_type, tokenizer=None, code_execution=None, sandbox=None, **kwargs): +def get_code_execution_model( + server_type, client_type="chat_completion", tokenizer=None, code_execution=None, sandbox=None, **kwargs +): """A helper function to make it easier to set server through cmd.""" - model = get_model(server_type=server_type, tokenizer=tokenizer, **kwargs) + model = get_model(server_type=server_type, client_type=client_type, tokenizer=tokenizer, **kwargs) if code_execution is None: code_execution = {} code_execution_config = CodeExecutionConfig(**code_execution) @@ -96,6 +96,7 @@ def get_parallel_thinking_model( def get_tool_calling_model( model, + client_type="chat_completion", tokenizer=None, additional_config=None, tool_modules: list[str] | None = None, @@ -103,7 +104,7 @@ def get_tool_calling_model( **kwargs, ): if isinstance(model, str): - model = get_model(model=model, tokenizer=tokenizer, **kwargs) + model = get_model(server_type=model, client_type=client_type, tokenizer=tokenizer, **kwargs) return ToolCallingWrapper( model, tool_modules=tool_modules, diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index 94aa096450..3e51c379e4 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -29,6 +29,247 @@ LOG = logging.getLogger(get_logger_name(__file__)) +class BaseClientHandler: + """Base client handler with clean public API""" + + def __init__(self, model_instance): + self.model = model_instance + self.defaults = None # Will be set in subclasses + self.setup_clients() # Public method + + def setup_clients(self): + """Public method for setting up API clients - override in subclasses""" + pass + + def get_supported_params(self) -> set: + """Public method to get supported parameters - override in subclasses""" + raise NotImplementedError + + def extract_and_validate_params(self, **kwargs) -> dict: + """Public method for parameter extraction and validation""" + supported = self.get_supported_params() + + # Let model filter/restrict parameters + model_supported = self.model.get_supported_params() + if model_supported: + supported = supported.intersection(model_supported) + + # Check for unsupported parameters + provided = set(kwargs.keys()) + unsupported = provided - supported + if unsupported: + raise ValueError( + f"Unsupported parameters for {self.__class__.__name__}: {unsupported}. Supported: {sorted(supported)}" + ) + + # Extract with defaults + params = {} + for param_name in supported: + if param_name in kwargs: + params[param_name] = kwargs[param_name] + else: + params[param_name] = getattr(self.defaults, param_name) + + return params + + def build_request_structure(self, prompt, params: dict) -> dict: + """Public method for building client-specific request structure""" + if isinstance(prompt, list): + return self.build_chat_request_structure(prompt, params) + else: + return self.build_completion_request_structure(prompt, params) + + def build_chat_request_structure(self, messages: list, params: dict) -> dict: + """Public method for chat request structure - override in subclasses""" + raise NotImplementedError + + def build_completion_request_structure(self, prompt: str, params: dict) -> dict: + """Public method for completion request structure - override in subclasses""" + raise NotImplementedError + + async def call_api_async(self, prompt, **kwargs): + """Public method for async API calls""" + # Extract and validate parameters + params = self.extract_and_validate_params(**kwargs) + + # Build request using two-stage process + request = self.build_request_structure(prompt, params) + request = self.model.apply_model_specific_params(request, params) + + # Make the actual API call + return await self.make_async_call(request, prompt) + + def call_api_sync(self, prompt, **kwargs): + """Public method for sync API calls""" + # Extract and validate parameters + params = self.extract_and_validate_params(**kwargs) + + # Build request using two-stage process + request = self.build_request_structure(prompt, params) + request = self.model.apply_model_specific_params(request, params) + + # Make the actual API call + return self.make_sync_call(request, prompt) + + async def make_async_call(self, request: dict, prompt): + """Public method for making async API calls - override in subclasses""" + raise NotImplementedError + + def make_sync_call(self, request: dict, prompt): + """Public method for making sync API calls - override in subclasses""" + raise NotImplementedError + + def parse_response(self, response, **kwargs) -> dict: + """Public method for parsing responses - override in subclasses""" + raise NotImplementedError + + +class ChatCompletionHandler(BaseClientHandler): + """Handler for chat completion and text completion APIs via litellm""" + + def __init__(self, model_instance): + from .defaults import CHAT_COMPLETION_PARAMS, GenerationDefaults + + super().__init__(model_instance) + self.defaults = GenerationDefaults() + self.supported_params = CHAT_COMPLETION_PARAMS + + def get_supported_params(self) -> set: + return self.supported_params + + def setup_clients(self): + """Setup litellm configuration""" + model_litellm = f"{self.model.MODEL_PROVIDER}/{self.model.model_name_or_path}" + self.litellm_kwargs = dict( + model=model_litellm, + max_retries=3, + api_key=self.model.api_key, + base_url=self.model.base_url, + ) + # Setup litellm sessions + httpx_limits = httpx.Limits(max_keepalive_connections=2048, max_connections=2048) + litellm.client_session = httpx.Client(limits=httpx_limits) + litellm.aclient_session = httpx.AsyncClient(limits=httpx_limits) + + def build_chat_request_structure(self, messages: list, params: dict) -> dict: + """Build chat completion request structure""" + return { + "messages": messages, + "max_tokens": params["tokens_to_generate"], + "temperature": params["temperature"], + "top_p": params["top_p"], + "seed": params["random_seed"], + "stop": params["stop_phrases"], + "logprobs": params["top_logprobs"] is not None, + "top_logprobs": params["top_logprobs"], + "stream": params["stream"], + "tools": params["tools"], + "timeout": params["timeout"], + } + + def build_completion_request_structure(self, prompt: str, params: dict) -> dict: + """Build text completion request structure""" + return { + "prompt": prompt, + "max_tokens": params["tokens_to_generate"], + "temperature": params["temperature"], + "top_p": params["top_p"], + "seed": params["random_seed"], + "stop": params["stop_phrases"], + "logprobs": params["top_logprobs"], + "stream": params["stream"], + "timeout": params["timeout"], + } + + async def make_async_call(self, request: dict, prompt): + """Make async API call via litellm""" + if isinstance(prompt, list): + return await litellm.acompletion(**request, **self.litellm_kwargs) + else: + return await litellm.atext_completion(**request, **self.litellm_kwargs) + + def make_sync_call(self, request: dict, prompt): + """Make sync API call via litellm""" + if isinstance(prompt, list): + return litellm.completion(**request, **self.litellm_kwargs) + else: + return litellm.text_completion(**request, **self.litellm_kwargs) + + def parse_response(self, response, **kwargs) -> dict: + """Parse response using existing BaseModel methods""" + if hasattr(response, "choices") and hasattr(response.choices[0], "message"): + return self.model.parse_chat_completion_response(response, **kwargs) + else: + return self.model.parse_completion_response(response, **kwargs) + + +class ResponsesHandler(BaseClientHandler): + """Handler for responses API using direct OpenAI client""" + + def __init__(self, model_instance): + from .defaults import RESPONSES_PARAMS, GenerationDefaults + + super().__init__(model_instance) + self.defaults = GenerationDefaults() + self.supported_params = RESPONSES_PARAMS + + def get_supported_params(self) -> set: + return self.supported_params + + def setup_clients(self): + """Setup OpenAI clients directly""" + from openai import AsyncOpenAI, OpenAI + + self.sync_client = OpenAI(base_url=self.model.base_url, api_key=self.model.api_key) + self.async_client = AsyncOpenAI(base_url=self.model.base_url, api_key=self.model.api_key) + + def build_chat_request_structure(self, messages: list, params: dict) -> dict: + """Build responses API request structure""" + return { + "input": messages, + "max_output_tokens": params["tokens_to_generate"], + "temperature": params["temperature"], + "top_p": params["top_p"], + "stream": params["stream"], + "tools": params["tools"], + "extra_body": { + "seed": params["random_seed"], + "reasoning_effort": params["reasoning_effort"], + "timeout": params["timeout"], + "stop": params["stop_phrases"], + "top_logprobs": params["top_logprobs"], + "top_k": params["top_k"], + "min_p": params["min_p"], + "repetition_penalty": params["repetition_penalty"], + **(params["extra_body"] or {}), + }, + } + + def build_completion_request_structure(self, prompt: str, params: dict) -> dict: + """Responses API doesn't support completion - raise error""" + raise ValueError("ResponsesHandler only supports message lists, not string prompts") + + async def make_async_call(self, request: dict, prompt): + """Make async call to responses API""" + return await self.async_client.responses.create(model=self.model.model_name_or_path, **request) + + def make_sync_call(self, request: dict, prompt): + """Make sync call to responses API""" + return self.sync_client.responses.create(model=self.model.model_name_or_path, **request) + + def parse_response(self, response, **kwargs) -> dict: + """Parse responses API response""" + return self.model.parse_responses_response(response, **kwargs) + + +# Global client handler registry +CLIENT_HANDLERS = { + "chat_completion": ChatCompletionHandler, + "responses": ResponsesHandler, + "completion": ChatCompletionHandler, # Same handler, different method selection +} + + class BaseModel: """Base model class for handling requests to the inference server. @@ -49,11 +290,11 @@ class BaseModel: def __init__( self, model: str, + client_type: str = "chat_completion", tokenizer: str | None = None, api_key: str | None = None, api_key_env_var: str | None = None, base_url: str | None = None, - max_retries: int = 3, use_v1_endpoint: bool = True, host: str = "127.0.0.1", port: str = "5000", @@ -64,21 +305,46 @@ def __init__( context_limit_retry_strategy: str | None = None, num_special_tokens_budget: int = 100, ): - self._tunnel = None + # Common model properties self.model_name_or_path = model + self.client_type = client_type self.server_host = host self.server_port = port - self.ssh_server = ssh_server - self.ssh_key_path = ssh_key_path + + # SSH tunnel setup (general networking) + self._setup_ssh_tunnel(ssh_server, ssh_key_path) + + # Base URL setup (general) + self.base_url = self._setup_base_url(base_url, use_v1_endpoint) + + # API key resolution (general) + self.api_key = self._resolve_api_key(api_key, api_key_env_var, base_url) + + # Context retry config (general) self.context_limit_retry_config = ContextLimitRetryConfig( enable_soft_fail=enable_soft_fail, strategy=context_limit_retry_strategy, num_special_tokens_budget=num_special_tokens_budget, ) - if ssh_server is None: - self.ssh_server = os.getenv("NEMO_SKILLS_SSH_SERVER") - if ssh_key_path is None: - self.ssh_key_path = os.getenv("NEMO_SKILLS_SSH_KEY_PATH") + + # Tokenizer setup (general) + if enable_soft_fail: + self.tokenizer = self._get_tokenizer(tokenizer) + else: + self.tokenizer = None + + # Initialize client handler LAST + if client_type not in CLIENT_HANDLERS: + raise ValueError(f"Unsupported client_type: {client_type}. Available: {list(CLIENT_HANDLERS.keys())}") + + handler_class = CLIENT_HANDLERS[client_type] + self.client_handler = handler_class(self) # Public attribute + + def _setup_ssh_tunnel(self, ssh_server: str | None, ssh_key_path: str | None): + """Setup SSH tunnel if needed""" + self._tunnel = None + self.ssh_server = ssh_server or os.getenv("NEMO_SKILLS_SSH_SERVER") + self.ssh_key_path = ssh_key_path or os.getenv("NEMO_SKILLS_SSH_KEY_PATH") if self.ssh_server and self.ssh_key_path: import sshtunnel @@ -99,35 +365,22 @@ def __init__( self.server_host = "127.0.0.1" self.server_port = str(self._tunnel.local_bind_port) + def _setup_base_url(self, base_url: str | None, use_v1_endpoint: bool) -> str: + """Setup base URL for API calls""" if base_url is None: v1_suffix = "/v1" if use_v1_endpoint else "" - self.base_url = f"http://{self.server_host}:{self.server_port}{v1_suffix}" + return f"http://{self.server_host}:{self.server_port}{v1_suffix}" elif base_url == "": - # We don't want to use base_url if it is an empty string - base_url = None - else: - self.base_url = base_url - - if enable_soft_fail: - self.tokenizer = self._get_tokenizer(tokenizer) + return None else: - self.tokenizer = None - - api_key = self._get_api_key(api_key, api_key_env_var, base_url) - if api_key is None: # self-hosted models don't need the key, but still require the parameter - api_key = "EMPTY" + return base_url - model_litellm = f"{self.MODEL_PROVIDER}/{model}" - # Passed to litellm every time we call it - self.litellm_kwargs = dict( - model=model_litellm, - max_retries=max_retries, - api_key=api_key, - base_url=self.base_url, - ) - httpx_limits = httpx.Limits(max_keepalive_connections=2048, max_connections=2048) - litellm.client_session = httpx.Client(limits=httpx_limits) - litellm.aclient_session = httpx.AsyncClient(limits=httpx_limits) + def _resolve_api_key(self, api_key: str | None, api_key_env_var: str | None, base_url: str) -> str | None: + """Resolve API key from various sources""" + resolved_key = self._get_api_key(api_key, api_key_env_var, base_url) + if resolved_key is None: # self-hosted models don't need the key + resolved_key = "EMPTY" + return resolved_key def _get_api_key(self, api_key: str | None, api_key_env_var: str | None, base_url: str) -> str | None: if api_key: # explicit cmd argument always takes precedence @@ -182,6 +435,110 @@ def _initialize_tokenizer(self, tokenizer: str | None) -> WrapperAutoTokenizer | if isinstance(tokenizer, str): return WrapperAutoTokenizer(tokenizer) + # Public methods for client handlers to call + def get_supported_params(self) -> set: + """Public method for models to restrict parameters - override in subclasses""" + return set() # Base implementation allows all + + def apply_model_specific_params(self, request: dict, params: dict) -> dict: + """Public method for model-specific parameter handling - override in subclasses""" + return request + + def parse_chat_completion_response(self, response, **kwargs) -> dict: + """Public method for parsing chat completion responses""" + return self._parse_chat_completion_response(response, **kwargs) + + def parse_completion_response(self, response, **kwargs) -> dict: + """Public method for parsing completion responses""" + return self._parse_completion_response(response, **kwargs) + + def parse_responses_response(self, response, **kwargs) -> dict: + """Public method for parsing responses API responses""" + # Move implementation from ResponsesModel here + result = {"generation": "", "num_generated_tokens": 0} + + # Get token usage + if hasattr(response, "usage"): + result["num_generated_tokens"] = getattr(response.usage, "output_tokens", 0) + + # Check for tool calls in the output array + tool_calls = [] + reasoning_content = "" + generation_text = "" + + if hasattr(response, "output") and response.output: + for output_item in response.output: + # Handle reasoning content + if hasattr(output_item, "type") and output_item.type == "reasoning": + if hasattr(output_item, "content") and output_item.content: + for content_item in output_item.content: + if hasattr(content_item, "text"): + reasoning_content += content_item.text + "\n" + + # Handle function calls + elif hasattr(output_item, "type") and output_item.type == "function_call": + tool_calls.append(output_item) + + # Handle message content + elif hasattr(output_item, "type") and output_item.type == "message": + if hasattr(output_item, "content") and output_item.content: + for content_item in output_item.content: + if hasattr(content_item, "text"): + generation_text += content_item.text + + # Set the appropriate response fields + if tool_calls: + result["tool_calls"] = tool_calls + result["generation"] = "" # No text generation when there are tool calls + else: + result["generation"] = generation_text + + # Add reasoning content if available + if reasoning_content: + result["reasoning_content"] = reasoning_content.strip() + + # Add finish reason if available + if hasattr(response, "status"): + result["finish_reason"] = response.status + + # Add serialized output for conversation history + result["serialized_output"] = self._serialize_response_output(response) + + if kwargs.get("include_response", False): + result["response"] = response + + return result + + def _serialize_response_output(self, response): + """Serialize response output objects using model_dump() for conversation history.""" + serialized_output = [] + + if hasattr(response, "output") and response.output: + for output_item in response.output: + try: + # Try to use model_dump() method if available (Pydantic models) + if hasattr(output_item, "model_dump"): + serialized_output.append(output_item.model_dump()) + # Fallback to dict conversion + elif hasattr(output_item, "__dict__"): + serialized_output.append(output_item.__dict__) + # Last resort: convert to string representation + else: + serialized_output.append({"content": str(output_item), "type": "unknown"}) + except Exception as e: + LOG.warning(f"Failed to serialize output item: {e}") + # Fallback serialization + serialized_output.append({"content": str(output_item), "type": "error", "error": str(e)}) + + return serialized_output + + def _handle_streaming_response(self, response): + """Handle streaming responses based on response type""" + if hasattr(response, "choices") and hasattr(response.choices[0], "message"): + return self._stream_chat_chunks_sync(response) + else: + return self._stream_completion_chunks_sync(response) + @abc.abstractmethod def _build_chat_request_params(self, **kwargs) -> dict: pass @@ -220,7 +577,7 @@ async def generate_async( include_response: bool = False, extra_body: dict = None, ) -> dict: - """Native async version of generate for single prompt.""" + """Unified async version of generate for single prompt.""" # Check tool calls are a list of dict if tools is not None: @@ -229,21 +586,11 @@ async def generate_async( if not isinstance(tool, dict): raise ValueError(f"Tool must be a dictionary, got {type(tool)}") - kwargs = { - "tokens_to_generate": tokens_to_generate, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "repetition_penalty": repetition_penalty, - "random_seed": random_seed, - "stop_phrases": stop_phrases, - "top_logprobs": top_logprobs, - "timeout": timeout, - "reasoning_effort": reasoning_effort, - "tools": tools, - "extra_body": extra_body, - } + # Build kwargs dict with only non-None values to avoid overriding defaults unnecessarily + kwargs = {} + for param, value in locals().items(): + if param not in ["self", "prompt", "kwargs"] and value is not None: + kwargs[param] = value # TODO: remove this after we no longer use gpt-oss or it's fixed in vllm max_retries = 2 @@ -251,28 +598,16 @@ async def generate_async( while retry_count <= max_retries: try: - if isinstance(prompt, list): - request_params = self._build_chat_request_params(messages=prompt, stream=stream, **kwargs) - response = await litellm.acompletion(**request_params, **self.litellm_kwargs) - if stream: - result = self._stream_chat_chunks_async(response) - else: - result = self._parse_chat_completion_response( - response, include_response=include_response, **kwargs - ) - - elif isinstance(prompt, str): - request_params = self._build_completion_request_params(prompt=prompt, stream=stream, **kwargs) - response = await litellm.atext_completion(**request_params, **self.litellm_kwargs) - if stream: - result = self._stream_completion_chunks_async(response) - else: - result = self._parse_completion_response(response, include_response=include_response, **kwargs) + # Delegate to client handler using public API + response = await self.client_handler.call_api_async(prompt, **kwargs) + + if stream: + return self._handle_streaming_response(response) else: - raise TypeError(f"Unsupported prompt type: {type(prompt)}") - if not stream: - self._maybe_apply_stop_phrase_removal(result, remove_stop_phrases, stop_phrases) - return result + result = self.client_handler.parse_response(response, **kwargs) + if remove_stop_phrases: + self._maybe_apply_stop_phrase_removal(result, remove_stop_phrases, stop_phrases) + return result except openai.BadRequestError as e: if "output messages (reasoning and final)" in str(e): @@ -310,7 +645,7 @@ def generate_sync( extra_body: dict = None, ) -> dict: """ - Synchronous version of generate for single prompt. + Unified synchronous version of generate for single prompt. See generate_async for full list of parameters. """ # Check tool calls are a list of dict @@ -320,40 +655,22 @@ def generate_sync( if not isinstance(tool, dict): raise ValueError(f"Tool must be a dictionary, got {type(tool)}") - kwargs = { - "tokens_to_generate": tokens_to_generate, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "repetition_penalty": repetition_penalty, - "random_seed": random_seed, - "stop_phrases": stop_phrases, - "top_logprobs": top_logprobs, - "timeout": timeout, - "reasoning_effort": reasoning_effort, - "tools": tools, - "extra_body": extra_body, - } - request_params = self._build_request_params(prompt=prompt, stream=stream, **kwargs) - if isinstance(prompt, list): - response = litellm.completion(**request_params, **self.litellm_kwargs) - if stream: - result = self._stream_chat_chunks_sync(response) - else: - result = self._parse_chat_completion_response(response, include_response=include_response, **kwargs) + # Build kwargs dict with only non-None values to avoid overriding defaults unnecessarily + kwargs = {} + for param, value in locals().items(): + if param not in ["self", "prompt", "kwargs"] and value is not None: + kwargs[param] = value - elif isinstance(prompt, str): - response = litellm.text_completion(**request_params, **self.litellm_kwargs) - if stream: - result = self._stream_completion_chunks_sync(response) - else: - result = self._parse_completion_response(response, include_response=include_response, **kwargs) - else: - raise TypeError(f"Unsupported prompt type: {type(prompt)}") + # Delegate to client handler using public API + response = self.client_handler.call_api_sync(prompt, **kwargs) - self._maybe_apply_stop_phrase_removal(result, remove_stop_phrases, stop_phrases) - return result + if stream: + return self._handle_streaming_response(response) + else: + result = self.client_handler.parse_response(response, **kwargs) + if remove_stop_phrases: + self._maybe_apply_stop_phrase_removal(result, remove_stop_phrases, stop_phrases) + return result def _parse_completion_response( self, response: "openai.types.Completion", include_response: bool = False, **kwargs diff --git a/nemo_skills/inference/model/defaults.py b/nemo_skills/inference/model/defaults.py new file mode 100644 index 0000000000..0f55beeb42 --- /dev/null +++ b/nemo_skills/inference/model/defaults.py @@ -0,0 +1,90 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import List, Optional + + +@dataclass +class GenerationDefaults: + """Shared defaults for generation parameters across all models and handlers""" + + tokens_to_generate: int = 512 + temperature: float = 0.0 + top_p: float = 0.95 + top_k: int = -1 + min_p: float = 0.0 + repetition_penalty: float = 1.0 + random_seed: Optional[int] = None + stop_phrases: Optional[List[str]] = None + top_logprobs: Optional[int] = None + timeout: Optional[float] = 14400 + remove_stop_phrases: bool = True + stream: bool = False + reasoning_effort: Optional[str] = None + tools: Optional[List[dict]] = None + include_response: bool = False + extra_body: Optional[dict] = None + + +# Base parameter sets for different client types +CHAT_COMPLETION_PARAMS = { + "tokens_to_generate", + "temperature", + "top_p", + "top_k", + "min_p", + "repetition_penalty", + "random_seed", + "stop_phrases", + "top_logprobs", + "timeout", + "stream", + "reasoning_effort", + "tools", + "extra_body", +} + +RESPONSES_PARAMS = { + "tokens_to_generate", + "temperature", + "top_p", + "top_k", + "min_p", + "repetition_penalty", + "random_seed", + "stop_phrases", + "top_logprobs", + "timeout", + "stream", + "reasoning_effort", + "tools", + "extra_body", +} + +COMPLETION_PARAMS = { + "tokens_to_generate", + "temperature", + "top_p", + "top_k", + "min_p", + "repetition_penalty", + "random_seed", + "stop_phrases", + "top_logprobs", + "timeout", + "stream", + "extra_body", + # Note: no 'reasoning_effort' or 'tools' for completion +} diff --git a/nemo_skills/inference/model/openai.py b/nemo_skills/inference/model/openai.py index 9c7faa85d1..f3ccf99ce2 100644 --- a/nemo_skills/inference/model/openai.py +++ b/nemo_skills/inference/model/openai.py @@ -17,6 +17,15 @@ import re from .base import BaseModel +from .defaults import CHAT_COMPLETION_PARAMS + +# OpenAI-specific parameter restrictions +OPENAI_REASONING_UNSUPPORTED = {"temperature", "top_p", "repetition_penalty", "top_logprobs"} + +OPENAI_GENERAL_UNSUPPORTED = { + "top_k", + "min_p", # OpenAI doesn't support these +} class OpenAIModel(BaseModel): @@ -27,6 +36,7 @@ def __init__( model: str | None = None, base_url: str | None = None, max_retries: int = 3, + client_type: str = "chat_completion", **kwargs, ): model = model or os.getenv("NEMO_SKILLS_OPENAI_MODEL") @@ -39,6 +49,7 @@ def __init__( super().__init__( model=model, + client_type=client_type, base_url=base_url, max_retries=max_retries, **kwargs, @@ -61,6 +72,40 @@ def _get_api_key(self, api_key: str | None, api_key_env_var: str | None, base_ur def _is_reasoning_model(self, model_name: str) -> bool: return re.match(r"^o\d", model_name) + def is_reasoning_model(self, model_name: str) -> bool: + """Public method to check if model is a reasoning model""" + return self._is_reasoning_model(model_name) + + def get_supported_params(self) -> set: + """OpenAI-specific parameter restrictions""" + base_params = CHAT_COMPLETION_PARAMS.copy() + + # Remove OpenAI-unsupported parameters + base_params -= OPENAI_GENERAL_UNSUPPORTED + + # Further restrictions for reasoning models + if self.is_reasoning_model(self.model_name_or_path): + base_params -= OPENAI_REASONING_UNSUPPORTED + + return base_params + + def apply_model_specific_params(self, request: dict, params: dict) -> dict: + """Apply OpenAI-specific parameter transformations""" + if self.is_reasoning_model(self.model_name_or_path): + # Reasoning model specific handling + if "messages" in request: + request["messages"] = [ + {**msg, "role": "developer"} if msg.get("role") == "system" else msg for msg in request["messages"] + ] + if params.get("reasoning_effort"): + request["reasoning_effort"] = params["reasoning_effort"] + else: + # Standard OpenAI model handling + if "repetition_penalty" in params and params["repetition_penalty"] != 1.0: + request["presence_penalty"] = params["repetition_penalty"] + + return request + def _build_completion_request_params(self, **kwargs) -> dict: kwargs = copy.deepcopy(kwargs) assert kwargs.pop("tools", None) is None, "tools are not supported by completion requests." diff --git a/nemo_skills/inference/model/tool_call.py b/nemo_skills/inference/model/tool_call.py index dfbdfbb077..b787dd58d3 100644 --- a/nemo_skills/inference/model/tool_call.py +++ b/nemo_skills/inference/model/tool_call.py @@ -34,7 +34,6 @@ from nemo_skills.utils import get_logger_name from .base import BaseModel -from .responses import ResponsesModel LOG = logging.getLogger(get_logger_name(__file__)) @@ -70,9 +69,9 @@ def __init__( self._setup_adapters() def _setup_adapters(self): - """Set up adapters based on model type.""" - # Detect model type and configure adapters - if isinstance(self.model, ResponsesModel): + """Set up adapters based on client type.""" + # Use client_type instead of model instance checks + if self.model.client_type == "responses": # Responses API model - uses flatter tool schema format self.schema_adapter = ResponsesSchemaAdapter() self.call_interpreter = ResponsesCallInterpreter() From 5688efea50e91bd8fcf8b540941c3636c787d34a Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Thu, 25 Sep 2025 13:05:57 -0700 Subject: [PATCH 11/19] FIX up responses conversion Signed-off-by: George Armstrong --- nemo_skills/inference/generate.py | 9 +- nemo_skills/inference/model/__init__.py | 2 +- nemo_skills/inference/model/base.py | 179 ++++++++--- nemo_skills/inference/model/defaults.py | 6 + nemo_skills/inference/model/openai.py | 4 +- nemo_skills/inference/model/responses.py | 375 ----------------------- 6 files changed, 161 insertions(+), 414 deletions(-) delete mode 100644 nemo_skills/inference/model/responses.py diff --git a/nemo_skills/inference/generate.py b/nemo_skills/inference/generate.py index 4361d394ca..40beacafbb 100644 --- a/nemo_skills/inference/generate.py +++ b/nemo_skills/inference/generate.py @@ -92,6 +92,8 @@ class GenerateSolutionsConfig: # Inference server configuration {server_params} server: dict = field(default_factory=dict) + # Client type for API calls (chat_completion, responses, completion) + client_type: str = "chat_completion" # Sandbox configuration {sandbox_params} sandbox: dict = field(default_factory=dict) # Prompt configuration - path to yaml files @@ -338,17 +340,20 @@ def setup_llm(self): self.sandbox = get_sandbox(**self.cfg.sandbox) if self.cfg.sandbox is not None else None if self.cfg.code_execution: - llm = get_code_execution_model(**self.cfg.server, tokenizer=self.tokenizer, sandbox=self.sandbox) + llm = get_code_execution_model( + **self.cfg.server, client_type=self.cfg.client_type, tokenizer=self.tokenizer, sandbox=self.sandbox + ) elif self.cfg.tool_modules is not None: llm = get_tool_calling_model( **self.cfg.server, + client_type=self.cfg.client_type, tool_modules=self.cfg.tool_modules, tool_overrides=self.cfg.tool_overrides, tokenizer=self.tokenizer, additional_config={"sandbox": self.cfg.sandbox}, ) else: - llm = get_model(**self.cfg.server, tokenizer=self.tokenizer) + llm = get_model(**self.cfg.server, client_type=self.cfg.client_type, tokenizer=self.tokenizer) if self.cfg.parallel_thinking.mode is not None: # We don't want to override these key variables which overlap with self.cfg diff --git a/nemo_skills/inference/model/__init__.py b/nemo_skills/inference/model/__init__.py index 25c7f43a26..9e0e5c7d88 100644 --- a/nemo_skills/inference/model/__init__.py +++ b/nemo_skills/inference/model/__init__.py @@ -104,7 +104,7 @@ def get_tool_calling_model( **kwargs, ): if isinstance(model, str): - model = get_model(server_type=model, client_type=client_type, tokenizer=tokenizer, **kwargs) + model = get_model(model=model, client_type=client_type, tokenizer=tokenizer, **kwargs) return ToolCallingWrapper( model, tool_modules=tool_modules, diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index 3e51c379e4..666bead3e4 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -47,6 +47,7 @@ def get_supported_params(self) -> set: def extract_and_validate_params(self, **kwargs) -> dict: """Public method for parameter extraction and validation""" + LOG.info(f"extract_and_validate_params called with kwargs: {list(kwargs.keys())}") supported = self.get_supported_params() # Let model filter/restrict parameters @@ -54,17 +55,29 @@ def extract_and_validate_params(self, **kwargs) -> dict: if model_supported: supported = supported.intersection(model_supported) - # Check for unsupported parameters + # Check for unsupported parameters, but only for parameters that differ from defaults provided = set(kwargs.keys()) - unsupported = provided - supported - if unsupported: + unsupported_non_default = set() + + for param_name in provided: + if param_name not in supported: + # Check if this parameter is set to its default value + default_value = getattr(self.defaults, param_name, None) + provided_value = kwargs[param_name] + + # Only consider it unsupported if it's not the default value + if default_value is None or provided_value != default_value: + unsupported_non_default.add(param_name) + + if unsupported_non_default: raise ValueError( - f"Unsupported parameters for {self.__class__.__name__}: {unsupported}. Supported: {sorted(supported)}" + f"Unsupported parameters for {self.__class__.__name__}: {unsupported_non_default}. Supported: {sorted(supported)}" ) - # Extract with defaults + # Extract with defaults - include all provided parameters, even unsupported ones if they're default values params = {} - for param_name in supported: + all_param_names = supported.union(provided) + for param_name in all_param_names: if param_name in kwargs: params[param_name] = kwargs[param_name] else: @@ -142,7 +155,7 @@ def setup_clients(self): model_litellm = f"{self.model.MODEL_PROVIDER}/{self.model.model_name_or_path}" self.litellm_kwargs = dict( model=model_litellm, - max_retries=3, + max_retries=getattr(self.model, "max_retries", 3), api_key=self.model.api_key, base_url=self.model.base_url, ) @@ -153,7 +166,7 @@ def setup_clients(self): def build_chat_request_structure(self, messages: list, params: dict) -> dict: """Build chat completion request structure""" - return { + request = { "messages": messages, "max_tokens": params["tokens_to_generate"], "temperature": params["temperature"], @@ -167,9 +180,23 @@ def build_chat_request_structure(self, messages: list, params: dict) -> dict: "timeout": params["timeout"], } + # Add non-standard parameters to extra_body + extra_body = params.get("extra_body", {}).copy() if params.get("extra_body") else {} + if params.get("top_k", -1) != -1: + extra_body["top_k"] = params["top_k"] + if params.get("min_p", 0.0) != 0.0: + extra_body["min_p"] = params["min_p"] + if params.get("repetition_penalty", 1.0) != 1.0: + extra_body["repetition_penalty"] = params["repetition_penalty"] + + if extra_body: + request["extra_body"] = extra_body + + return request + def build_completion_request_structure(self, prompt: str, params: dict) -> dict: """Build text completion request structure""" - return { + request = { "prompt": prompt, "max_tokens": params["tokens_to_generate"], "temperature": params["temperature"], @@ -181,6 +208,20 @@ def build_completion_request_structure(self, prompt: str, params: dict) -> dict: "timeout": params["timeout"], } + # Add non-standard parameters to extra_body + extra_body = params.get("extra_body", {}).copy() if params.get("extra_body") else {} + if params.get("top_k", -1) != -1: + extra_body["top_k"] = params["top_k"] + if params.get("min_p", 0.0) != 0.0: + extra_body["min_p"] = params["min_p"] + if params.get("repetition_penalty", 1.0) != 1.0: + extra_body["repetition_penalty"] = params["repetition_penalty"] + + if extra_body: + request["extra_body"] = extra_body + + return request + async def make_async_call(self, request: dict, prompt): """Make async API call via litellm""" if isinstance(prompt, list): @@ -225,26 +266,47 @@ def setup_clients(self): def build_chat_request_structure(self, messages: list, params: dict) -> dict: """Build responses API request structure""" - return { + # Use proper list of dicts format for vLLM servers + request = { "input": messages, "max_output_tokens": params["tokens_to_generate"], "temperature": params["temperature"], "top_p": params["top_p"], "stream": params["stream"], - "tools": params["tools"], - "extra_body": { - "seed": params["random_seed"], - "reasoning_effort": params["reasoning_effort"], - "timeout": params["timeout"], - "stop": params["stop_phrases"], - "top_logprobs": params["top_logprobs"], - "top_k": params["top_k"], - "min_p": params["min_p"], - "repetition_penalty": params["repetition_penalty"], - **(params["extra_body"] or {}), - }, } + # Only include tools if they are provided + if params["tools"] is not None: + request["tools"] = params["tools"] + + # Add non-standard parameters to extra_body (OpenAI responses API requirement) + extra_body = {} + if params["random_seed"] is not None: + extra_body["seed"] = params["random_seed"] + if params["reasoning_effort"] is not None: + extra_body["reasoning_effort"] = params["reasoning_effort"] + if params["timeout"] is not None: + extra_body["timeout"] = params["timeout"] + if params["stop_phrases"] is not None: + extra_body["stop"] = params["stop_phrases"] + if params["top_logprobs"] is not None: + extra_body["top_logprobs"] = params["top_logprobs"] + if params["top_k"] != -1: # Only include if not default + extra_body["top_k"] = params["top_k"] + if params["min_p"] != 0.0: # Only include if not default + extra_body["min_p"] = params["min_p"] + if params["repetition_penalty"] != 1.0: # Only include if not default + extra_body["repetition_penalty"] = params["repetition_penalty"] + + # Add any additional extra_body parameters + if params["extra_body"]: + extra_body.update(params["extra_body"]) + + if extra_body: + request["extra_body"] = extra_body + + return request + def build_completion_request_structure(self, prompt: str, params: dict) -> dict: """Responses API doesn't support completion - raise error""" raise ValueError("ResponsesHandler only supports message lists, not string prompts") @@ -396,7 +458,7 @@ def _get_api_key(self, api_key: str | None, api_key_env_var: str | None, base_ur return api_key def __del__(self): - if self._tunnel: + if hasattr(self, "_tunnel") and self._tunnel: self._tunnel.stop() def _maybe_apply_stop_phrase_removal( @@ -457,9 +519,24 @@ def parse_responses_response(self, response, **kwargs) -> dict: # Move implementation from ResponsesModel here result = {"generation": "", "num_generated_tokens": 0} - # Get token usage + # Debug logging to understand response structure + LOG.debug(f"Responses API response type: {type(response)}") + LOG.debug(f"Response has usage: {hasattr(response, 'usage')}") if hasattr(response, "usage"): - result["num_generated_tokens"] = getattr(response.usage, "output_tokens", 0) + LOG.debug(f"Usage object: {response.usage}") + LOG.debug(f"Usage type: {type(response.usage)}") + if response.usage: + LOG.debug(f"Usage attributes: {dir(response.usage)}") + + # Get token usage - ensure it's always an integer + if hasattr(response, "usage") and response.usage: + tokens = getattr(response.usage, "output_tokens", None) + if tokens is None: + # Try alternative field names for token usage + tokens = getattr(response.usage, "completion_tokens", None) + result["num_generated_tokens"] = tokens if tokens is not None else 0 + else: + result["num_generated_tokens"] = 0 # Check for tool calls in the output array tool_calls = [] @@ -507,6 +584,10 @@ def parse_responses_response(self, response, **kwargs) -> dict: if kwargs.get("include_response", False): result["response"] = response + # Ensure num_generated_tokens is never None for metrics compatibility + if result["num_generated_tokens"] is None: + result["num_generated_tokens"] = 0 + return result def _serialize_response_output(self, response): @@ -586,11 +667,25 @@ async def generate_async( if not isinstance(tool, dict): raise ValueError(f"Tool must be a dictionary, got {type(tool)}") - # Build kwargs dict with only non-None values to avoid overriding defaults unnecessarily - kwargs = {} - for param, value in locals().items(): - if param not in ["self", "prompt", "kwargs"] and value is not None: - kwargs[param] = value + # Build kwargs dict explicitly to avoid capturing unwanted local variables + kwargs = { + "tokens_to_generate": tokens_to_generate, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "repetition_penalty": repetition_penalty, + "random_seed": random_seed, + "stop_phrases": stop_phrases, + "top_logprobs": top_logprobs, + "timeout": timeout, + "reasoning_effort": reasoning_effort, + "tools": tools, + "extra_body": extra_body, + "remove_stop_phrases": remove_stop_phrases, + "stream": stream, + "include_response": include_response, + } # TODO: remove this after we no longer use gpt-oss or it's fixed in vllm max_retries = 2 @@ -655,11 +750,25 @@ def generate_sync( if not isinstance(tool, dict): raise ValueError(f"Tool must be a dictionary, got {type(tool)}") - # Build kwargs dict with only non-None values to avoid overriding defaults unnecessarily - kwargs = {} - for param, value in locals().items(): - if param not in ["self", "prompt", "kwargs"] and value is not None: - kwargs[param] = value + # Build kwargs dict explicitly to avoid capturing unwanted local variables + kwargs = { + "tokens_to_generate": tokens_to_generate, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "repetition_penalty": repetition_penalty, + "random_seed": random_seed, + "stop_phrases": stop_phrases, + "top_logprobs": top_logprobs, + "timeout": timeout, + "reasoning_effort": reasoning_effort, + "tools": tools, + "extra_body": extra_body, + "remove_stop_phrases": remove_stop_phrases, + "stream": stream, + "include_response": include_response, + } # Delegate to client handler using public API response = self.client_handler.call_api_sync(prompt, **kwargs) diff --git a/nemo_skills/inference/model/defaults.py b/nemo_skills/inference/model/defaults.py index 0f55beeb42..2d7a5218d7 100644 --- a/nemo_skills/inference/model/defaults.py +++ b/nemo_skills/inference/model/defaults.py @@ -54,6 +54,8 @@ class GenerationDefaults: "reasoning_effort", "tools", "extra_body", + "remove_stop_phrases", + "include_response", } RESPONSES_PARAMS = { @@ -71,6 +73,8 @@ class GenerationDefaults: "reasoning_effort", "tools", "extra_body", + "remove_stop_phrases", + "include_response", } COMPLETION_PARAMS = { @@ -86,5 +90,7 @@ class GenerationDefaults: "timeout", "stream", "extra_body", + "remove_stop_phrases", + "include_response", # Note: no 'reasoning_effort' or 'tools' for completion } diff --git a/nemo_skills/inference/model/openai.py b/nemo_skills/inference/model/openai.py index f3ccf99ce2..fdd27b7baa 100644 --- a/nemo_skills/inference/model/openai.py +++ b/nemo_skills/inference/model/openai.py @@ -47,11 +47,13 @@ def __init__( if base_url is None: base_url = os.getenv("NEMO_SKILLS_OPENAI_BASE_URL", f"http://{host}:{port}/v1") + # Store max_retries for client handlers to use + self.max_retries = max_retries + super().__init__( model=model, client_type=client_type, base_url=base_url, - max_retries=max_retries, **kwargs, ) diff --git a/nemo_skills/inference/model/responses.py b/nemo_skills/inference/model/responses.py deleted file mode 100644 index 93350e125e..0000000000 --- a/nemo_skills/inference/model/responses.py +++ /dev/null @@ -1,375 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import logging - -from openai import AsyncOpenAI, OpenAI - -from nemo_skills.utils import get_logger_name - -from .base import BaseModel -from .context_retry import with_context_retry - -LOG = logging.getLogger(get_logger_name(__file__)) - - -class ResponsesModel(BaseModel): - """Model implementation using OpenAI Responses API via LiteLLM. - - This model uses the responses API endpoint instead of chat completions, - which is useful for models like gpt-oss that support the responses format. - """ - - MODEL_PROVIDER = "openai" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - # Initialize OpenAI clients for responses API - self.client = OpenAI(base_url=self.base_url, api_key=self.litellm_kwargs["api_key"]) - self.async_client = AsyncOpenAI(base_url=self.base_url, api_key=self.litellm_kwargs["api_key"]) - - def _build_completion_request_params( - self, - prompt: str, - tokens_to_generate: int = 512, - temperature: float = 0.0, - top_p: float = 0.95, - top_k: int = -1, - min_p: float = 0.0, - repetition_penalty: float = 1.0, - random_seed: int = None, - top_logprobs: int | None = None, - timeout: int | None = None, - stop_phrases: list[str] | None = None, - stream: bool = False, - reasoning_effort: str | None = None, - extra_body: dict = None, - tools: list[dict] | None = None, - ) -> dict: - """Build request parameters for responses API with string input.""" - request = { - "input": prompt, - "max_output_tokens": tokens_to_generate, - "temperature": temperature, - "top_p": top_p, - "stream": stream, - "extra_body": { - "seed": random_seed, - "reasoning_effort": reasoning_effort, - "timeout": timeout, - "stop": stop_phrases, - "top_logprobs": top_logprobs, - "top_k": top_k, - "min_p": min_p, - "repetition_penalty": repetition_penalty, - **(extra_body or {}), - }, - } - - # Only include tools if they are provided - if tools is not None: - request["tools"] = tools - - return request - - def _build_chat_request_params( - self, - messages: list[dict], - stream: bool, - tokens_to_generate: int = 512, - temperature: float = 0.0, - top_p: float = 0.95, - top_k: int = -1, - min_p: float = 0.0, - repetition_penalty: float = 1.0, - random_seed: int = 0, - stop_phrases: list[str] | None = None, - timeout: int | None = None, - top_logprobs: int | None = None, - reasoning_effort: str | None = None, - tools: list[dict] | None = None, - extra_body: dict = None, - ) -> dict: - """Build request parameters for responses API with messages input.""" - request = { - "input": messages, - "max_output_tokens": tokens_to_generate, - "temperature": temperature, - "top_p": top_p, - "stream": stream, - "extra_body": { - "seed": random_seed, - "reasoning_effort": reasoning_effort, - "timeout": timeout, - "stop": stop_phrases, - "top_logprobs": top_logprobs, - "top_k": top_k, - "min_p": min_p, - "repetition_penalty": repetition_penalty, - **(extra_body or {}), - }, - } - - # Only include tools if they are provided - if tools is not None: - request["tools"] = tools - - return request - - def _parse_responses_response(self, response, include_response: bool = False, **kwargs) -> dict: - """Parse responses API response into standard format.""" - - result = {"generation": "", "num_generated_tokens": 0} - - # Get token usage - if hasattr(response, "usage"): - result["num_generated_tokens"] = getattr(response.usage, "output_tokens", 0) - - # Check for tool calls in the output array - tool_calls = [] - reasoning_content = "" - generation_text = "" - - if hasattr(response, "output") and response.output: - for output_item in response.output: - # Handle reasoning content - if hasattr(output_item, "type") and output_item.type == "reasoning": - if hasattr(output_item, "content") and output_item.content: - for content_item in output_item.content: - if hasattr(content_item, "text"): - reasoning_content += content_item.text + "\n" - - # Handle function calls - elif hasattr(output_item, "type") and output_item.type == "function_call": - tool_calls.append(output_item) - - # Handle message content - elif hasattr(output_item, "type") and output_item.type == "message": - if hasattr(output_item, "content") and output_item.content: - for content_item in output_item.content: - if hasattr(content_item, "text"): - generation_text += content_item.text - - # Set the appropriate response fields - if tool_calls: - result["tool_calls"] = tool_calls - result["generation"] = "" # No text generation when there are tool calls - else: - result["generation"] = generation_text - - # Add reasoning content if available - if reasoning_content: - result["reasoning_content"] = reasoning_content.strip() - - # Add finish reason if available - if hasattr(response, "status"): - result["finish_reason"] = response.status - - # Add serialized output for conversation history - result["serialized_output"] = self._serialize_response_output(response) - - if include_response: - result["response"] = response - - return result - - def _serialize_response_output(self, response) -> list[dict]: - """Serialize response output objects using model_dump() for conversation history.""" - serialized_output = [] - - if hasattr(response, "output") and response.output: - for output_item in response.output: - try: - # Try to use model_dump() method if available (Pydantic models) - if hasattr(output_item, "model_dump"): - serialized_output.append(output_item.model_dump()) - # Fallback to dict conversion - elif hasattr(output_item, "__dict__"): - serialized_output.append(output_item.__dict__) - # Last resort: convert to string representation - else: - serialized_output.append({"content": str(output_item), "type": "unknown"}) - except Exception as e: - LOG.warning(f"Failed to serialize output item: {e}") - # Fallback serialization - serialized_output.append({"content": str(output_item), "type": "error", "error": str(e)}) - - return serialized_output - - @with_context_retry - async def generate_async( - self, - prompt: str | list[dict], - tokens_to_generate: int | None = None, - temperature: float = 0.0, - top_p: float = 0.95, - top_k: int = -1, - min_p: float = 0.0, - repetition_penalty: float = 1.0, - random_seed: int = None, - stop_phrases: list[str] | None = None, - top_logprobs: int | None = None, - timeout: float | int | None = 14400, - remove_stop_phrases: bool = True, - stream: bool = False, - reasoning_effort: str | None = None, - tools: list[dict] | None = None, - include_response: bool = False, - extra_body: dict = None, - ) -> dict: - """Generate response using responses API.""" - - # Check tool calls are a list of dict - if tools is not None: - for tool in tools: - if not isinstance(tool, dict): - raise ValueError(f"Tool must be a dictionary, got {type(tool)}") - - kwargs = { - "tokens_to_generate": tokens_to_generate, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "repetition_penalty": repetition_penalty, - "random_seed": random_seed, - "stop_phrases": stop_phrases, - "top_logprobs": top_logprobs, - "timeout": timeout, - "reasoning_effort": reasoning_effort, - "tools": tools, - "extra_body": extra_body, - } - - request_params = self._build_request_params(prompt=prompt, stream=stream, **kwargs) - - # Use OpenAI client for responses API - LOG.info(f"Making responses API call with params: {request_params}") - LOG.info(f"Full litellm_kwargs: {self.litellm_kwargs}") - LOG.info(f"Model name from litellm_kwargs: {self.litellm_kwargs['model']}") - LOG.info(f"Base URL: {self.base_url}") - - # Use the original model name (without litellm prefix) for OpenAI client - model_name = self.model_name_or_path # Just gpt-oss-20b - LOG.info(f"About to call responses.create with model: {model_name}") - - response = await self.async_client.responses.create(model=model_name, **request_params) - - LOG.info(f"Response from server: {response}") - LOG.info(f"Response type: {type(response)}") - - if stream: - # Handle streaming responses - return self._stream_responses_chunks_async(response) - else: - result = self._parse_responses_response(response, include_response=include_response, **kwargs) - self._maybe_apply_stop_phrase_removal(result, remove_stop_phrases, stop_phrases) - return result - - @with_context_retry - def generate_sync( - self, - prompt: str | list[dict], - tokens_to_generate: int | None = None, - temperature: float = 0.0, - top_p: float = 0.95, - top_k: int = -1, - min_p: float = 0.0, - repetition_penalty: float = 1.0, - random_seed: int = None, - stop_phrases: list[str] | None = None, - top_logprobs: int | None = None, - timeout: float | int | None = 14400, - remove_stop_phrases: bool = True, - stream: bool = False, - reasoning_effort: str | None = None, - tools: list[dict] | None = None, - include_response: bool = False, - extra_body: dict = None, - ) -> dict: - """Synchronous version of generate using responses API.""" - - # Check tool calls are a list of dict - if tools is not None: - for tool in tools: - if not isinstance(tool, dict): - raise ValueError(f"Tool must be a dictionary, got {type(tool)}") - - kwargs = { - "tokens_to_generate": tokens_to_generate, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "repetition_penalty": repetition_penalty, - "random_seed": random_seed, - "stop_phrases": stop_phrases, - "top_logprobs": top_logprobs, - "timeout": timeout, - "reasoning_effort": reasoning_effort, - "tools": tools, - "extra_body": extra_body, - } - - request_params = self._build_request_params(prompt=prompt, stream=stream, **kwargs) - - # Use OpenAI client for responses API - LOG.info(f"Making sync responses API call with params: {request_params}") - LOG.info(f"Model name: {self.litellm_kwargs['model']}") - - # Use the original model name (without litellm prefix) for OpenAI client - model_name = self.model_name_or_path # Just gpt-oss-20b - - response = self.client.responses.create(model=model_name, **request_params) - - LOG.info(f"Response from server: {response}") - LOG.info(f"Response type: {type(response)}") - - if stream: - # Handle streaming responses - return self._stream_responses_chunks_sync(response) - else: - result = self._parse_responses_response(response, include_response=include_response, **kwargs) - self._maybe_apply_stop_phrase_removal(result, remove_stop_phrases, stop_phrases) - return result - - def _stream_responses_chunks_sync(self, response): - """Synchronous version of stream responses chunks.""" - for chunk in response: - result = self._process_responses_chunk(chunk) - if result: - yield result - - async def _stream_responses_chunks_async(self, response): - """Async version of stream responses chunks.""" - async for chunk in response: - result = self._process_responses_chunk(chunk) - if result: - yield result - - def _process_responses_chunk(self, chunk): - """Process a single responses API chunk.""" - # This will depend on the actual streaming format from responses API - # For now, implement basic text streaming - if hasattr(chunk, "output_text"): - return {"generation": chunk.output_text or ""} - elif hasattr(chunk, "output") and chunk.output: - if isinstance(chunk.output, list) and len(chunk.output) > 0: - first_output = chunk.output[0] - if hasattr(first_output, "text"): - return {"generation": first_output.text or ""} - - # Fallback - return {"generation": ""} From 4e3a5a44e561cdd3d18096ea24729b416893b1d9 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Thu, 25 Sep 2025 14:53:28 -0700 Subject: [PATCH 12/19] MAINT switch to use_responses_api parameter Signed-off-by: George Armstrong MAINT switch to use_responses_api parameter Signed-off-by: George Armstrong --- nemo_skills/inference/generate.py | 13 ++++--- nemo_skills/inference/model/__init__.py | 43 +++++++++++++++++++----- nemo_skills/inference/model/base.py | 37 ++++++++++++++------ nemo_skills/inference/model/openai.py | 4 +-- nemo_skills/inference/model/tool_call.py | 2 +- 5 files changed, 73 insertions(+), 26 deletions(-) diff --git a/nemo_skills/inference/generate.py b/nemo_skills/inference/generate.py index 40beacafbb..09d06e77f6 100644 --- a/nemo_skills/inference/generate.py +++ b/nemo_skills/inference/generate.py @@ -92,8 +92,8 @@ class GenerateSolutionsConfig: # Inference server configuration {server_params} server: dict = field(default_factory=dict) - # Client type for API calls (chat_completion, responses, completion) - client_type: str = "chat_completion" + # Use responses API instead of chat completion API + use_responses_api: bool = False # Sandbox configuration {sandbox_params} sandbox: dict = field(default_factory=dict) # Prompt configuration - path to yaml files @@ -341,19 +341,22 @@ def setup_llm(self): if self.cfg.code_execution: llm = get_code_execution_model( - **self.cfg.server, client_type=self.cfg.client_type, tokenizer=self.tokenizer, sandbox=self.sandbox + **self.cfg.server, + use_responses_api=self.cfg.use_responses_api, + tokenizer=self.tokenizer, + sandbox=self.sandbox, ) elif self.cfg.tool_modules is not None: llm = get_tool_calling_model( **self.cfg.server, - client_type=self.cfg.client_type, + use_responses_api=self.cfg.use_responses_api, tool_modules=self.cfg.tool_modules, tool_overrides=self.cfg.tool_overrides, tokenizer=self.tokenizer, additional_config={"sandbox": self.cfg.sandbox}, ) else: - llm = get_model(**self.cfg.server, client_type=self.cfg.client_type, tokenizer=self.tokenizer) + llm = get_model(**self.cfg.server, use_responses_api=self.cfg.use_responses_api, tokenizer=self.tokenizer) if self.cfg.parallel_thinking.mode is not None: # We don't want to override these key variables which overlap with self.cfg diff --git a/nemo_skills/inference/model/__init__.py b/nemo_skills/inference/model/__init__.py index 9e0e5c7d88..dfc7626357 100644 --- a/nemo_skills/inference/model/__init__.py +++ b/nemo_skills/inference/model/__init__.py @@ -50,20 +50,36 @@ } -def get_model(server_type, client_type="chat_completion", tokenizer=None, **kwargs): - """A helper function to make it easier to set server through cmd.""" +def get_model(server_type, use_responses_api=False, tokenizer=None, **kwargs): + """A helper function to make it easier to set server through cmd. + + Args: + server_type: Type of server to use (e.g., 'openai', 'vllm', etc.) + use_responses_api: Whether to use responses API instead of chat completion API + tokenizer: Tokenizer to use + **kwargs: Additional arguments to pass to the model + """ model_class = models[server_type.lower()] if server_type == "trtllm" and kwargs.get("enable_soft_fail", False): if kwargs.get("context_limit_retry_strategy", None) is not None: raise ValueError("context_limit_retry_strategy is not supported for trtllm") - return model_class(client_type=client_type, tokenizer=tokenizer, **kwargs) + return model_class(use_responses_api=use_responses_api, tokenizer=tokenizer, **kwargs) def get_code_execution_model( - server_type, client_type="chat_completion", tokenizer=None, code_execution=None, sandbox=None, **kwargs + server_type, use_responses_api=False, tokenizer=None, code_execution=None, sandbox=None, **kwargs ): - """A helper function to make it easier to set server through cmd.""" - model = get_model(server_type=server_type, client_type=client_type, tokenizer=tokenizer, **kwargs) + """A helper function to make it easier to set server through cmd. + + Args: + server_type: Type of server to use (e.g., 'openai', 'vllm', etc.) + use_responses_api: Whether to use responses API instead of chat completion API + tokenizer: Tokenizer to use + code_execution: Code execution configuration + sandbox: Sandbox to use for code execution + **kwargs: Additional arguments to pass to the model + """ + model = get_model(server_type=server_type, use_responses_api=use_responses_api, tokenizer=tokenizer, **kwargs) if code_execution is None: code_execution = {} code_execution_config = CodeExecutionConfig(**code_execution) @@ -96,15 +112,26 @@ def get_parallel_thinking_model( def get_tool_calling_model( model, - client_type="chat_completion", + use_responses_api=False, tokenizer=None, additional_config=None, tool_modules: list[str] | None = None, tool_overrides: dict | None = None, **kwargs, ): + """A helper function to create a tool calling model. + + Args: + model: Model name/path or model instance + use_responses_api: Whether to use responses API instead of chat completion API + tokenizer: Tokenizer to use + additional_config: Additional configuration + tool_modules: List of tool modules to use + tool_overrides: Tool overrides + **kwargs: Additional arguments to pass to the model + """ if isinstance(model, str): - model = get_model(model=model, client_type=client_type, tokenizer=tokenizer, **kwargs) + model = get_model(model=model, use_responses_api=use_responses_api, tokenizer=tokenizer, **kwargs) return ToolCallingWrapper( model, tool_modules=tool_modules, diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index 666bead3e4..3237f507c6 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -336,14 +336,23 @@ class BaseModel: """Base model class for handling requests to the inference server. Args: - host: Optional[str] = '127.0.0.1' - Host of the inference server. - port: Optional[str] = '5000' - Port of the inference server. - Only required if handle_code_execution is True. - ssh_server: Optional[str] = None - SSH server for tunneling requests. + model: str - Model name or path to use for inference. + use_responses_api: bool = False - Whether to use responses API instead of chat completion API. + tokenizer: str | None = None - Tokenizer to use for the model. + api_key: str | None = None - API key for authentication. + api_key_env_var: str | None = None - Environment variable name containing API key. + base_url: str | None = None - Base URL for the API server. + use_v1_endpoint: bool = True - Whether to use v1 endpoint format. + host: str = '127.0.0.1' - Host of the inference server. + port: str = '5000' - Port of the inference server. + ssh_server: str | None = None - SSH server for tunneling requests. Useful if server is running on slurm cluster to which there is an ssh access Can also be specified through NEMO_SKILLS_SSH_SERVER env var. - ssh_key_path: Optional[str] = None - Path to the ssh key for tunneling. + ssh_key_path: str | None = None - Path to the ssh key for tunneling. Can also be specified through NEMO_SKILLS_SSH_KEY_PATH env var. + enable_soft_fail: bool = False - Enable soft failure handling. + context_limit_retry_strategy: str | None = None - Context limit retry strategy. + num_special_tokens_budget: int = 100 - Budget for special tokens. """ # Litellm provider name @@ -352,7 +361,7 @@ class BaseModel: def __init__( self, model: str, - client_type: str = "chat_completion", + use_responses_api: bool = False, tokenizer: str | None = None, api_key: str | None = None, api_key_env_var: str | None = None, @@ -369,7 +378,7 @@ def __init__( ): # Common model properties self.model_name_or_path = model - self.client_type = client_type + self.use_responses_api = use_responses_api self.server_host = host self.server_port = port @@ -396,10 +405,18 @@ def __init__( self.tokenizer = None # Initialize client handler LAST - if client_type not in CLIENT_HANDLERS: - raise ValueError(f"Unsupported client_type: {client_type}. Available: {list(CLIENT_HANDLERS.keys())}") + # Determine client type based on use_responses_api flag + if use_responses_api: + selected_client_type = "responses" + else: + selected_client_type = "chat_completion" + + if selected_client_type not in CLIENT_HANDLERS: + raise ValueError( + f"Unsupported client handler: {selected_client_type}. Available: {list(CLIENT_HANDLERS.keys())}" + ) - handler_class = CLIENT_HANDLERS[client_type] + handler_class = CLIENT_HANDLERS[selected_client_type] self.client_handler = handler_class(self) # Public attribute def _setup_ssh_tunnel(self, ssh_server: str | None, ssh_key_path: str | None): diff --git a/nemo_skills/inference/model/openai.py b/nemo_skills/inference/model/openai.py index fdd27b7baa..7eec904f99 100644 --- a/nemo_skills/inference/model/openai.py +++ b/nemo_skills/inference/model/openai.py @@ -36,7 +36,7 @@ def __init__( model: str | None = None, base_url: str | None = None, max_retries: int = 3, - client_type: str = "chat_completion", + use_responses_api: bool = False, **kwargs, ): model = model or os.getenv("NEMO_SKILLS_OPENAI_MODEL") @@ -52,7 +52,7 @@ def __init__( super().__init__( model=model, - client_type=client_type, + use_responses_api=use_responses_api, base_url=base_url, **kwargs, ) diff --git a/nemo_skills/inference/model/tool_call.py b/nemo_skills/inference/model/tool_call.py index b787dd58d3..9eb81984c6 100644 --- a/nemo_skills/inference/model/tool_call.py +++ b/nemo_skills/inference/model/tool_call.py @@ -71,7 +71,7 @@ def __init__( def _setup_adapters(self): """Set up adapters based on client type.""" # Use client_type instead of model instance checks - if self.model.client_type == "responses": + if self.model.use_responses_api: # Responses API model - uses flatter tool schema format self.schema_adapter = ResponsesSchemaAdapter() self.call_interpreter = ResponsesCallInterpreter() From 8428242a3c9891d18048451161a4076603840e81 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Thu, 25 Sep 2025 15:21:01 -0700 Subject: [PATCH 13/19] MAINT cleanup Signed-off-by: George Armstrong --- nemo_skills/inference/model/base.py | 29 +++++++++------------------ nemo_skills/inference/model/openai.py | 4 +--- 2 files changed, 11 insertions(+), 22 deletions(-) diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index 3237f507c6..0f2827ffc1 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -20,6 +20,7 @@ import httpx import litellm import openai +from openai import AsyncOpenAI, OpenAI from nemo_skills.utils import get_logger_name @@ -43,11 +44,10 @@ def setup_clients(self): def get_supported_params(self) -> set: """Public method to get supported parameters - override in subclasses""" - raise NotImplementedError + raise NotImplementedError() def extract_and_validate_params(self, **kwargs) -> dict: """Public method for parameter extraction and validation""" - LOG.info(f"extract_and_validate_params called with kwargs: {list(kwargs.keys())}") supported = self.get_supported_params() # Let model filter/restrict parameters @@ -94,11 +94,11 @@ def build_request_structure(self, prompt, params: dict) -> dict: def build_chat_request_structure(self, messages: list, params: dict) -> dict: """Public method for chat request structure - override in subclasses""" - raise NotImplementedError + raise NotImplementedError() def build_completion_request_structure(self, prompt: str, params: dict) -> dict: """Public method for completion request structure - override in subclasses""" - raise NotImplementedError + raise NotImplementedError() async def call_api_async(self, prompt, **kwargs): """Public method for async API calls""" @@ -126,15 +126,15 @@ def call_api_sync(self, prompt, **kwargs): async def make_async_call(self, request: dict, prompt): """Public method for making async API calls - override in subclasses""" - raise NotImplementedError + raise NotImplementedError() def make_sync_call(self, request: dict, prompt): """Public method for making sync API calls - override in subclasses""" - raise NotImplementedError + raise NotImplementedError() def parse_response(self, response, **kwargs) -> dict: """Public method for parsing responses - override in subclasses""" - raise NotImplementedError + raise NotImplementedError() class ChatCompletionHandler(BaseClientHandler): @@ -259,8 +259,6 @@ def get_supported_params(self) -> set: def setup_clients(self): """Setup OpenAI clients directly""" - from openai import AsyncOpenAI, OpenAI - self.sync_client = OpenAI(base_url=self.model.base_url, api_key=self.model.api_key) self.async_client = AsyncOpenAI(base_url=self.model.base_url, api_key=self.model.api_key) @@ -345,6 +343,7 @@ class BaseModel: use_v1_endpoint: bool = True - Whether to use v1 endpoint format. host: str = '127.0.0.1' - Host of the inference server. port: str = '5000' - Port of the inference server. + max_retries: int = 3 - Maximum number of retries for API calls. ssh_server: str | None = None - SSH server for tunneling requests. Useful if server is running on slurm cluster to which there is an ssh access Can also be specified through NEMO_SKILLS_SSH_SERVER env var. @@ -369,6 +368,7 @@ def __init__( use_v1_endpoint: bool = True, host: str = "127.0.0.1", port: str = "5000", + max_retries: int = 3, ssh_server: str | None = None, ssh_key_path: str | None = None, # Context limit retry config variables @@ -379,6 +379,7 @@ def __init__( # Common model properties self.model_name_or_path = model self.use_responses_api = use_responses_api + self.max_retries = max_retries self.server_host = host self.server_port = port @@ -533,18 +534,8 @@ def parse_completion_response(self, response, **kwargs) -> dict: def parse_responses_response(self, response, **kwargs) -> dict: """Public method for parsing responses API responses""" - # Move implementation from ResponsesModel here result = {"generation": "", "num_generated_tokens": 0} - # Debug logging to understand response structure - LOG.debug(f"Responses API response type: {type(response)}") - LOG.debug(f"Response has usage: {hasattr(response, 'usage')}") - if hasattr(response, "usage"): - LOG.debug(f"Usage object: {response.usage}") - LOG.debug(f"Usage type: {type(response.usage)}") - if response.usage: - LOG.debug(f"Usage attributes: {dir(response.usage)}") - # Get token usage - ensure it's always an integer if hasattr(response, "usage") and response.usage: tokens = getattr(response.usage, "output_tokens", None) diff --git a/nemo_skills/inference/model/openai.py b/nemo_skills/inference/model/openai.py index 7eec904f99..5879985d2e 100644 --- a/nemo_skills/inference/model/openai.py +++ b/nemo_skills/inference/model/openai.py @@ -47,13 +47,11 @@ def __init__( if base_url is None: base_url = os.getenv("NEMO_SKILLS_OPENAI_BASE_URL", f"http://{host}:{port}/v1") - # Store max_retries for client handlers to use - self.max_retries = max_retries - super().__init__( model=model, use_responses_api=use_responses_api, base_url=base_url, + max_retries=max_retries, **kwargs, ) From 51fca89eb943100c3ff32f72eb3d20bd2fb740ae Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Fri, 26 Sep 2025 11:02:07 -0700 Subject: [PATCH 14/19] Revert server type changes Signed-off-by: George Armstrong --- nemo_skills/pipeline/utils/server.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nemo_skills/pipeline/utils/server.py b/nemo_skills/pipeline/utils/server.py index 0946b0bd98..f3b9b3ea55 100644 --- a/nemo_skills/pipeline/utils/server.py +++ b/nemo_skills/pipeline/utils/server.py @@ -123,7 +123,7 @@ def get_server_command( # check if the model path is mounted if not vllm, sglang, or trtllm; # vllm, sglang, and trtllm can also pass model name as "model_path" so we need special processing - if server_type not in ["vllm", "sglang", "trtllm", "responses"]: + if server_type not in ["vllm", "sglang", "trtllm"]: check_if_mounted(cluster_config, model_path) # the model path will be mounted, so generally it will start with / @@ -152,7 +152,7 @@ def get_server_command( f" --micro-batch-size 1 " # that's a training argument, ignored here, but required to specify.. f" {server_args} " ) - elif server_type in ["vllm", "responses"]: + elif server_type == "vllm": server_entrypoint = server_entrypoint or "-m nemo_skills.inference.server.serve_vllm" start_vllm_cmd = ( f"python3 {server_entrypoint} " From 3ea3bb648a6f0535a636eb174489cb6d087146a9 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Fri, 26 Sep 2025 11:54:28 -0700 Subject: [PATCH 15/19] ENH error if no serialized output Signed-off-by: George Armstrong --- nemo_skills/mcp/adapters.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/nemo_skills/mcp/adapters.py b/nemo_skills/mcp/adapters.py index f5f398e163..18f8e82695 100644 --- a/nemo_skills/mcp/adapters.py +++ b/nemo_skills/mcp/adapters.py @@ -176,11 +176,7 @@ class ResponsesConversationManager(ConversationManager): def add_assistant_response(self, conversation: List[Dict[str, Any]], response: Dict[str, Any]) -> None: """Add an assistant response to the conversation using serialized output.""" # Use the serialized output from the responses API - if "serialized_output" in response: - conversation.extend(response["serialized_output"]) - else: - # Fallback to basic message format - conversation.append({"role": "assistant", "content": response["generation"]}) + conversation.extend(response["serialized_output"]) def add_tool_results( self, From f8c636716aee0ceefdcfbd129489ae50de01bdf7 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Fri, 26 Sep 2025 11:55:00 -0700 Subject: [PATCH 16/19] ENH improve non default arg errors Signed-off-by: George Armstrong --- nemo_skills/inference/model/base.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index 0f2827ffc1..ab752eb295 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -61,13 +61,9 @@ def extract_and_validate_params(self, **kwargs) -> dict: for param_name in provided: if param_name not in supported: - # Check if this parameter is set to its default value - default_value = getattr(self.defaults, param_name, None) - provided_value = kwargs[param_name] - - # Only consider it unsupported if it's not the default value - if default_value is None or provided_value != default_value: - unsupported_non_default.add(param_name) + if hasattr(self.defaults, param_name) and kwargs[param_name] == getattr(self.defaults, param_name): + continue # Default value is allowed even if unsupported + unsupported_non_default.add(param_name) if unsupported_non_default: raise ValueError( From ab7ad5459440d68fa8aced8587c675765bccd97c Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Fri, 26 Sep 2025 11:57:48 -0700 Subject: [PATCH 17/19] Apply suggestion from @coderabbitai[bot] Moves timeout to the main body parameters Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: George Armstrong --- nemo_skills/inference/model/base.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index ab752eb295..8ca2f44033 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -269,6 +269,14 @@ def build_chat_request_structure(self, messages: list, params: dict) -> dict: "stream": params["stream"], } + # Explicitly pass timeout to client, not inside extra_body + if params["timeout"] is not None: + request["timeout"] = params["timeout"] + + # Temporary guard until streaming parsing for Responses is implemented + if params["stream"]: + raise ValueError("stream=True is not supported for ResponsesHandler yet.") + # Only include tools if they are provided if params["tools"] is not None: request["tools"] = params["tools"] @@ -279,8 +287,6 @@ def build_chat_request_structure(self, messages: list, params: dict) -> dict: extra_body["seed"] = params["random_seed"] if params["reasoning_effort"] is not None: extra_body["reasoning_effort"] = params["reasoning_effort"] - if params["timeout"] is not None: - extra_body["timeout"] = params["timeout"] if params["stop_phrases"] is not None: extra_body["stop"] = params["stop_phrases"] if params["top_logprobs"] is not None: @@ -301,6 +307,8 @@ def build_chat_request_structure(self, messages: list, params: dict) -> dict: return request + return request + def build_completion_request_structure(self, prompt: str, params: dict) -> dict: """Responses API doesn't support completion - raise error""" raise ValueError("ResponsesHandler only supports message lists, not string prompts") From 8e137b8e867a0dfc49d13531409620b347bae861 Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Fri, 26 Sep 2025 13:03:12 -0700 Subject: [PATCH 18/19] log args Signed-off-by: George Armstrong --- nemo_skills/inference/model/tool_call.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_skills/inference/model/tool_call.py b/nemo_skills/inference/model/tool_call.py index 9eb81984c6..93e4c1b65b 100644 --- a/nemo_skills/inference/model/tool_call.py +++ b/nemo_skills/inference/model/tool_call.py @@ -95,7 +95,7 @@ async def _execute_tool_call(self, tool_call, request_id: str): try: tool_args = json.loads(tool_args) except json.decoder.JSONDecodeError as e: - LOG.exception(e) + LOG.exception(f"Failed to parse tool arguments {tool_args}: {e}") return {"error": "Tool argument parsing failed."} ## TODO(sanyamk): Only exceptions related to tool execution here, all others must fail. From 39b947ef000fe73db74561dcd78411f035c81c7f Mon Sep 17 00:00:00 2001 From: George Armstrong Date: Fri, 26 Sep 2025 14:11:39 -0700 Subject: [PATCH 19/19] Apply suggestion from @coderabbitai[bot] Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: George Armstrong --- nemo_skills/inference/model/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index 8ca2f44033..a42f65add0 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -394,7 +394,7 @@ def __init__( self.base_url = self._setup_base_url(base_url, use_v1_endpoint) # API key resolution (general) - self.api_key = self._resolve_api_key(api_key, api_key_env_var, base_url) + self.api_key = self._resolve_api_key(api_key, api_key_env_var, self.base_url) # Context retry config (general) self.context_limit_retry_config = ContextLimitRetryConfig(