llama-stack 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

llama_stack/providers/inline/agents/meta_reference/responses/tool_executor.py

@@ -7,8 +7,18 @@
 import asyncio
 import json
 from collections.abc import AsyncIterator
+from typing import Any
 
-from llama_stack.apis.agents.openai_responses import (
+from opentelemetry import trace
+
+from llama_stack.core.datatypes import VectorStoresConfig
+from llama_stack.log import get_logger
+from llama_stack_api import (
+    ImageContentItem,
+    OpenAIChatCompletionContentPartImageParam,
+    OpenAIChatCompletionContentPartTextParam,
+    OpenAIChatCompletionToolCall,
+    OpenAIImageURL,
     OpenAIResponseInputToolFileSearch,
     OpenAIResponseInputToolMCP,
     OpenAIResponseObjectStreamResponseFileSearchCallCompleted,
@@ -23,26 +33,18 @@ from llama_stack.apis.agents.openai_responses import (
     OpenAIResponseOutputMessageFileSearchToolCall,
     OpenAIResponseOutputMessageFileSearchToolCallResults,
     OpenAIResponseOutputMessageWebSearchToolCall,
-)
-from llama_stack.apis.common.content_types import (
-    ImageContentItem,
-    TextContentItem,
-)
-from llama_stack.apis.inference import (
-    OpenAIChatCompletionContentPartImageParam,
-    OpenAIChatCompletionContentPartTextParam,
-    OpenAIChatCompletionToolCall,
-    OpenAIImageURL,
     OpenAIToolMessageParam,
+    TextContentItem,
+    ToolGroups,
+    ToolInvocationResult,
+    ToolRuntime,
+    VectorIO,
 )
-from llama_stack.apis.tools import ToolGroups, ToolInvocationResult, ToolRuntime
-from llama_stack.apis.vector_io import VectorIO
-from llama_stack.log import get_logger
-from llama_stack.providers.utils.telemetry import tracing
 
 from .types import ChatCompletionContext, ToolExecutionResult
 
 logger = get_logger(name=__name__, category="agents::meta_reference")
+tracer = trace.get_tracer(__name__)
 
 
 class ToolExecutor:
@@ -51,10 +53,12 @@ class ToolExecutor:
         tool_groups_api: ToolGroups,
         tool_runtime_api: ToolRuntime,
         vector_io_api: VectorIO,
+        vector_stores_config=None,
     ):
         self.tool_groups_api = tool_groups_api
         self.tool_runtime_api = tool_runtime_api
         self.vector_io_api = vector_io_api
+        self.vector_stores_config = vector_stores_config
 
     async def execute_tool_call(
         self,
@@ -67,7 +71,7 @@ class ToolExecutor:
     ) -> AsyncIterator[ToolExecutionResult]:
         tool_call_id = tool_call.id
         function = tool_call.function
-        tool_kwargs = json.loads(function.arguments) if function.arguments else {}
+        tool_kwargs = json.loads(function.arguments) if function and function.arguments else {}
 
         if not function or not tool_call_id or not function.name:
             yield ToolExecutionResult(sequence_number=sequence_number)
@@ -84,7 +88,16 @@ class ToolExecutor:
         error_exc, result = await self._execute_tool(function.name, tool_kwargs, ctx, mcp_tool_to_server)
 
         # Emit completion events for tool execution
-        has_error = error_exc or (result and ((result.error_code and result.error_code > 0) or result.error_message))
+        has_error = bool(
+            error_exc
+            or (
+                result
+                and (
+                    ((error_code := getattr(result, "error_code", None)) and error_code > 0)
+                    or getattr(result, "error_message", None)
+                )
+            )
+        )
         async for event_result in self._emit_completion_events(
             function.name, ctx, sequence_number, output_index, item_id, has_error, mcp_tool_to_server
         ):
@@ -101,7 +114,9 @@ class ToolExecutor:
             sequence_number=sequence_number,
             final_output_message=output_message,
             final_input_message=input_message,
-            citation_files=result.metadata.get("citation_files") if result and result.metadata else None,
+            citation_files=(
+                metadata.get("citation_files") if result and (metadata := getattr(result, "metadata", None)) else None
+            ),
         )
 
     async def _execute_knowledge_search_via_vector_store(
@@ -136,13 +151,35 @@ class ToolExecutor:
         for results in all_results:
             search_results.extend(results)
 
-        content_items = []
-        content_items.append(
-            TextContentItem(
-                text=f"knowledge_search tool found {len(search_results)} chunks:\nBEGIN of knowledge_search tool results.\n"
-            )
+        # Get templates from vector stores config, fallback to constants
+
+        # Check if annotations are enabled
+        enable_annotations = (
+            self.vector_stores_config
+            and self.vector_stores_config.annotation_prompt_params
+            and self.vector_stores_config.annotation_prompt_params.enable_annotations
         )
 
+        # Get templates
+        header_template = self.vector_stores_config.file_search_params.header_template
+        footer_template = self.vector_stores_config.file_search_params.footer_template
+        context_template = self.vector_stores_config.context_prompt_params.context_template
+
+        # Get annotation templates (use defaults if annotations disabled)
+        if enable_annotations:
+            chunk_annotation_template = self.vector_stores_config.annotation_prompt_params.chunk_annotation_template
+            annotation_instruction_template = (
+                self.vector_stores_config.annotation_prompt_params.annotation_instruction_template
+            )
+        else:
+            # Use defaults from VectorStoresConfig when annotations disabled
+            default_config = VectorStoresConfig()
+            chunk_annotation_template = default_config.annotation_prompt_params.chunk_annotation_template
+            annotation_instruction_template = default_config.annotation_prompt_params.annotation_instruction_template
+
+        content_items = []
+        content_items.append(TextContentItem(text=header_template.format(num_chunks=len(search_results))))
+
         unique_files = set()
         for i, result_item in enumerate(search_results):
             chunk_text = result_item.content[0].text if result_item.content else ""
@@ -154,22 +191,23 @@ class ToolExecutor:
             if result_item.attributes:
                 metadata_text += f", attributes: {result_item.attributes}"
 
-            text_content = f"[{i + 1}] {metadata_text} (cite as <|{file_id}|>)\n{chunk_text}\n"
+            text_content = chunk_annotation_template.format(
+                index=i + 1, metadata_text=metadata_text, file_id=file_id, chunk_text=chunk_text
+            )
             content_items.append(TextContentItem(text=text_content))
             unique_files.add(file_id)
 
-        content_items.append(TextContentItem(text="END of knowledge_search tool results.\n"))
+        content_items.append(TextContentItem(text=footer_template))
 
-        citation_instruction = ""
+        annotation_instruction = ""
         if unique_files:
-            citation_instruction = (
-                " Cite sources immediately at the end of sentences before punctuation, using `<|file-id|>` format (e.g., 'This is a fact <|file-Cn3MSNn72ENTiiq11Qda4A|>.'). "
-                "Do not add extra punctuation. Use only the file IDs provided (do not invent new ones)."
-            )
+            annotation_instruction = annotation_instruction_template
 
         content_items.append(
             TextContentItem(
-                text=f'The above results were retrieved to help answer the user\'s query: "{query}". Use them as supporting information only in answering this query.{citation_instruction}\n',
+                text=context_template.format(
+                    query=query, num_chunks=len(search_results), annotation_instruction=annotation_instruction
+                )
             )
         )
 
@@ -188,8 +226,9 @@ class ToolExecutor:
 
             citation_files[file_id] = filename
 
+        # Cast to proper InterleavedContent type (list invariance)
         return ToolInvocationResult(
-            content=content_items,
+            content=content_items,  # type: ignore[arg-type]
             metadata={
                 "document_ids": [r.file_id for r in search_results],
                 "chunks": [r.content[0].text if r.content else "" for r in search_results],
@@ -209,51 +248,60 @@ class ToolExecutor:
     ) -> AsyncIterator[ToolExecutionResult]:
         """Emit progress events for tool execution start."""
         # Emit in_progress event based on tool type (only for tools with specific streaming events)
-        progress_event = None
         if mcp_tool_to_server and function_name in mcp_tool_to_server:
             sequence_number += 1
-            progress_event = OpenAIResponseObjectStreamResponseMcpCallInProgress(
-                item_id=item_id,
-                output_index=output_index,
+            yield ToolExecutionResult(
+                stream_event=OpenAIResponseObjectStreamResponseMcpCallInProgress(
+                    item_id=item_id,
+                    output_index=output_index,
+                    sequence_number=sequence_number,
+                ),
                 sequence_number=sequence_number,
             )
         elif function_name == "web_search":
             sequence_number += 1
-            progress_event = OpenAIResponseObjectStreamResponseWebSearchCallInProgress(
-                item_id=item_id,
-                output_index=output_index,
+            yield ToolExecutionResult(
+                stream_event=OpenAIResponseObjectStreamResponseWebSearchCallInProgress(
+                    item_id=item_id,
+                    output_index=output_index,
+                    sequence_number=sequence_number,
+                ),
                 sequence_number=sequence_number,
             )
         elif function_name == "knowledge_search":
             sequence_number += 1
-            progress_event = OpenAIResponseObjectStreamResponseFileSearchCallInProgress(
-                item_id=item_id,
-                output_index=output_index,
+            yield ToolExecutionResult(
+                stream_event=OpenAIResponseObjectStreamResponseFileSearchCallInProgress(
+                    item_id=item_id,
+                    output_index=output_index,
+                    sequence_number=sequence_number,
+                ),
                 sequence_number=sequence_number,
             )
 
-        if progress_event:
-            yield ToolExecutionResult(stream_event=progress_event, sequence_number=sequence_number)
-
         # For web search, emit searching event
         if function_name == "web_search":
             sequence_number += 1
-            searching_event = OpenAIResponseObjectStreamResponseWebSearchCallSearching(
-                item_id=item_id,
-                output_index=output_index,
+            yield ToolExecutionResult(
+                stream_event=OpenAIResponseObjectStreamResponseWebSearchCallSearching(
+                    item_id=item_id,
+                    output_index=output_index,
+                    sequence_number=sequence_number,
+                ),
                 sequence_number=sequence_number,
             )
-            yield ToolExecutionResult(stream_event=searching_event, sequence_number=sequence_number)
 
         # For file search, emit searching event
         if function_name == "knowledge_search":
             sequence_number += 1
-            searching_event = OpenAIResponseObjectStreamResponseFileSearchCallSearching(
-                item_id=item_id,
-                output_index=output_index,
+            yield ToolExecutionResult(
+                stream_event=OpenAIResponseObjectStreamResponseFileSearchCallSearching(
+                    item_id=item_id,
+                    output_index=output_index,
+                    sequence_number=sequence_number,
+                ),
                 sequence_number=sequence_number,
             )
-            yield ToolExecutionResult(stream_event=searching_event, sequence_number=sequence_number)
 
     async def _execute_tool(
         self,
@@ -261,7 +309,7 @@ class ToolExecutor:
         tool_kwargs: dict,
         ctx: ChatCompletionContext,
         mcp_tool_to_server: dict[str, OpenAIResponseInputToolMCP] | None = None,
-    ) -> tuple[Exception | None, any]:
+    ) -> tuple[Exception | None, Any]:
         """Execute the tool and return error exception and result."""
         error_exc = None
         result = None
@@ -276,23 +324,30 @@ class ToolExecutor:
                     "server_url": mcp_tool.server_url,
                     "tool_name": function_name,
                 }
-                async with tracing.span("invoke_mcp_tool", attributes):
+                # TODO: follow semantic conventions for Open Telemetry tool spans
+                # https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span
+                with tracer.start_as_current_span("invoke_mcp_tool", attributes=attributes):
                     result = await invoke_mcp_tool(
                         endpoint=mcp_tool.server_url,
-                        headers=mcp_tool.headers or {},
                         tool_name=function_name,
                         kwargs=tool_kwargs,
+                        headers=mcp_tool.headers,
+                        authorization=mcp_tool.authorization,
                     )
             elif function_name == "knowledge_search":
-                response_file_search_tool = next(
-                    (t for t in ctx.response_tools if isinstance(t, OpenAIResponseInputToolFileSearch)),
-                    None,
+                response_file_search_tool = (
+                    next(
+                        (t for t in ctx.response_tools if isinstance(t, OpenAIResponseInputToolFileSearch)),
+                        None,
+                    )
+                    if ctx.response_tools
+                    else None
                 )
                 if response_file_search_tool:
                     # Use vector_stores.search API instead of knowledge_search tool
                     # to support filters and ranking_options
                     query = tool_kwargs.get("query", "")
-                    async with tracing.span("knowledge_search", {}):
+                    with tracer.start_as_current_span("knowledge_search"):
                         result = await self._execute_knowledge_search_via_vector_store(
                             query=query,
                             response_file_search_tool=response_file_search_tool,
@@ -301,7 +356,9 @@ class ToolExecutor:
                 attributes = {
                     "tool_name": function_name,
                 }
-                async with tracing.span("invoke_tool", attributes):
+                # TODO: follow semantic conventions for Open Telemetry tool spans
+                # https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span
+                with tracer.start_as_current_span("invoke_tool", attributes=attributes):
                     result = await self.tool_runtime_api.invoke_tool(
                         tool_name=function_name,
                         kwargs=tool_kwargs,
@@ -322,35 +379,34 @@ class ToolExecutor:
         mcp_tool_to_server: dict[str, OpenAIResponseInputToolMCP] | None = None,
     ) -> AsyncIterator[ToolExecutionResult]:
         """Emit completion or failure events for tool execution."""
-        completion_event = None
-
         if mcp_tool_to_server and function_name in mcp_tool_to_server:
             sequence_number += 1
             if has_error:
-                completion_event = OpenAIResponseObjectStreamResponseMcpCallFailed(
+                mcp_failed_event = OpenAIResponseObjectStreamResponseMcpCallFailed(
                     sequence_number=sequence_number,
                 )
+                yield ToolExecutionResult(stream_event=mcp_failed_event, sequence_number=sequence_number)
             else:
-                completion_event = OpenAIResponseObjectStreamResponseMcpCallCompleted(
+                mcp_completed_event = OpenAIResponseObjectStreamResponseMcpCallCompleted(
                     sequence_number=sequence_number,
                 )
+                yield ToolExecutionResult(stream_event=mcp_completed_event, sequence_number=sequence_number)
         elif function_name == "web_search":
             sequence_number += 1
-            completion_event = OpenAIResponseObjectStreamResponseWebSearchCallCompleted(
+            web_completion_event = OpenAIResponseObjectStreamResponseWebSearchCallCompleted(
                 item_id=item_id,
                 output_index=output_index,
                 sequence_number=sequence_number,
             )
+            yield ToolExecutionResult(stream_event=web_completion_event, sequence_number=sequence_number)
         elif function_name == "knowledge_search":
             sequence_number += 1
-            completion_event = OpenAIResponseObjectStreamResponseFileSearchCallCompleted(
+            file_completion_event = OpenAIResponseObjectStreamResponseFileSearchCallCompleted(
                 item_id=item_id,
                 output_index=output_index,
                 sequence_number=sequence_number,
             )
-
-        if completion_event:
-            yield ToolExecutionResult(stream_event=completion_event, sequence_number=sequence_number)
+            yield ToolExecutionResult(stream_event=file_completion_event, sequence_number=sequence_number)
 
     async def _build_result_messages(
         self,
@@ -360,18 +416,19 @@ class ToolExecutor:
         tool_kwargs: dict,
         ctx: ChatCompletionContext,
         error_exc: Exception | None,
-        result: any,
+        result: Any,
         has_error: bool,
         mcp_tool_to_server: dict[str, OpenAIResponseInputToolMCP] | None = None,
-    ) -> tuple[any, any]:
+    ) -> tuple[Any, Any]:
         """Build output and input messages from tool execution results."""
         from llama_stack.providers.utils.inference.prompt_adapter import (
             interleaved_content_as_str,
         )
 
         # Build output message
+        message: Any
         if mcp_tool_to_server and function.name in mcp_tool_to_server:
-            from llama_stack.apis.agents.openai_responses import (
+            from llama_stack_api import (
                 OpenAIResponseOutputMessageMCPCall,
             )
 
@@ -383,10 +440,14 @@ class ToolExecutor:
             )
             if error_exc:
                 message.error = str(error_exc)
-            elif (result and result.error_code and result.error_code > 0) or (result and result.error_message):
-                message.error = f"Error (code {result.error_code}): {result.error_message}"
-            elif result and result.content:
-                message.output = interleaved_content_as_str(result.content)
+            elif (result and (error_code := getattr(result, "error_code", None)) and error_code > 0) or (
+                result and getattr(result, "error_message", None)
+            ):
+                ec = getattr(result, "error_code", "unknown")
+                em = getattr(result, "error_message", "")
+                message.error = f"Error (code {ec}): {em}"
+            elif result and (content := getattr(result, "content", None)):
+                message.output = interleaved_content_as_str(content)
         else:
             if function.name == "web_search":
                 message = OpenAIResponseOutputMessageWebSearchToolCall(
@@ -401,17 +462,17 @@ class ToolExecutor:
                     queries=[tool_kwargs.get("query", "")],
                     status="completed",
                 )
-                if result and "document_ids" in result.metadata:
+                if result and (metadata := getattr(result, "metadata", None)) and "document_ids" in metadata:
                     message.results = []
-                    for i, doc_id in enumerate(result.metadata["document_ids"]):
-                        text = result.metadata["chunks"][i] if "chunks" in result.metadata else None
-                        score = result.metadata["scores"][i] if "scores" in result.metadata else None
+                    for i, doc_id in enumerate(metadata["document_ids"]):
+                        text = metadata["chunks"][i] if "chunks" in metadata else None
+                        score = metadata["scores"][i] if "scores" in metadata else None
                         message.results.append(
                             OpenAIResponseOutputMessageFileSearchToolCallResults(
                                 file_id=doc_id,
                                 filename=doc_id,
-                                text=text,
-                                score=score,
+                                text=text if text is not None else "",
+                                score=score if score is not None else 0.0,
                                 attributes={},
                             )
                         )
@@ -421,27 +482,32 @@ class ToolExecutor:
                 raise ValueError(f"Unknown tool {function.name} called")
 
         # Build input message
-        input_message = None
-        if result and result.content:
-            if isinstance(result.content, str):
-                content = result.content
-            elif isinstance(result.content, list):
-                content = []
-                for item in result.content:
+        input_message: OpenAIToolMessageParam | None = None
+        if result and (result_content := getattr(result, "content", None)):
+            # all the mypy contortions here are still unsatisfactory with random Any typing
+            if isinstance(result_content, str):
+                msg_content: str | list[Any] = result_content
+            elif isinstance(result_content, list):
+                content_list: list[Any] = []
+                for item in result_content:
+                    part: Any
                     if isinstance(item, TextContentItem):
                         part = OpenAIChatCompletionContentPartTextParam(text=item.text)
                     elif isinstance(item, ImageContentItem):
                         if item.image.data:
-                            url = f"data:image;base64,{item.image.data}"
+                            url_value = f"data:image;base64,{item.image.data}"
                         else:
-                            url = item.image.url
-                        part = OpenAIChatCompletionContentPartImageParam(image_url=OpenAIImageURL(url=url))
+                            url_value = str(item.image.url) if item.image.url else ""
+                        part = OpenAIChatCompletionContentPartImageParam(image_url=OpenAIImageURL(url=url_value))
                     else:
                         raise ValueError(f"Unknown result content type: {type(item)}")
-                    content.append(part)
+                    content_list.append(part)
+                msg_content = content_list
             else:
-                raise ValueError(f"Unknown result content type: {type(result.content)}")
-            input_message = OpenAIToolMessageParam(content=content, tool_call_id=tool_call_id)
+                raise ValueError(f"Unknown result content type: {type(result_content)}")
+            # OpenAIToolMessageParam accepts str | list[TextParam] but we may have images
+            # This is runtime-safe as the API accepts it, but mypy complains
+            input_message = OpenAIToolMessageParam(content=msg_content, tool_call_id=tool_call_id)  # type: ignore[arg-type]
         else:
             text = str(error_exc) if error_exc else "Tool execution failed"
             input_message = OpenAIToolMessageParam(content=text, tool_call_id=tool_call_id)

llama_stack/providers/inline/agents/meta_reference/responses/types.py

@@ -5,13 +5,18 @@
 # the root directory of this source tree.
 
 from dataclasses import dataclass
+from typing import cast
 
 from openai.types.chat import ChatCompletionToolParam
 from pydantic import BaseModel
 
-from llama_stack.apis.agents.openai_responses import (
+from llama_stack_api import (
+    OpenAIChatCompletionToolCall,
+    OpenAIMessageParam,
+    OpenAIResponseFormatParam,
     OpenAIResponseInput,
     OpenAIResponseInputTool,
+    OpenAIResponseInputToolChoice,
     OpenAIResponseInputToolFileSearch,
     OpenAIResponseInputToolFunction,
     OpenAIResponseInputToolMCP,
@@ -24,8 +29,8 @@ from llama_stack.apis.agents.openai_responses import (
     OpenAIResponseOutputMessageMCPListTools,
     OpenAIResponseTool,
     OpenAIResponseToolMCP,
+    OpenAITokenLogProb,
 )
-from llama_stack.apis.inference import OpenAIChatCompletionToolCall, OpenAIMessageParam, OpenAIResponseFormatParam
 
 
 class ToolExecutionResult(BaseModel):
@@ -51,6 +56,7 @@ class ChatCompletionResult:
     message_item_id: str  # For streaming events
     tool_call_item_ids: dict[int, str]  # For streaming events
     content_part_emitted: bool  # Tracking state
+    logprobs: list[OpenAITokenLogProb] | None = None
 
     @property
     def content_text(self) -> str:
@@ -100,17 +106,19 @@ class ToolContext(BaseModel):
                 if isinstance(tool, OpenAIResponseToolMCP):
                     previous_tools_by_label[tool.server_label] = tool
             # collect tool definitions which are the same in current and previous requests:
-            tools_to_process = []
+            tools_to_process: list[OpenAIResponseInputTool] = []
             matched: dict[str, OpenAIResponseInputToolMCP] = {}
-            for tool in self.current_tools:
+            # Mypy confuses OpenAIResponseInputTool (Input union) with OpenAIResponseTool (output union)
+            # which differ only in MCP type (InputToolMCP vs ToolMCP). Code is correct.
+            for tool in cast(list[OpenAIResponseInputTool], self.current_tools):  # type: ignore[assignment]
                 if isinstance(tool, OpenAIResponseInputToolMCP) and tool.server_label in previous_tools_by_label:
                     previous_tool = previous_tools_by_label[tool.server_label]
                     if previous_tool.allowed_tools == tool.allowed_tools:
                         matched[tool.server_label] = tool
                     else:
-                        tools_to_process.append(tool)
+                        tools_to_process.append(tool)  # type: ignore[arg-type]
                 else:
-                    tools_to_process.append(tool)
+                    tools_to_process.append(tool)  # type: ignore[arg-type]
             # tools that are not the same or were not previously defined need to be processed:
             self.tools_to_process = tools_to_process
             # for all matched definitions, get the mcp_list_tools objects from the previous output:
@@ -119,9 +127,11 @@ class ToolContext(BaseModel):
             ]
             # reconstruct the tool to server mappings that can be reused:
             for listing in self.previous_tool_listings:
+                # listing is OpenAIResponseOutputMessageMCPListTools which has tools: list[MCPListToolsTool]
                 definition = matched[listing.server_label]
-                for tool in listing.tools:
-                    self.previous_tools[tool.name] = definition
+                for mcp_tool in listing.tools:
+                    # mcp_tool is MCPListToolsTool which has a name: str field
+                    self.previous_tools[mcp_tool.name] = definition
 
     def available_tools(self) -> list[OpenAIResponseTool]:
         if not self.current_tools:
@@ -139,6 +149,8 @@ class ToolContext(BaseModel):
                     server_label=tool.server_label,
                     allowed_tools=tool.allowed_tools,
                 )
+            # Exhaustive check - all tool types should be handled above
+            raise AssertionError(f"Unexpected tool type: {type(tool)}")
 
         return [convert_tool(tool) for tool in self.current_tools]
 
@@ -151,6 +163,7 @@ class ChatCompletionContext(BaseModel):
     temperature: float | None
     response_format: OpenAIResponseFormatParam
     tool_context: ToolContext | None
+    tool_choice: OpenAIResponseInputToolChoice | None = None
     approval_requests: list[OpenAIResponseMCPApprovalRequest] = []
     approval_responses: dict[str, OpenAIResponseMCPApprovalResponse] = {}
 
@@ -163,6 +176,7 @@ class ChatCompletionContext(BaseModel):
         response_format: OpenAIResponseFormatParam,
         tool_context: ToolContext,
         inputs: list[OpenAIResponseInput] | str,
+        tool_choice: OpenAIResponseInputToolChoice | None = None,
     ):
         super().__init__(
             model=model,
@@ -171,6 +185,7 @@ class ChatCompletionContext(BaseModel):
             temperature=temperature,
             response_format=response_format,
             tool_context=tool_context,
+            tool_choice=tool_choice,
         )
         if not isinstance(inputs, str):
             self.approval_requests = [input for input in inputs if input.type == "mcp_approval_request"]