unique_toolkit 1.34.0__tar.gz → 1.35.0__tar.gz

{unique_toolkit-1.34.0 → unique_toolkit-1.35.0}/CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.35.0] - 2025-12-04
+- Add `LoopIterationRunner` abstraction and support for planning before every loop iteration.
+
+## [1.34.1] - 2025-12-02
+- Update code interpreter tool instructions.
+
 ## [1.34.0] - 2025-12-02
 - Add option to upload code interpreter generated files to the chat.
 

{unique_toolkit-1.34.0 → unique_toolkit-1.35.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 1.34.0
+Version: 1.35.0
 Summary: 
 License: Proprietary
 Author: Cedric Klinkert
@@ -121,6 +121,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.35.0] - 2025-12-04
+- Add `LoopIterationRunner` abstraction and support for planning before every loop iteration.
+
+## [1.34.1] - 2025-12-02
+- Update code interpreter tool instructions.
+
 ## [1.34.0] - 2025-12-02
 - Add option to upload code interpreter generated files to the chat.
 

{unique_toolkit-1.34.0 → unique_toolkit-1.35.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unique_toolkit"
-version = "1.34.0"
+version = "1.35.0"
 description = ""
 authors = [
     "Cedric Klinkert <cedric.klinkert@unique.ch>",

unique_toolkit-1.35.0/unique_toolkit/agentic/loop_runner/__init__.py

@@ -0,0 +1,19 @@
+from unique_toolkit.agentic.loop_runner.base import LoopIterationRunner
+from unique_toolkit.agentic.loop_runner.middleware import (
+    PlanningConfig,
+    PlanningMiddleware,
+    PlanningSchemaConfig,
+)
+from unique_toolkit.agentic.loop_runner.runners import (
+    BasicLoopIterationRunner,
+    BasicLoopIterationRunnerConfig,
+)
+
+__all__ = [
+    "LoopIterationRunner",
+    "PlanningConfig",
+    "PlanningMiddleware",
+    "PlanningSchemaConfig",
+    "BasicLoopIterationRunnerConfig",
+    "BasicLoopIterationRunner",
+]

unique_toolkit-1.35.0/unique_toolkit/agentic/loop_runner/_stream_handler_utils.py

@@ -0,0 +1,57 @@
+from typing import Any, Required
+
+from openai.types.chat import ChatCompletionNamedToolChoiceParam
+from typing_extensions import TypedDict
+
+from unique_toolkit import LanguageModelToolDescription
+from unique_toolkit.agentic.loop_runner.base import _LoopIterationRunnerKwargs
+from unique_toolkit.chat.functions import LanguageModelStreamResponse
+from unique_toolkit.chat.service import LanguageModelMessages
+from unique_toolkit.content import ContentChunk
+
+
+class _StreamingHandlerKwargs(TypedDict, total=False):
+    messages: Required[LanguageModelMessages]
+    model_name: Required[str]
+    tools: list[LanguageModelToolDescription]
+    content_chunks: list[ContentChunk]
+    start_text: str
+    debug_info: dict[str, Any]
+    temperature: float
+    tool_choice: ChatCompletionNamedToolChoiceParam
+    other_options: dict[str, Any]
+
+
+def _extract_streaming_kwargs(
+    kwargs: _LoopIterationRunnerKwargs,
+) -> _StreamingHandlerKwargs:
+    res = _StreamingHandlerKwargs(
+        messages=kwargs["messages"],
+        model_name=kwargs["model"].name,
+    )
+
+    for k in [
+        "tools",
+        "content_chunks",
+        "start_text",
+        "debug_info",
+        "temperature",
+        "other_options",
+    ]:
+        if k in kwargs:
+            res[k] = kwargs[k]
+
+    return res
+
+
+async def stream_response(
+    loop_runner_kwargs: _LoopIterationRunnerKwargs,
+    **kwargs,
+) -> LanguageModelStreamResponse:
+    streaming_handler = loop_runner_kwargs["streaming_handler"]
+    streaming_hander_kwargs = _extract_streaming_kwargs(loop_runner_kwargs)
+    streaming_hander_kwargs.update(**kwargs)
+
+    return await streaming_handler.complete_with_references_async(
+        **streaming_hander_kwargs
+    )

unique_toolkit-1.35.0/unique_toolkit/agentic/loop_runner/base.py

@@ -0,0 +1,38 @@
+from typing import Any, Protocol, Required, Unpack
+
+from openai.types.chat import ChatCompletionNamedToolChoiceParam
+from typing_extensions import TypedDict
+
+from unique_toolkit import LanguageModelToolDescription
+from unique_toolkit.chat.functions import LanguageModelStreamResponse
+from unique_toolkit.chat.service import LanguageModelMessages
+from unique_toolkit.content import ContentChunk
+from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.protocols.support import (
+    SupportCompleteWithReferences,
+)
+
+
+class _LoopIterationRunnerKwargs(TypedDict, total=False):
+    iteration_index: Required[int]
+    streaming_handler: Required[SupportCompleteWithReferences]
+    messages: Required[LanguageModelMessages]
+    model: Required[LanguageModelInfo]
+    tools: list[LanguageModelToolDescription]
+    content_chunks: list[ContentChunk]
+    start_text: str
+    debug_info: dict[str, Any]
+    temperature: float
+    tool_choices: list[ChatCompletionNamedToolChoiceParam]
+    other_options: dict[str, Any]
+
+
+class LoopIterationRunner(Protocol):
+    """
+    A loop iteration runner is responsible for running a single iteration of the loop, and returning the stream response for that iteration.
+    """
+
+    async def __call__(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse: ...

unique_toolkit-1.35.0/unique_toolkit/agentic/loop_runner/middleware/__init__.py

@@ -0,0 +1,7 @@
+from unique_toolkit.agentic.loop_runner.middleware.planning import (
+    PlanningConfig,
+    PlanningMiddleware,
+    PlanningSchemaConfig,
+)
+
+__all__ = ["PlanningConfig", "PlanningMiddleware", "PlanningSchemaConfig"]

unique_toolkit-1.35.0/unique_toolkit/agentic/loop_runner/middleware/planning/__init__.py

@@ -0,0 +1,9 @@
+from unique_toolkit.agentic.loop_runner.middleware.planning.planning import (
+    PlanningConfig,
+    PlanningMiddleware,
+)
+from unique_toolkit.agentic.loop_runner.middleware.planning.schema import (
+    PlanningSchemaConfig,
+)
+
+__all__ = ["PlanningMiddleware", "PlanningConfig", "PlanningSchemaConfig"]

unique_toolkit-1.35.0/unique_toolkit/agentic/loop_runner/middleware/planning/planning.py

@@ -0,0 +1,85 @@
+import json
+import logging
+from typing import Unpack
+
+from pydantic import BaseModel
+
+from unique_toolkit import LanguageModelService
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit.agentic.history_manager.history_manager import HistoryManager
+from unique_toolkit.agentic.loop_runner.base import (
+    LoopIterationRunner,
+    _LoopIterationRunnerKwargs,
+)
+from unique_toolkit.agentic.loop_runner.middleware.planning.schema import (
+    PlanningSchemaConfig,
+    get_planning_schema,
+)
+from unique_toolkit.agentic.tools.utils import failsafe_async
+from unique_toolkit.chat.service import LanguageModelStreamResponse
+from unique_toolkit.language_model import (
+    LanguageModelAssistantMessage,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class PlanningConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    planning_schema_config: PlanningSchemaConfig = PlanningSchemaConfig()
+
+
+class PlanningMiddleware(LoopIterationRunner):
+    def __init__(
+        self,
+        *,
+        loop_runner: LoopIterationRunner,
+        config: PlanningConfig,
+        llm_service: LanguageModelService,
+        history_manager: HistoryManager | None = None,
+    ) -> None:
+        self._config = config
+        self._loop_runner = loop_runner
+        self._history_manager = history_manager
+        self._llm_service = llm_service
+
+    @failsafe_async(failure_return_value=None, logger=_LOGGER)
+    async def _run_plan_step(
+        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
+    ) -> LanguageModelAssistantMessage | None:
+        planning_schema = get_planning_schema(self._config.planning_schema_config)
+
+        response = await self._llm_service.complete_async(
+            messages=kwargs["messages"],
+            model_name=kwargs["model"].name,
+            structured_output_model=planning_schema,
+            other_options=kwargs.get("other_options", {}),
+        )
+
+        if response.choices[0].message.parsed is None:
+            _LOGGER.info("Error parsing planning response")
+            return None
+
+        return LanguageModelAssistantMessage(
+            content=json.dumps(response.choices[0].message.parsed)
+        )
+
+    async def __call__(
+        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
+    ) -> LanguageModelStreamResponse:
+        assistant_message = await self._run_plan_step(**kwargs)
+
+        if assistant_message is None:
+            _LOGGER.info(
+                "Error executing planning step, proceeding without planning step"
+            )
+            return await self._loop_runner(**kwargs)
+
+        if self._history_manager is not None:
+            self._history_manager.add_assistant_message(assistant_message)
+
+        kwargs["messages"] = (
+            kwargs["messages"].builder().append(assistant_message).build()
+        )
+        return await self._loop_runner(**kwargs)

unique_toolkit-1.35.0/unique_toolkit/agentic/loop_runner/middleware/planning/schema.py

@@ -0,0 +1,82 @@
+import json
+from typing import Annotated, Any
+
+from pydantic import BaseModel, Field, RootModel, create_model
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+
+_PLANNING_SCHEMA_DESCRIPTION = """
+Think about the next step to take.
+
+Instructions:
+- Consider the user input and the context of the conversation.
+- Consider any previous tool calls, their results and the instructions related to the available tool calls.
+- Consider any failed tool calls.
+Goals:
+- Output a plan for the next step. It MUST be justified, meaning that you MUST explain why you choose to take this step.
+- You MUST recover from any failed tool calls.
+- You MUST explain what tool calls to call next and why.
+- If ready to answer the user, justify why you have gathered enough information/ tried all possible ways and failed.
+- If ready to answer the user, REMEMBER and mention any previous instructions you have in the history. This is a CRUCIAL step.
+
+IMPORTANT:
+- Tools will be available after the planning step.
+""".strip()
+
+_DEFAULT_PLANNING_PARAM_DESCRIPTION = """
+Next step description:
+- Decide what to do next.
+- Justify it THOROUGLY.
+""".strip()
+
+
+class DefaultPlanningSchemaConfig(BaseModel):
+    """
+    Configuration for the default planning schema, which is a simple json with a single field: "planning".
+    """
+
+    model_config = get_configuration_dict()
+
+    description: str = Field(
+        default=_PLANNING_SCHEMA_DESCRIPTION,
+        description="Description of the planning schema. This will correspond to the description of the model in the json schema.",
+    )
+    plan_param_description: str = Field(
+        default=_DEFAULT_PLANNING_PARAM_DESCRIPTION,
+        description="The description of the `planning` parameter.",
+    )
+
+
+class PlanningSchemaConfig(RootModel[DefaultPlanningSchemaConfig | str]):
+    model_config = get_configuration_dict()
+
+    root: (
+        Annotated[
+            DefaultPlanningSchemaConfig,
+            Field(
+                description="Configuration for the default planning schema, which is a simple json dict with a single `plan` field.",
+                title="Default Planning Schema",
+            ),
+        ]
+        | Annotated[
+            str,
+            Field(
+                description="Custom JSON Schema as string for the planning schema.",
+                title="Custom Planning Schema",
+            ),
+        ]
+    ) = Field(default=DefaultPlanningSchemaConfig())
+
+
+def get_planning_schema(config: PlanningSchemaConfig) -> dict[str, Any]:
+    if isinstance(config.root, DefaultPlanningSchemaConfig):
+        return create_model(
+            "Planning",
+            plan=(
+                str,
+                Field(description=config.root.plan_param_description),
+            ),
+            __doc__=config.root.description,
+        ).model_json_schema()
+
+    return json.loads(config.root)

unique_toolkit-1.35.0/unique_toolkit/agentic/loop_runner/runners/__init__.py

@@ -0,0 +1,6 @@
+from unique_toolkit.agentic.loop_runner.runners.basic import (
+    BasicLoopIterationRunner,
+    BasicLoopIterationRunnerConfig,
+)
+
+__all__ = ["BasicLoopIterationRunnerConfig", "BasicLoopIterationRunner"]

unique_toolkit-1.35.0/unique_toolkit/agentic/loop_runner/runners/basic.py

@@ -0,0 +1,94 @@
+import logging
+from typing import Unpack, override
+
+from pydantic import BaseModel
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit.agentic.loop_runner._stream_handler_utils import stream_response
+from unique_toolkit.agentic.loop_runner.base import (
+    LoopIterationRunner,
+    _LoopIterationRunnerKwargs,
+)
+from unique_toolkit.chat.functions import LanguageModelStreamResponse
+from unique_toolkit.protocols.support import (
+    ResponsesLanguageModelStreamResponse,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class BasicLoopIterationRunnerConfig(BaseModel):
+    model_config = get_configuration_dict()
+    max_loop_iterations: int
+
+
+class BasicLoopIterationRunner(LoopIterationRunner):
+    def __init__(self, config: BasicLoopIterationRunnerConfig) -> None:
+        self._config = config
+
+    async def _handle_last_iteration(
+        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
+    ) -> LanguageModelStreamResponse:
+        _LOGGER.info(
+            "Reached last iteration, removing tools and producing final response"
+        )
+
+        return await stream_response(
+            loop_runner_kwargs=kwargs,
+            tools=None,
+        )
+
+    async def _handle_normal_iteration(
+        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
+    ) -> LanguageModelStreamResponse:
+        _LOGGER.info("Running loop iteration %d", kwargs["iteration_index"])
+
+        return await stream_response(loop_runner_kwargs=kwargs)
+
+    async def _handle_forced_tools_iteration(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse:
+        assert "tool_choices" in kwargs
+
+        tool_choices = kwargs["tool_choices"]
+        _LOGGER.info("Forcing tools calls: %s", tool_choices)
+
+        responses: list[LanguageModelStreamResponse] = []
+
+        for opt in tool_choices:
+            responses.append(
+                await stream_response(
+                    loop_runner_kwargs=kwargs,
+                    tool_choice=opt,
+                )
+            )
+
+        # Merge responses and refs:
+        tool_calls = []
+        references = []
+        for r in responses:
+            if r.tool_calls:
+                tool_calls.extend(r.tool_calls)
+            references.extend(r.message.references)
+
+        response = responses[0]
+        response.tool_calls = tool_calls if len(tool_calls) > 0 else None
+        response.message.references = references
+
+        return response
+
+    @override
+    async def __call__(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse | ResponsesLanguageModelStreamResponse:
+        tool_choices = kwargs.get("tool_choices", [])
+        iteration_index = kwargs["iteration_index"]
+
+        if len(tool_choices) > 0 and iteration_index == 0:
+            return await self._handle_forced_tools_iteration(**kwargs)
+        elif iteration_index == self._config.max_loop_iterations - 1:
+            return await self._handle_last_iteration(**kwargs)
+        else:
+            return await self._handle_normal_iteration(**kwargs)

{unique_toolkit-1.34.0 → unique_toolkit-1.35.0}/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/config.py

@@ -11,27 +11,41 @@ DEFAULT_TOOL_DESCRIPTION = "Use this tool to run python code, e.g to generate pl
 
 DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT = """
 Use this tool to run python code, e.g to generate plots, process excel files, perform calculations, etc.
+
 Instructions:
-- All files uploaded to the chat are available in the code interpreter under the path `/mnt/data/<filename>
-- All files generated through code should be saved in the `/mnt/data` folder
 
-CRUCIAL Instructions for displaying images and files in the chat:
-Once files are generated in the `/mnt/data` folder you MUST reference them in the chat using markdown syntax in order to display them in the chat.
+Uploaded and generated files:
+- All files uploaded to the chat are available at the path `/mnt/data/<filename>`.
+- All files generated through code MUST be saved in the `/mnt/data` folder.
 
+CRUCIAL Instructions for displaying images and files in the chat:
+- Once files are generated in the `/mnt/data` folder you MUST reference them in the chat using markdown syntax in order to display them in the chat.
 WHENEVER you reference a generated file, you MUST use the following format:
 ```
-**Descriptive Title of Graph/Chart/File** (<- linebreak is important) (You must choose a good user friendly title, Other markdown syntax such as `#` can be used too)
+**Descriptive Title of Graph/Chart/File**  (linebreak is important. You must choose a good user friendly title, Other markdown syntax such as `#` can be used too)
 [*Generating your {Graph/Chart/File}…*](sandbox:/mnt/data/<filename>)
 ```
 IMPORTANT: Do NOT append a leading `!` even when displaying an image.
 Always use a line break between the title and the markdown!
-
 - Files with image file extensions are displayed directly in the chat, while other file extensions are shown as download links.
 - Not using syntax above will FAIL to show images to the user. 
 - YOU MUST use the syntax above to display images, otherwise the image will not be displayed in the chat.
 - Only the following file types are allowed to be uploaded to the platform, anything else will FAIL: PDF, DOCX, XLSX, PPTX, CSV, HTML, MD, TXT, PNG, JPG, JPEG.
+- You MUST always use this syntax, otherwise the files will not be displayed in the chat.
+
+
+# Displaying Dataframes/Tables:
+- Whenever asked to display a dataframe/table, it is CRITICAL to represent it faithfully as a markdown table in your response.
 
-You MUST always use this syntax, otherwise the files will not be displayed in the chat.
+Handling User Queries:
+- Whenever the user query requires using the python tool, you must always think first about the steps required.
+- Use the tool multiple times:
+    - You MUST NOT guess anything about the structure of the data / files uploaded. Rather, you MUST perform some data exploration first.
+        - Example: User uploads an excel files and asks a question about it. First Read the data, explore the columns, columns types, etc. Then use the tool to answer the user's query.
+        In this case, you can simply call the tool multiple times.
+        - REMEMBER that you can always read the content of text, csv files if needed. In this case, you MUST always limit the amount of data displayed.
+- If a tool step fails, you must recover as much as possible.
+- After exhausting all possible solutions without success, inform the user of what was tried and request clarification/help.
 """.strip()
 
 DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT = ""

{unique_toolkit-1.34.0 → unique_toolkit-1.35.0}/unique_toolkit/language_model/functions.py

@@ -46,7 +46,7 @@ def complete(
     timeout: int = DEFAULT_COMPLETE_TIMEOUT,
     tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
     other_options: dict | None = None,
-    structured_output_model: type[BaseModel] | None = None,
+    structured_output_model: type[BaseModel] | dict[str, Any] | None = None,
     structured_output_enforce_schema: bool = False,
     user_id: str | None = None,
 ) -> LanguageModelResponse:
@@ -110,7 +110,7 @@ async def complete_async(
     timeout: int = DEFAULT_COMPLETE_TIMEOUT,
     tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
     other_options: dict | None = None,
-    structured_output_model: type[BaseModel] | None = None,
+    structured_output_model: type[BaseModel] | dict[str, Any] | None = None,
     structured_output_enforce_schema: bool = False,
 ) -> LanguageModelResponse:
     """Call the completion endpoint asynchronously without streaming the response.
@@ -214,9 +214,21 @@ def _to_search_context(
 
 def _add_response_format_to_options(
     options: dict,
-    structured_output_model: type[BaseModel],
+    structured_output_model: type[BaseModel] | dict[str, Any],
     structured_output_enforce_schema: bool = False,
 ) -> dict:
+    if isinstance(structured_output_model, dict):
+        name = structured_output_model.get("title", "DefaultName")
+        options["responseFormat"] = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": name,
+                "strict": structured_output_enforce_schema,
+                "schema": structured_output_model,
+            },
+        }
+        return options
+
     options["responseFormat"] = {
         "type": "json_schema",
         "json_schema": {
@@ -235,7 +247,7 @@ def _prepare_completion_params_util(
     tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
     other_options: dict | None = None,
     content_chunks: list[ContentChunk] | None = None,
-    structured_output_model: type[BaseModel] | None = None,
+    structured_output_model: type[BaseModel] | dict[str, Any] | None = None,
     structured_output_enforce_schema: bool = False,
 ) -> tuple[dict, str, dict, SearchContext | None]:
     """Prepare common parameters for completion requests.
@@ -283,7 +295,7 @@ def _prepare_openai_completion_params_util(
     tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
     other_options: dict | None = None,
     content_chunks: list[ContentChunk] | None = None,
-    structured_output_model: type[BaseModel] | None = None,
+    structured_output_model: type[BaseModel] | dict[str, Any] | None = None,
     structured_output_enforce_schema: bool = False,
 ) -> tuple[dict, str, SearchContext | None]:
     """Prepare common parameters for completion requests.
@@ -355,7 +367,7 @@ def _prepare_all_completions_params_util(
     other_options: dict | None = None,
     content_chunks: list[ContentChunk] | None = None,
     tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
-    structured_output_model: type[BaseModel] | None = None,
+    structured_output_model: type[BaseModel] | dict[str, Any] | None = None,
     structured_output_enforce_schema: bool = False,
 ) -> tuple[
     dict,

{unique_toolkit-1.34.0 → unique_toolkit-1.35.0}/unique_toolkit/language_model/service.py

@@ -226,7 +226,7 @@ class LanguageModelService:
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: Optional[list[LanguageModelTool | LanguageModelToolDescription]] = None,
-        structured_output_model: Optional[Type[BaseModel]] = None,
+        structured_output_model: Optional[Type[BaseModel] | dict[str, Any]] = None,
         structured_output_enforce_schema: bool = False,
         other_options: Optional[dict] = None,
     ) -> LanguageModelResponse:
@@ -254,7 +254,7 @@ class LanguageModelService:
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: Optional[list[LanguageModelTool | LanguageModelToolDescription]] = None,
-        structured_output_model: Optional[Type[BaseModel]] = None,
+        structured_output_model: Optional[Type[BaseModel] | dict[str, Any]] = None,
         structured_output_enforce_schema: bool = False,
         other_options: Optional[dict] = None,
     ) -> LanguageModelResponse:
@@ -286,7 +286,7 @@ class LanguageModelService:
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
         tools: Optional[list[LanguageModelTool | LanguageModelToolDescription]] = None,
-        structured_output_model: Optional[Type[BaseModel]] = None,
+        structured_output_model: Optional[Type[BaseModel] | dict[str, Any]] = None,
         structured_output_enforce_schema: bool = False,
         other_options: Optional[dict] = None,
     ) -> LanguageModelResponse: