feat: add support for structured outputs in OpenRouterChatGenerator (#2406)

* Add support

* Updates

* Update workflow

* Fix license

* Fix linting

* Update python version

* Update files

* Fix linting

* Update chat_generator.py

* Fix tools

* Update haystack version

* Updates

Author: Amnah199

integrations/openrouter/examples/openrouter_with_structured_outputs.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+# This example demonstrates how to use the OpenRouterChatGenerator component
+# with structured outputs.
+# To run this example, you will need to
+# set `OPENROUTER_API_KEY` environment variable
+
+from haystack.dataclasses import ChatMessage
+from pydantic import BaseModel
+
+from haystack_integrations.components.generators.openrouter import OpenRouterChatGenerator
+
+
+class NobelPrizeInfo(BaseModel):
+    recipient_name: str
+    award_year: int
+    category: str
+    achievement_description: str
+    nationality: str
+
+
+chat_messages = [
+    ChatMessage.from_user(
+        "In 2021, American scientist David Julius received the Nobel Prize in"
+        " Physiology or Medicine for his groundbreaking discoveries on how the human body"
+        " senses temperature and touch."
+    )
+]
+component = OpenRouterChatGenerator(generation_kwargs={"response_format": NobelPrizeInfo})
+results = component.run(chat_messages)
+
+# print(results)

integrations/openrouter/pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: CPython",
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["haystack-ai>=2.13.1"]
+dependencies = ["haystack-ai>=2.19.0"]
 
 [project.urls]
 Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/openrouter#readme"
@@ -154,4 +154,4 @@ addopts = "--strict-markers"
 markers = [
   "integration: integration tests",
 ]
-log_cli = true
+log_cli = true

integrations/openrouter/src/haystack_integrations/components/generators/openrouter/chat/chat_generator.py

@@ -2,12 +2,12 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, Optional
 
 from haystack import component, default_to_dict, logging
 from haystack.components.generators.chat import OpenAIChatGenerator
 from haystack.dataclasses import ChatMessage, StreamingCallbackT
-from haystack.tools import Tool, Toolset, _check_duplicate_tool_names
+from haystack.tools import ToolsType, _check_duplicate_tool_names, flatten_tools_or_toolsets, serialize_tools_or_toolset
 from haystack.utils import serialize_callable
 from haystack.utils.auth import Secret
 
@@ -64,7 +64,7 @@ def __init__(
         streaming_callback: Optional[StreamingCallbackT] = None,
         api_base_url: Optional[str] = "https://openrouter.ai/api/v1",
         generation_kwargs: Optional[Dict[str, Any]] = None,
-        tools: Optional[Union[List[Tool], Toolset]] = None,
+        tools: Optional[ToolsType] = None,
         timeout: Optional[float] = None,
         extra_headers: Optional[Dict[str, Any]] = None,
         max_retries: Optional[int] = None,
@@ -98,6 +98,14 @@ def __init__(
                 events as they become available, with the stream terminated by a data: [DONE] message.
             - `safe_prompt`: Whether to inject a safety prompt before all conversations.
             - `random_seed`: The seed to use for random sampling.
+            - `response_format`: A JSON schema or a Pydantic model that enforces the structure of the model's response.
+                If provided, the output will always be validated against this
+                format (unless the model returns a tool call).
+                For details, see the [OpenAI Structured Outputs documentation](https://platform.openai.com/docs/guides/structured-outputs).
+                Notes:
+                - This parameter accepts Pydantic models and JSON schemas for latest models starting from GPT-4o.
+                - For structured outputs with streaming,
+                  the `response_format` must be a JSON schema and not a Pydantic model.
         :param tools:
             A list of tools or a Toolset for which the model can prepare calls. This parameter can accept either a
             list of `Tool` objects or a `Toolset` instance.
@@ -148,7 +156,7 @@ def to_dict(self) -> Dict[str, Any]:
             api_base_url=self.api_base_url,
             generation_kwargs=self.generation_kwargs,
             api_key=self.api_key.to_dict(),
-            tools=[tool.to_dict() for tool in self.tools] if self.tools else None,
+            tools=serialize_tools_or_toolset(self.tools),
             extra_headers=self.extra_headers,
             timeout=self.timeout,
             max_retries=self.max_retries,
@@ -158,46 +166,64 @@ def to_dict(self) -> Dict[str, Any]:
     def _prepare_api_call(
         self,
         *,
-        messages: List[ChatMessage],
+        messages: list[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
-        generation_kwargs: Optional[Dict[str, Any]] = None,
-        tools: Optional[Union[List[Tool], Toolset]] = None,
+        generation_kwargs: Optional[dict[str, Any]] = None,
+        tools: Optional[ToolsType] = None,
         tools_strict: Optional[bool] = None,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         # update generation kwargs by merging with the generation kwargs passed to the run method
         generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
         extra_headers = {**(self.extra_headers or {})}
 
+        is_streaming = streaming_callback is not None
+        num_responses = generation_kwargs.pop("n", 1)
+
+        if is_streaming and num_responses > 1:
+            msg = "Cannot stream multiple responses, please set n=1."
+            raise ValueError(msg)
+        response_format = generation_kwargs.pop("response_format", None)
+
         # adapt ChatMessage(s) to the format expected by the OpenAI API
         openai_formatted_messages = [message.to_openai_dict_format() for message in messages]
 
-        tools = tools or self.tools
-        if isinstance(tools, Toolset):
-            tools = list(tools)
+        flattened_tools = flatten_tools_or_toolsets(tools or self.tools)
         tools_strict = tools_strict if tools_strict is not None else self.tools_strict
-        _check_duplicate_tool_names(list(tools or []))
+        _check_duplicate_tool_names(flattened_tools)
 
         openai_tools = {}
-        if tools:
-            tool_definitions = [
-                {"type": "function", "function": {**t.tool_spec, **({"strict": tools_strict} if tools_strict else {})}}
-                for t in tools
-            ]
+        if flattened_tools:
+            tool_definitions = []
+            for t in flattened_tools:
+                function_spec = {**t.tool_spec}
+                if tools_strict:
+                    function_spec["strict"] = True
+                    function_spec["parameters"]["additionalProperties"] = False
+                tool_definitions.append({"type": "function", "function": function_spec})
             openai_tools = {"tools": tool_definitions}
 
-        is_streaming = streaming_callback is not None
-        num_responses = generation_kwargs.pop("n", 1)
-        if is_streaming and num_responses > 1:
-            msg = "Cannot stream multiple responses, please set n=1."
-            raise ValueError(msg)
-
-        return {
+        base_args = {
             "model": self.model,
-            "messages": openai_formatted_messages,  # type: ignore[arg-type] # openai expects list of specific message types
-            "stream": streaming_callback is not None,
+            "messages": openai_formatted_messages,
             "n": num_responses,
             **openai_tools,
-            "extra_body": {**generation_kwargs},
             "extra_headers": {**extra_headers},
-            "openai_endpoint": "create",
+            "extra_body": {**generation_kwargs},
         }
+
+        if response_format and not is_streaming:
+            # for structured outputs without streaming, we use openai's parse endpoint
+            # Note: `stream` cannot be passed to chat.completions.parse
+            # we pass a key `openai_endpoint` as a hint to the run method to use the parse endpoint
+            # this key will be removed before the API call is made
+            return {**base_args, "response_format": response_format, "openai_endpoint": "parse"}
+
+        # for structured outputs with streaming, we use openai's create endpoint
+        # we pass a key `openai_endpoint` as a hint to the run method to use the create endpoint
+        # this key will be removed before the API call is made
+        final_args = {**base_args, "stream": is_streaming, "openai_endpoint": "create"}
+
+        # We only set the response_format parameter if it's not None since None is not a valid value in the API.
+        if response_format:
+            final_args["response_format"] = response_format
+        return final_args

integrations/openrouter/tests/test_openrouter_chat_generator.py

@@ -1,3 +1,4 @@
+import json
 import os
 from datetime import datetime
 from unittest.mock import patch
@@ -16,10 +17,22 @@
 from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
 from openai.types.chat.chat_completion_chunk import ChoiceDelta, ChoiceDeltaToolCall, ChoiceDeltaToolCallFunction
 from openai.types.completion_usage import CompletionTokensDetails, CompletionUsage, PromptTokensDetails
+from pydantic import BaseModel
 
 from haystack_integrations.components.generators.openrouter.chat.chat_generator import OpenRouterChatGenerator
 
 
+class CalendarEvent(BaseModel):
+    event_name: str
+    event_date: str
+    event_location: str
+
+
+@pytest.fixture
+def calendar_event_model():
+    return CalendarEvent
+
+
 class CollectorCallback:
     """
     Callback to collect streaming chunks for testing purposes.
@@ -440,6 +453,41 @@ def test_pipeline_with_openrouter_chat_generator(self, tools):
             == results["tool_invoker"]["tool_messages"][0].tool_call_result.result
         )
 
+    @pytest.mark.skipif(
+        not os.environ.get("OPENROUTER_API_KEY", None),
+        reason="Export an env var called OPENROUTER_API_KEY containing the OpenRouter API key to run this test.",
+    )
+    @pytest.mark.integration
+    def test_live_run_with_response_format_json_schema(self):
+        response_schema = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": "CapitalCity",
+                "strict": True,
+                "schema": {
+                    "title": "CapitalCity",
+                    "type": "object",
+                    "properties": {
+                        "city": {"title": "City", "type": "string"},
+                        "country": {"title": "Country", "type": "string"},
+                    },
+                    "required": ["city", "country"],
+                    "additionalProperties": False,
+                },
+            },
+        }
+
+        chat_messages = [ChatMessage.from_user("What's the capital of France?")]
+        comp = OpenRouterChatGenerator(generation_kwargs={"response_format": response_schema})
+        results = comp.run(chat_messages)
+        assert len(results["replies"]) == 1
+        message: ChatMessage = results["replies"][0]
+        msg = json.loads(message.text)
+        assert "Paris" in msg["city"]
+        assert isinstance(msg["country"], str)
+        assert "France" in msg["country"]
+        assert message.meta["finish_reason"] == "stop"
+
     def test_serde_in_pipeline(self, monkeypatch):
         """
         Test serialization/deserialization of OpenRouterChatGenerator in a Pipeline,
@@ -539,6 +587,24 @@ def test_serde_in_pipeline(self, monkeypatch):
         assert loaded_generator.tools[0].description == generator.tools[0].description
         assert loaded_generator.tools[0].parameters == generator.tools[0].parameters
 
+    @pytest.mark.skipif(
+        not os.environ.get("OPENROUTER_API_KEY", None),
+        reason="Export an env var called OPENROUTER_API_KEY containing the OpenRouter API key to run this test.",
+    )
+    @pytest.mark.integration
+    def test_live_run_with_response_format_pydantic_model(self, calendar_event_model):
+        chat_messages = [
+            ChatMessage.from_user("The marketing summit takes place on October12th at the Hilton Hotel downtown.")
+        ]
+        component = OpenRouterChatGenerator(generation_kwargs={"response_format": calendar_event_model})
+        results = component.run(chat_messages)
+        assert len(results["replies"]) == 1
+        message: ChatMessage = results["replies"][0]
+        msg = json.loads(message.text)
+        assert "Marketing Summit" in msg["event_name"]
+        assert isinstance(msg["event_date"], str)
+        assert isinstance(msg["event_location"], str)
+
 
 class TestChatCompletionChunkConversion:
     def test_handle_stream_response(self):