grasp_agents 0.5.9__py3-none-any.whl → 0.5.11__py3-none-any.whl

grasp_agents/llm_policy_executor.py

@@ -1,6 +1,6 @@
 import asyncio
 import json
-from collections.abc import AsyncIterator, Coroutine, Sequence
+from collections.abc import AsyncIterator, Coroutine, Mapping, Sequence
 from itertools import starmap
 from logging import getLogger
 from typing import Any, Generic, Protocol, final
@@ -36,7 +36,8 @@ class ToolCallLoopTerminator(Protocol[CtxT]):
         self,
         conversation: Messages,
         *,
-        ctx: RunContext[CtxT] | None,
+        ctx: RunContext[CtxT],
+        call_id: str,
         **kwargs: Any,
     ) -> bool: ...
 
@@ -46,7 +47,8 @@ class MemoryManager(Protocol[CtxT]):
         self,
         memory: LLMAgentMemory,
         *,
-        ctx: RunContext[CtxT] | None,
+        ctx: RunContext[CtxT],
+        call_id: str,
         **kwargs: Any,
     ) -> None: ...
 
@@ -54,9 +56,12 @@ class MemoryManager(Protocol[CtxT]):
 class LLMPolicyExecutor(Generic[CtxT]):
     def __init__(
         self,
+        *,
         agent_name: str,
         llm: LLM[LLMSettings, Converters],
         tools: list[BaseTool[BaseModel, Any, CtxT]] | None,
+        response_schema: Any | None = None,
+        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
         max_turns: int,
         react_mode: bool = False,
         final_answer_type: type[BaseModel] = BaseModel,
@@ -70,12 +75,15 @@ class LLMPolicyExecutor(Generic[CtxT]):
         self._final_answer_as_tool_call = final_answer_as_tool_call
         self._final_answer_tool = self.get_final_answer_tool()
 
-        _tools: list[BaseTool[BaseModel, Any, CtxT]] | None = tools
+        tools_list: list[BaseTool[BaseModel, Any, CtxT]] | None = tools
         if tools and final_answer_as_tool_call:
-            _tools = tools + [self._final_answer_tool]
+            tools_list = tools + [self._final_answer_tool]
+        self._tools = {t.name: t for t in tools_list} if tools_list else None
+
+        self._response_schema = response_schema
+        self._response_schema_by_xml_tag = response_schema_by_xml_tag
 
         self._llm = llm
-        self._llm.tools = _tools
 
         self._max_turns = max_turns
         self._react_mode = react_mode
@@ -91,9 +99,21 @@ class LLMPolicyExecutor(Generic[CtxT]):
     def llm(self) -> LLM[LLMSettings, Converters]:
         return self._llm
 
+    @property
+    def response_schema(self) -> Any | None:
+        return self._response_schema
+
+    @response_schema.setter
+    def response_schema(self, value: Any | None) -> None:
+        self._response_schema = value
+
+    @property
+    def response_schema_by_xml_tag(self) -> Mapping[str, Any] | None:
+        return self._response_schema_by_xml_tag
+
     @property
     def tools(self) -> dict[str, BaseTool[BaseModel, Any, CtxT]]:
-        return self._llm.tools or {}
+        return self._tools or {}
 
     @property
     def max_turns(self) -> int:
@@ -104,11 +124,14 @@ class LLMPolicyExecutor(Generic[CtxT]):
         self,
         conversation: Messages,
         *,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
+        call_id: str,
         **kwargs: Any,
     ) -> bool:
         if self.tool_call_loop_terminator:
-            return self.tool_call_loop_terminator(conversation, ctx=ctx, **kwargs)
+            return self.tool_call_loop_terminator(
+                conversation, ctx=ctx, call_id=call_id, **kwargs
+            )
 
         return False
 
@@ -117,21 +140,26 @@ class LLMPolicyExecutor(Generic[CtxT]):
         self,
         memory: LLMAgentMemory,
         *,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
+        call_id: str,
         **kwargs: Any,
     ) -> None:
         if self.memory_manager:
-            self.memory_manager(memory=memory, ctx=ctx, **kwargs)
+            self.memory_manager(memory=memory, ctx=ctx, call_id=call_id, **kwargs)
 
     async def generate_message(
         self,
         memory: LLMAgentMemory,
-        call_id: str,
+        *,
         tool_choice: ToolChoice | None = None,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
+        call_id: str,
     ) -> AssistantMessage:
         completion = await self.llm.generate_completion(
             memory.message_history,
+            response_schema=self.response_schema,
+            response_schema_by_xml_tag=self.response_schema_by_xml_tag,
+            tools=self.tools,
             tool_choice=tool_choice,
             n_choices=1,
             proc_name=self.agent_name,
@@ -139,7 +167,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         )
         memory.update(completion.messages)
         self._process_completion(
-            completion, call_id=call_id, ctx=ctx, print_messages=True
+            completion, ctx=ctx, call_id=call_id, print_messages=True
         )
 
         return completion.messages[0]
@@ -147,9 +175,10 @@ class LLMPolicyExecutor(Generic[CtxT]):
     async def generate_message_stream(
         self,
         memory: LLMAgentMemory,
-        call_id: str,
+        *,
         tool_choice: ToolChoice | None = None,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
+        call_id: str,
     ) -> AsyncIterator[
         CompletionChunkEvent[CompletionChunk]
         | CompletionEvent
@@ -160,6 +189,9 @@ class LLMPolicyExecutor(Generic[CtxT]):
 
         llm_event_stream = self.llm.generate_completion_stream(
             memory.message_history,
+            response_schema=self.response_schema,
+            response_schema_by_xml_tag=self.response_schema_by_xml_tag,
+            tools=self.tools,
             tool_choice=tool_choice,
             n_choices=1,
             proc_name=self.agent_name,
@@ -181,22 +213,22 @@ class LLMPolicyExecutor(Generic[CtxT]):
         memory.update(completion.messages)
 
         self._process_completion(
-            completion, call_id=call_id, print_messages=True, ctx=ctx
+            completion, print_messages=True, ctx=ctx, call_id=call_id
         )
 
     async def call_tools(
         self,
         calls: Sequence[ToolCall],
         memory: LLMAgentMemory,
+        ctx: RunContext[CtxT],
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
     ) -> Sequence[ToolMessage]:
         # TODO: Add image support
         corouts: list[Coroutine[Any, Any, BaseModel]] = []
         for call in calls:
             tool = self.tools[call.tool_name]
             args = json.loads(call.tool_arguments)
-            corouts.append(tool(ctx=ctx, **args))
+            corouts.append(tool(ctx=ctx, call_id=call_id, **args))
 
         outs = await asyncio.gather(*corouts)
         tool_messages = list(
@@ -216,11 +248,11 @@ class LLMPolicyExecutor(Generic[CtxT]):
         self,
         calls: Sequence[ToolCall],
         memory: LLMAgentMemory,
+        ctx: RunContext[CtxT],
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[ToolMessageEvent]:
         tool_messages = await self.call_tools(
-            calls, memory=memory, call_id=call_id, ctx=ctx
+            calls, memory=memory, ctx=ctx, call_id=call_id
         )
         for tool_message, call in zip(tool_messages, calls, strict=True):
             yield ToolMessageEvent(
@@ -245,7 +277,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
                 return final_answer_message
 
     async def _generate_final_answer(
-        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT] | None = None
+        self, memory: LLMAgentMemory, ctx: RunContext[CtxT], call_id: str
     ) -> AssistantMessage:
         user_message = UserMessage.from_text(
             "Exceeded the maximum number of turns: provide a final answer now!"
@@ -258,7 +290,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
 
         tool_choice = NamedToolChoice(name=self._final_answer_tool.name)
         await self.generate_message(
-            memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
+            memory, tool_choice=tool_choice, ctx=ctx, call_id=call_id
         )
 
         final_answer_message = self._extract_final_answer_from_tool_calls(memory=memory)
@@ -268,7 +300,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         return final_answer_message
 
     async def _generate_final_answer_stream(
-        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT] | None = None
+        self, memory: LLMAgentMemory, ctx: RunContext[CtxT], call_id: str
     ) -> AsyncIterator[Event[Any]]:
         user_message = UserMessage.from_text(
             "Exceeded the maximum number of turns: provide a final answer now!",
@@ -284,7 +316,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
 
         tool_choice = NamedToolChoice(name=self._final_answer_tool.name)
         async for event in self.generate_message_stream(
-            memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
+            memory, tool_choice=tool_choice, ctx=ctx, call_id=call_id
         ):
             yield event
 
@@ -296,7 +328,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         )
 
     async def execute(
-        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT] | None = None
+        self, memory: LLMAgentMemory, ctx: RunContext[CtxT], call_id: str
     ) -> AssistantMessage | Sequence[AssistantMessage]:
         # 1. Generate the first message:
         #    In ReAct mode, we generate the first message without tool calls
@@ -306,7 +338,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         if self.tools:
             tool_choice = "none" if self._react_mode else "auto"
         gen_message = await self.generate_message(
-            memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
+            memory, tool_choice=tool_choice, ctx=ctx, call_id=call_id
         )
         if not self.tools:
             return gen_message
@@ -319,7 +351,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
             # If a final answer is not provided via a tool call, we use
             # _terminate_tool_call_loop to determine whether to exit the loop.
             if not self._final_answer_as_tool_call and self._terminate_tool_call_loop(
-                memory.message_history, ctx=ctx, num_turns=turns
+                memory.message_history, ctx=ctx, call_id=call_id, num_turns=turns
             ):
                 return gen_message
 
@@ -338,7 +370,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
                 # Otherwise, we simply return the last generated message.
                 if self._final_answer_as_tool_call:
                     final_answer = await self._generate_final_answer(
-                        memory, call_id=call_id, ctx=ctx
+                        memory, ctx=ctx, call_id=call_id
                     )
                 else:
                     final_answer = gen_message
@@ -351,11 +383,11 @@ class LLMPolicyExecutor(Generic[CtxT]):
 
             if gen_message.tool_calls:
                 await self.call_tools(
-                    gen_message.tool_calls, memory=memory, call_id=call_id, ctx=ctx
+                    gen_message.tool_calls, memory=memory, ctx=ctx, call_id=call_id
                 )
 
             # Apply memory management (e.g. compacting or pruning memory)
-            self._manage_memory(memory, ctx=ctx, num_turns=turns)
+            self._manage_memory(memory, ctx=ctx, call_id=call_id, num_turns=turns)
 
             # 4. Generate the next message based on the updated memory.
             #    In ReAct mode, we set tool_choice to "none" if we just called tools,
@@ -370,7 +402,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
                 tool_choice = "required"
 
             gen_message = await self.generate_message(
-                memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
+                memory, tool_choice=tool_choice, ctx=ctx, call_id=call_id
             )
 
             turns += 1
@@ -378,13 +410,13 @@ class LLMPolicyExecutor(Generic[CtxT]):
     async def execute_stream(
         self,
         memory: LLMAgentMemory,
+        ctx: RunContext[CtxT],
         call_id: str,
-        ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         tool_choice: ToolChoice = "none" if self._react_mode else "auto"
         gen_message: AssistantMessage | None = None
         async for event in self.generate_message_stream(
-            memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
+            memory, tool_choice=tool_choice, ctx=ctx, call_id=call_id
         ):
             if isinstance(event, GenMessageEvent):
                 gen_message = event.data
@@ -399,7 +431,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
 
         while True:
             if not self._final_answer_as_tool_call and self._terminate_tool_call_loop(
-                memory.message_history, ctx=ctx, num_turns=turns
+                memory.message_history, ctx=ctx, call_id=call_id, num_turns=turns
             ):
                 return
 
@@ -418,7 +450,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
             if turns >= self.max_turns:
                 if self._final_answer_as_tool_call:
                     async for event in self._generate_final_answer_stream(
-                        memory, call_id=call_id, ctx=ctx
+                        memory, ctx=ctx, call_id=call_id
                     ):
                         yield event
                 logger.info(
@@ -433,11 +465,11 @@ class LLMPolicyExecutor(Generic[CtxT]):
                     )
 
                 async for event in self.call_tools_stream(
-                    gen_message.tool_calls, memory=memory, call_id=call_id, ctx=ctx
+                    gen_message.tool_calls, memory=memory, ctx=ctx, call_id=call_id
                 ):
                     yield event
 
-            self._manage_memory(memory, ctx=ctx, num_turns=turns)
+            self._manage_memory(memory, ctx=ctx, call_id=call_id, num_turns=turns)
 
             if self._react_mode and gen_message.tool_calls:
                 tool_choice = "none"
@@ -447,7 +479,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
                 tool_choice = "required"
 
             async for event in self.generate_message_stream(
-                memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
+                memory, tool_choice=tool_choice, ctx=ctx, call_id=call_id
             ):
                 yield event
                 if isinstance(event, GenMessageEvent):
@@ -464,7 +496,11 @@ class LLMPolicyExecutor(Generic[CtxT]):
             )
 
             async def run(
-                self, inp: BaseModel, ctx: RunContext[Any] | None = None
+                self,
+                inp: BaseModel,
+                *,
+                ctx: RunContext[Any] | None = None,
+                call_id: str | None = None,
             ) -> None:
                 return None
 
@@ -473,22 +509,22 @@ class LLMPolicyExecutor(Generic[CtxT]):
     def _process_completion(
         self,
         completion: Completion,
-        call_id: str,
+        *,
         print_messages: bool = False,
-        ctx: RunContext[CtxT] | None = None,
+        ctx: RunContext[CtxT],
+        call_id: str,
     ) -> None:
-        if ctx is not None:
-            ctx.completions[self.agent_name].append(completion)
-            ctx.usage_tracker.update(
+        ctx.completions[self.agent_name].append(completion)
+        ctx.usage_tracker.update(
+            agent_name=self.agent_name,
+            completions=[completion],
+            model_name=self.llm.model_name,
+        )
+        if ctx.printer and print_messages:
+            usages = [None] * (len(completion.messages) - 1) + [completion.usage]
+            ctx.printer.print_messages(
+                completion.messages,
+                usages=usages,
                 agent_name=self.agent_name,
-                completions=[completion],
-                model_name=self.llm.model_name,
+                call_id=call_id,
             )
-            if ctx.printer and print_messages:
-                usages = [None] * (len(completion.messages) - 1) + [completion.usage]
-                ctx.printer.print_messages(
-                    completion.messages,
-                    usages=usages,
-                    agent_name=self.agent_name,
-                    call_id=call_id,
-                )

grasp_agents/openai/converters.py

@@ -96,8 +96,10 @@ class OpenAIConverters(Converters):
         return from_api_tool_message(raw_message, name=name, **kwargs)
 
     @staticmethod
-    def to_tool(tool: BaseTool[BaseModel, Any, Any], **kwargs: Any) -> OpenAIToolParam:
-        return to_api_tool(tool, **kwargs)
+    def to_tool(
+        tool: BaseTool[BaseModel, Any, Any], strict: bool | None = None, **kwargs: Any
+    ) -> OpenAIToolParam:
+        return to_api_tool(tool, strict=strict, **kwargs)
 
     @staticmethod
     def to_tool_choice(

grasp_agents/openai/openai_llm.py

@@ -3,9 +3,9 @@ import logging
 import os
 from collections.abc import AsyncIterator, Iterable, Mapping
 from copy import deepcopy
+from dataclasses import dataclass, field
 from typing import Any, Literal
 
-import httpx
 from openai import AsyncOpenAI, AsyncStream
 from openai._types import NOT_GIVEN  # type: ignore[import]
 from openai.lib.streaming.chat import (
@@ -15,8 +15,7 @@ from openai.lib.streaming.chat import ChatCompletionStreamState
 from openai.lib.streaming.chat import ChunkEvent as OpenAIChunkEvent
 from pydantic import BaseModel
 
-from ..cloud_llm import APIProvider, CloudLLM, CloudLLMSettings, LLMRateLimiter
-from ..http_client import AsyncHTTPClientParams
+from ..cloud_llm import APIProvider, CloudLLM, CloudLLMSettings
 from ..typing.tool import BaseTool
 from . import (
     OpenAICompletion,
@@ -61,7 +60,7 @@ def get_openai_compatible_providers() -> list[APIProvider]:
 
 
 class OpenAILLMSettings(CloudLLMSettings, total=False):
-    reasoning_effort: Literal["low", "medium", "high"] | None
+    reasoning_effort: Literal["disable", "minimal", "low", "medium", "high"] | None
 
     parallel_tool_calls: bool
 
@@ -90,105 +89,75 @@ class OpenAILLMSettings(CloudLLMSettings, total=False):
     # TODO: support audio
 
 
+@dataclass(frozen=True)
 class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
-    def __init__(
-        self,
-        # Base LLM args
-        model_name: str,
-        llm_settings: OpenAILLMSettings | None = None,
-        tools: list[BaseTool[BaseModel, Any, Any]] | None = None,
-        response_schema: Any | None = None,
-        response_schema_by_xml_tag: Mapping[str, Any] | None = None,
-        apply_response_schema_via_provider: bool = False,
-        model_id: str | None = None,
-        # Custom LLM provider
-        api_provider: APIProvider | None = None,
-        # Connection settings
-        max_client_retries: int = 2,
-        async_http_client: httpx.AsyncClient | None = None,
-        async_http_client_params: (
-            dict[str, Any] | AsyncHTTPClientParams | None
-        ) = None,
-        async_openai_client_params: dict[str, Any] | None = None,
-        # Rate limiting
-        rate_limiter: LLMRateLimiter | None = None,
-        # LLM response retries: try to regenerate to pass validation
-        max_response_retries: int = 1,
-    ) -> None:
+    converters: OpenAIConverters = field(default_factory=OpenAIConverters)
+    async_openai_client_params: dict[str, Any] | None = None
+    client: AsyncOpenAI = field(init=False)
+
+    def __post_init__(self):
+        super().__post_init__()
+
         openai_compatible_providers = get_openai_compatible_providers()
 
-        model_name_parts = model_name.split("/", 1)
-        if api_provider is not None:
-            provider_model_name = model_name
+        _api_provider = self.api_provider
+
+        model_name_parts = self.model_name.split("/", 1)
+        if _api_provider is not None:
+            _model_name = self.model_name
         elif len(model_name_parts) == 2:
             compat_providers_map = {
                 provider["name"]: provider for provider in openai_compatible_providers
             }
-            provider_name, provider_model_name = model_name_parts
+            provider_name, _model_name = model_name_parts
             if provider_name not in compat_providers_map:
                 raise ValueError(
                     f"API provider '{provider_name}' is not a supported OpenAI "
                     f"compatible provider. Supported providers are: "
                     f"{', '.join(compat_providers_map.keys())}"
                 )
-            api_provider = compat_providers_map[provider_name]
+            _api_provider = compat_providers_map[provider_name]
         else:
             raise ValueError(
                 "Model name must be in the format 'provider/model_name' or "
                 "you must provide an 'api_provider' argument."
             )
 
-        if llm_settings is not None:
-            stream_options = llm_settings.get("stream_options") or {}
+        if self.llm_settings is not None:
+            stream_options = self.llm_settings.get("stream_options") or {}
             stream_options["include_usage"] = True
-            _llm_settings = deepcopy(llm_settings)
+            _llm_settings = deepcopy(self.llm_settings)
             _llm_settings["stream_options"] = stream_options
         else:
             _llm_settings = OpenAILLMSettings(stream_options={"include_usage": True})
 
-        super().__init__(
-            model_name=provider_model_name,
-            model_id=model_id,
-            llm_settings=_llm_settings,
-            converters=OpenAIConverters(),
-            tools=tools,
-            response_schema=response_schema,
-            response_schema_by_xml_tag=response_schema_by_xml_tag,
-            apply_response_schema_via_provider=apply_response_schema_via_provider,
-            api_provider=api_provider,
-            async_http_client=async_http_client,
-            async_http_client_params=async_http_client_params,
-            rate_limiter=rate_limiter,
-            max_client_retries=max_client_retries,
-            max_response_retries=max_response_retries,
-        )
-
         response_schema_support: bool = any(
-            fnmatch.fnmatch(self._model_name, pat)
-            for pat in api_provider.get("response_schema_support") or []
+            fnmatch.fnmatch(_model_name, pat)
+            for pat in _api_provider.get("response_schema_support") or []
         )
-        if apply_response_schema_via_provider:
-            if self._tools:
-                for tool in self._tools.values():
-                    tool.strict = True
-            if not response_schema_support:
-                raise ValueError(
-                    "Native response schema validation is not supported for model "
-                    f"'{self._model_name}' by the API provider. Please set "
-                    "apply_response_schema_via_provider=False."
-                )
+        if self.apply_response_schema_via_provider and not response_schema_support:
+            raise ValueError(
+                "Native response schema validation is not supported for model "
+                f"'{_model_name}' by the API provider. Please set "
+                "apply_response_schema_via_provider=False."
+            )
 
-        _async_openai_client_params = deepcopy(async_openai_client_params or {})
-        if self._async_http_client is not None:
-            _async_openai_client_params["http_client"] = self._async_http_client
+        _async_openai_client_params = deepcopy(self.async_openai_client_params or {})
+        if self.async_http_client is not None:
+            _async_openai_client_params["http_client"] = self.async_http_client
 
-        self._client: AsyncOpenAI = AsyncOpenAI(
-            base_url=self.api_provider.get("base_url"),
-            api_key=self.api_provider.get("api_key"),
-            max_retries=max_client_retries,
+        _client = AsyncOpenAI(
+            base_url=_api_provider.get("base_url"),
+            api_key=_api_provider.get("api_key"),
+            max_retries=self.max_client_retries,
             **_async_openai_client_params,
         )
 
+        object.__setattr__(self, "model_name", _model_name)
+        object.__setattr__(self, "api_provider", _api_provider)
+        object.__setattr__(self, "llm_settings", _llm_settings)
+        object.__setattr__(self, "client", _client)
+
     async def _get_completion(
         self,
         api_messages: Iterable[OpenAIMessageParam],
@@ -203,9 +172,9 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         response_format = api_response_schema or NOT_GIVEN
         n = n_choices or NOT_GIVEN
 
-        if self._apply_response_schema_via_provider:
-            return await self._client.beta.chat.completions.parse(
-                model=self._model_name,
+        if self.apply_response_schema_via_provider:
+            return await self.client.beta.chat.completions.parse(
+                model=self.model_name,
                 messages=api_messages,
                 tools=tools,
                 tool_choice=tool_choice,
@@ -214,8 +183,8 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
                 **api_llm_settings,
             )
 
-        return await self._client.chat.completions.create(
-            model=self._model_name,
+        return await self.client.chat.completions.create(
+            model=self.model_name,
             messages=api_messages,
             tools=tools,
             tool_choice=tool_choice,
@@ -238,10 +207,10 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         response_format = api_response_schema or NOT_GIVEN
         n = n_choices or NOT_GIVEN
 
-        if self._apply_response_schema_via_provider:
+        if self.apply_response_schema_via_provider:
             stream_manager: OpenAIAsyncChatCompletionStreamManager[Any] = (
-                self._client.beta.chat.completions.stream(
-                    model=self._model_name,
+                self.client.beta.chat.completions.stream(
+                    model=self.model_name,
                     messages=api_messages,
                     tools=tools,
                     tool_choice=tool_choice,
@@ -257,8 +226,8 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         else:
             stream_generator: AsyncStream[
                 OpenAICompletionChunk
-            ] = await self._client.chat.completions.create(
-                model=self._model_name,
+            ] = await self.client.chat.completions.create(
+                model=self.model_name,
                 messages=api_messages,
                 tools=tools,
                 tool_choice=tool_choice,
@@ -271,16 +240,20 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
                     yield completion_chunk
 
     def combine_completion_chunks(
-        self, completion_chunks: list[OpenAICompletionChunk]
+        self,
+        completion_chunks: list[OpenAICompletionChunk],
+        response_schema: Any | None = None,
+        tools: Mapping[str, BaseTool[BaseModel, Any, Any]] | None = None,
     ) -> OpenAICompletion:
         response_format = NOT_GIVEN
         input_tools = NOT_GIVEN
-        if self._apply_response_schema_via_provider:
-            if self._response_schema:
-                response_format = self._response_schema
-            if self._tools:
+        if self.apply_response_schema_via_provider:
+            if response_schema:
+                response_format = response_schema
+            if tools:
                 input_tools = [
-                    self._converters.to_tool(tool) for tool in self._tools.values()
+                    self.converters.to_tool(tool, strict=True)
+                    for tool in tools.values()
                 ]
         state = ChatCompletionStreamState[Any](
             input_tools=input_tools, response_format=response_format