Add Responses API

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

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,