unique_toolkit 1.34.1__tar.gz → 1.35.1__tar.gz

{unique_toolkit-1.34.1 → unique_toolkit-1.35.1}/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.1] - 2025-12-05
+- Improve efficiency of token reducer if tool calls overshoot max token limit
+
+## [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.
 

{unique_toolkit-1.34.1 → unique_toolkit-1.35.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 1.34.1
+Version: 1.35.1
 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.1] - 2025-12-05
+- Improve efficiency of token reducer if tool calls overshoot max token limit
+
+## [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.
 

{unique_toolkit-1.34.1 → unique_toolkit-1.35.1}/pyproject.toml

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

{unique_toolkit-1.34.1 → unique_toolkit-1.35.1}/unique_toolkit/agentic/history_manager/loop_token_reducer.py

@@ -63,6 +63,10 @@ class LoopTokenReducer:
         self._content_service = ContentService.from_event(event)
         self._user_message = event.payload.user_message
         self._chat_id = event.payload.chat_id
+        self._effective_token_limit = int(
+            self._language_model.token_limits.token_limit_input
+            * (1 - MAX_INPUT_TOKENS_SAFETY_PERCENTAGE)
+        )
 
     def _get_encoder(self, language_model: LMI) -> tiktoken.Encoding:
         name = language_model.encoder_name or "cl100k_base"
@@ -95,7 +99,7 @@ class LoopTokenReducer:
 
         while self._exceeds_token_limit(token_count):
             token_count_before_reduction = token_count
-            loop_history = self._handle_token_limit_exceeded(loop_history)
+            loop_history = self._handle_token_limit_exceeded(loop_history, token_count)
             messages = self._construct_history(
                 history_from_db,
                 loop_history,
@@ -120,14 +124,10 @@ class LoopTokenReducer:
             len(chunks) > 1
             for chunks in self._reference_manager.get_chunks_of_all_tools()
         )
-        max_tokens = int(
-            self._language_model.token_limits.token_limit_input
-            * (1 - MAX_INPUT_TOKENS_SAFETY_PERCENTAGE)
-        )
         # TODO: This is not fully correct at the moment as the token_count
         # include system_prompt and user question already
         # TODO: There is a problem if we exceed but only have one chunk per tool call
-        exceeds_limit = token_count > max_tokens
+        exceeds_limit = token_count > self._effective_token_limit
 
         return has_multiple_chunks_for_a_tool_call and exceeds_limit
 
@@ -171,16 +171,21 @@ class LoopTokenReducer:
         return constructed_history
 
     def _handle_token_limit_exceeded(
-        self, loop_history: list[LanguageModelMessage]
+        self, loop_history: list[LanguageModelMessage], token_count: int
     ) -> list[LanguageModelMessage]:
         """Handle case where token limit is exceeded by reducing sources in tool responses."""
+        overshoot_factor = (
+            token_count / self._effective_token_limit
+            if self._effective_token_limit > 0
+            else 1.0
+        )
         self._logger.warning(
-            f"Length of messages is exceeds limit of {self._language_model.token_limits.token_limit_input} tokens. "
-            "Reducing number of sources per tool call.",
+            f"Length of messages exceeds limit of {self._effective_token_limit} tokens "
+            f"(overshoot factor: {overshoot_factor:.2f}x). Reducing number of sources per tool call.",
         )
 
         return self._reduce_message_length_by_reducing_sources_in_tool_response(
-            loop_history
+            loop_history, overshoot_factor
         )
 
     def _replace_user_message(
@@ -312,10 +317,18 @@ class LoopTokenReducer:
     def _reduce_message_length_by_reducing_sources_in_tool_response(
         self,
         history: list[LanguageModelMessage],
+        overshoot_factor: float,
     ) -> list[LanguageModelMessage]:
         """
-        Reduce the message length by removing the last source result of each tool call.
-        If there is only one source for a tool call, the tool call message is returned unchanged.
+        Reduce the message length by removing sources from each tool call based on overshoot.
+
+        The number of chunks to keep per tool call is calculated as:
+        chunks_to_keep = num_sources / (overshoot_factor * 0.75)
+
+        This ensures more aggressive reduction when we're significantly over the limit.
+        Using 0.75 factor provides a safety margin to avoid over-reduction.
+        E.g., if overshoot_factor = 2 (2x over limit), keep 1/1.5 = 2/3 of chunks.
+        Always keeps at least 1 chunk.
         """
         history_reduced: list[LanguageModelMessage] = []
         content_chunks_reduced: list[ContentChunk] = []
@@ -328,6 +341,7 @@ class LoopTokenReducer:
                     message,  # type: ignore
                     chunk_offset,
                     source_offset,
+                    overshoot_factor,
                 )
                 content_chunks_reduced.extend(result.reduced_chunks)
                 history_reduced.append(result.message)
@@ -350,10 +364,15 @@ class LoopTokenReducer:
         message: LanguageModelToolMessage,
         chunk_offset: int,
         source_offset: int,
+        overshoot_factor: float,
     ) -> SourceReductionResult:
         """
-        Reduce the sources in the tool message by removing the last source.
-        If there is only one source, the message is returned unchanged.
+        Reduce the sources in the tool message based on overshoot factor.
+
+        Chunks to keep = num_sources / (overshoot_factor * 0.75)
+        This ensures fewer chunks are kept when overshoot is larger.
+        E.g., if overshoot_factor = 2 (2x over limit), keep 1/1.5 = 2/3 of chunks
+        Always keeps at least 1 chunk.
         """
         tool_chunks = self._reference_manager.get_chunks_of_tool(message.tool_call_id)
         num_sources = len(tool_chunks)
@@ -366,16 +385,27 @@ class LoopTokenReducer:
                 source_offset=source_offset,
             )
 
-        # Reduce chunks, keeping all but the last one if multiple exist
-        if num_sources == 1:
+        # Calculate how many chunks to keep based on overshoot
+        # Use 0.75 safety margin for aggressive reduction, but only when overshoot is
+        # significant enough (>= ~1.33). Otherwise, the margin would prevent reduction.
+        divisor = (
+            overshoot_factor * 0.75
+            if overshoot_factor * 0.75 >= 1.0
+            else overshoot_factor
+        )
+        chunks_to_keep = max(1, int(num_sources / divisor))
+
+        # Reduce chunks
+        if chunks_to_keep >= num_sources:
+            # No reduction needed for this tool call
             reduced_chunks = tool_chunks
             content_chunks_reduced = self._reference_manager.get_chunks()[
                 chunk_offset : chunk_offset + num_sources
             ]
         else:
-            reduced_chunks = tool_chunks[:-1]
+            reduced_chunks = tool_chunks[:chunks_to_keep]
             content_chunks_reduced = self._reference_manager.get_chunks()[
-                chunk_offset : chunk_offset + num_sources - 1
+                chunk_offset : chunk_offset + chunks_to_keep
             ]
             self._reference_manager.replace_chunks_of_tool(
                 message.tool_call_id, reduced_chunks
@@ -392,7 +422,7 @@ class LoopTokenReducer:
             message=new_message,
             reduced_chunks=content_chunks_reduced,
             chunk_offset=chunk_offset + num_sources,
-            source_offset=source_offset + num_sources - (1 if num_sources != 1 else 0),
+            source_offset=source_offset + len(reduced_chunks),
         )
 
     def _create_tool_call_message_with_reduced_sources(

unique_toolkit-1.35.1/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.1/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.1/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.1/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.1/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.1/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.1/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.1/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.1/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.1 → unique_toolkit-1.35.1}/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,