diff --git a/nemo_skills/inference/generate.py b/nemo_skills/inference/generate.py index 4361d394ca..09d06e77f6 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) + # 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 @@ -338,17 +340,23 @@ 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, + 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, + 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, 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 6ecfaa6e25..dfc7626357 100644 --- a/nemo_skills/inference/model/__init__.py +++ b/nemo_skills/inference/model/__init__.py @@ -50,18 +50,36 @@ } -def get_model(server_type, 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(tokenizer=tokenizer, **kwargs) + return model_class(use_responses_api=use_responses_api, tokenizer=tokenizer, **kwargs) -def get_code_execution_model(server_type, 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) +def get_code_execution_model( + 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. + + 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) @@ -94,14 +112,26 @@ def get_parallel_thinking_model( def get_tool_calling_model( model, + 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, 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 94aa096450..a42f65add0 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 @@ -29,18 +30,332 @@ 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, but only for parameters that differ from defaults + provided = set(kwargs.keys()) + unsupported_non_default = set() + + for param_name in provided: + if param_name not in supported: + 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( + f"Unsupported parameters for {self.__class__.__name__}: {unsupported_non_default}. Supported: {sorted(supported)}" + ) + + # Extract with defaults - include all provided parameters, even unsupported ones if they're default values + params = {} + 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: + 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=getattr(self.model, "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""" + request = { + "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"], + } + + # 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""" + request = { + "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"], + } + + # 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): + 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""" + 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""" + # 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"], + } + + # 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"] + + # 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["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 + + 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") + + 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. 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. + 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. - 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 @@ -49,14 +364,15 @@ class BaseModel: def __init__( self, model: str, + use_responses_api: bool = False, 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", + max_retries: int = 3, ssh_server: str | None = None, ssh_key_path: str | None = None, # Context limit retry config variables @@ -64,21 +380,55 @@ 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.use_responses_api = use_responses_api + self.max_retries = max_retries 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, self.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 + # 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[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): + """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 +449,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 + return base_url - 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" - - 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 @@ -143,7 +480,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( @@ -182,6 +519,119 @@ 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""" + result = {"generation": "", "num_generated_tokens": 0} + + # 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 = [] + 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 + + # 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): + """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 +670,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,6 +679,7 @@ async def generate_async( if not isinstance(tool, dict): raise ValueError(f"Tool must be a dictionary, got {type(tool)}") + # Build kwargs dict explicitly to avoid capturing unwanted local variables kwargs = { "tokens_to_generate": tokens_to_generate, "temperature": temperature, @@ -243,6 +694,9 @@ async def generate_async( "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 @@ -251,28 +705,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 +752,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,6 +762,7 @@ def generate_sync( if not isinstance(tool, dict): raise ValueError(f"Tool must be a dictionary, got {type(tool)}") + # Build kwargs dict explicitly to avoid capturing unwanted local variables kwargs = { "tokens_to_generate": tokens_to_generate, "temperature": temperature, @@ -334,26 +777,21 @@ def generate_sync( "reasoning_effort": reasoning_effort, "tools": tools, "extra_body": extra_body, + "remove_stop_phrases": remove_stop_phrases, + "stream": stream, + "include_response": include_response, } - 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) - 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..2d7a5218d7 --- /dev/null +++ b/nemo_skills/inference/model/defaults.py @@ -0,0 +1,96 @@ +# 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", + "remove_stop_phrases", + "include_response", +} + +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", + "remove_stop_phrases", + "include_response", +} + +COMPLETION_PARAMS = { + "tokens_to_generate", + "temperature", + "top_p", + "top_k", + "min_p", + "repetition_penalty", + "random_seed", + "stop_phrases", + "top_logprobs", + "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 9c7faa85d1..5879985d2e 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, + use_responses_api: bool = False, **kwargs, ): model = model or os.getenv("NEMO_SKILLS_OPENAI_MODEL") @@ -39,6 +49,7 @@ def __init__( super().__init__( model=model, + use_responses_api=use_responses_api, 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 bba2f64f9c..93e4c1b65b 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,25 @@ 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 client type.""" + # Use client_type instead of model instance checks + if self.model.use_responses_api: + # 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. @@ -75,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. @@ -103,7 +123,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 +153,18 @@ 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) + tool_results = await self._execute_tool_calls(tool_calls_message["tool_calls"], request_id=request_id) + + # 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 6e385d9757..18f8e82695 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,25 @@ 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_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_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.") + + # ============================== # ADAPTER IMPLEMENTATIONS # ============================== @@ -116,3 +135,112 @@ 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_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_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 (completion models need this) + if conversation and conversation[-1]["role"] == "assistant": + conversation[-1].update(tool_calls_message) + + # Add the formatted tool result messages directly - they're already in the right format + conversation.extend(formatted_results) + + +class ResponsesConversationManager(ConversationManager): + """Manages conversation history for responses API models.""" + + 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 + conversation.extend(response["serialized_output"]) + + def add_tool_results( + 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 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) + + +# ============================== +# 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 + ]