data-designer-engine 0.4.0rc3__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

data_designer/engine/model_provider.py

@@ -8,6 +8,7 @@ from functools import cached_property
 from pydantic import BaseModel, field_validator, model_validator
 from typing_extensions import Self
 
+from data_designer.config.mcp import MCPProviderT
 from data_designer.config.models import ModelProvider
 from data_designer.engine.errors import NoModelProvidersError, UnknownProviderError
 
@@ -75,3 +76,70 @@ def resolve_model_provider_registry(
         providers=model_providers,
         default=default_provider_name or model_providers[0].name,
     )
+
+
+class MCPProviderRegistry(BaseModel):
+    """Registry for MCP providers.
+
+    Unlike ModelProviderRegistry, MCPProviderRegistry can be empty since MCP providers
+    are optional. Users only need to register MCP providers if they want to use MCP tools
+    for generation.
+
+    Attributes:
+        providers: List of MCP providers (both MCPProvider and LocalStdioMCPProvider).
+    """
+
+    providers: list[MCPProviderT] = []
+
+    @field_validator("providers", mode="after")
+    @classmethod
+    def validate_providers_have_unique_names(cls, v: list[MCPProviderT]) -> list[MCPProviderT]:
+        names = set()
+        dupes = set()
+        for provider in v:
+            if provider.name in names:
+                dupes.add(provider.name)
+            names.add(provider.name)
+
+        if len(dupes) > 0:
+            raise ValueError(f"MCP providers must have unique names, found duplicates: {dupes}")
+        return v
+
+    @cached_property
+    def _providers_dict(self) -> dict[str, MCPProviderT]:
+        return {p.name: p for p in self.providers}
+
+    def get_provider(self, name: str) -> MCPProviderT:
+        """Get an MCP provider by name.
+
+        Args:
+            name: The name of the MCP provider.
+
+        Returns:
+            The MCP provider with the given name.
+
+        Raises:
+            UnknownProviderError: If no provider with the given name is registered.
+        """
+        try:
+            return self._providers_dict[name]
+        except KeyError:
+            raise UnknownProviderError(f"No MCP provider named {name!r} registered")
+
+    def is_empty(self) -> bool:
+        """Check if the registry has no providers."""
+        return len(self.providers) == 0
+
+
+def resolve_mcp_provider_registry(
+    mcp_providers: list[MCPProviderT] | None = None,
+) -> MCPProviderRegistry:
+    """Create an MCPProviderRegistry from a list of MCP providers.
+
+    Args:
+        mcp_providers: Optional list of MCP providers. If None or empty, returns an empty registry.
+
+    Returns:
+        An MCPProviderRegistry containing the provided MCP providers.
+    """
+    return MCPProviderRegistry(providers=mcp_providers or [])

data_designer/engine/models/facade.py

@@ -9,6 +9,7 @@ from copy import deepcopy
 from typing import TYPE_CHECKING, Any
 
 from data_designer.config.models import GenerationType, ModelConfig, ModelProvider
+from data_designer.engine.mcp.errors import MCPConfigurationError
 from data_designer.engine.model_provider import ModelProviderRegistry
 from data_designer.engine.models.errors import (
     GenerationValidationFailureError,
@@ -25,6 +26,15 @@ from data_designer.lazy_heavy_imports import litellm
 if TYPE_CHECKING:
     import litellm
 
+    from data_designer.engine.mcp.facade import MCPFacade
+    from data_designer.engine.mcp.registry import MCPRegistry
+
+
+def _identity(x: Any) -> Any:
+    """Identity function for default parser. Module-level for pickling compatibility."""
+    return x
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -34,10 +44,13 @@ class ModelFacade:
         model_config: ModelConfig,
         secret_resolver: SecretResolver,
         model_provider_registry: ModelProviderRegistry,
-    ):
+        *,
+        mcp_registry: MCPRegistry | None = None,
+    ) -> None:
         self._model_config = model_config
         self._secret_resolver = secret_resolver
         self._model_provider_registry = model_provider_registry
+        self._mcp_registry = mcp_registry
         self._litellm_deployment = self._get_litellm_deployment(model_config)
         self._router = CustomRouter([self._litellm_deployment], **LiteLLMRouterDefaultKwargs().model_dump())
         self._usage_stats = ModelUsageStats()
@@ -104,6 +117,17 @@ class ModelFacade:
             kwargs["extra_headers"] = self.model_provider.extra_headers
         return kwargs
 
+    def _get_mcp_facade(self, tool_alias: str | None) -> MCPFacade | None:
+        if tool_alias is None:
+            return None
+        if self._mcp_registry is None:
+            raise MCPConfigurationError(f"Tool alias {tool_alias!r} specified but no MCPRegistry configured.")
+
+        try:
+            return self._mcp_registry.get_mcp(tool_alias=tool_alias)
+        except ValueError as exc:
+            raise MCPConfigurationError(f"Tool alias {tool_alias!r} is not registered.") from exc
+
     @catch_llm_exceptions
     def generate_text_embeddings(
         self, input_texts: list[str], skip_usage_tracking: bool = False, **kwargs
@@ -142,9 +166,10 @@ class ModelFacade:
         self,
         prompt: str,
         *,
-        parser: Callable[[str], Any],
+        parser: Callable[[str], Any] = _identity,
         system_prompt: str | None = None,
         multi_modal_context: list[dict[str, Any]] | None = None,
+        tool_alias: str | None = None,
         max_correction_steps: int = 0,
         max_conversation_restarts: int = 0,
         skip_usage_tracking: bool = False,
@@ -170,7 +195,10 @@ class ModelFacade:
                 no system message is provided and the model should use its default system
                 prompt.
             parser (func(str) -> Any): A function applied to the LLM response which processes
-                an LLM response into some output object.
+                an LLM response into some output object. Default: identity function.
+            tool_alias (str | None): Optional tool configuration alias. When provided,
+                the model may call permitted tools from the configured MCP providers.
+                The alias must reference a ToolConfig registered in the MCPRegistry.
             max_correction_steps (int): Maximum number of correction rounds permitted
                 within a single conversation. Note, many rounds can lead to increasing
                 context size without necessarily improving performance -- small language
@@ -186,25 +214,61 @@ class ModelFacade:
         Returns:
             A tuple containing:
                 - The parsed output object from the parser.
-                - The full trace of ChatMessage entries in the conversation, including any
-                  corrections and reasoning traces. Callers can decide whether to store this.
+                - The full trace of ChatMessage entries in the conversation, including any tool calls,
+                  corrections, and reasoning traces. Callers can decide whether to store this.
 
         Raises:
             GenerationValidationFailureError: If the maximum number of retries or
                 correction steps are met and the last response failures on
                 generation validation.
+            MCPConfigurationError: If tool_alias is specified but no MCPRegistry is configured.
         """
         output_obj = None
+        tool_schemas = None
+        tool_call_turns = 0
         curr_num_correction_steps = 0
         curr_num_restarts = 0
 
-        starting_messages = prompt_to_messages(
+        mcp_facade = self._get_mcp_facade(tool_alias)
+
+        # Checkpoint for restarts - updated after tool calls so we don't repeat them
+        restart_checkpoint = prompt_to_messages(
             user_prompt=prompt, system_prompt=system_prompt, multi_modal_context=multi_modal_context
         )
-        messages: list[ChatMessage] = deepcopy(starting_messages)
+        checkpoint_tool_call_turns = 0
+        messages: list[ChatMessage] = deepcopy(restart_checkpoint)
+
+        if mcp_facade is not None:
+            tool_schemas = mcp_facade.get_tool_schemas()
 
         while True:
-            completion_response = self.completion(messages, skip_usage_tracking=skip_usage_tracking, **kwargs)
+            completion_kwargs = dict(kwargs)
+            if tool_schemas is not None:
+                completion_kwargs["tools"] = tool_schemas
+
+            completion_response = self.completion(
+                messages,
+                skip_usage_tracking=skip_usage_tracking,
+                **completion_kwargs,
+            )
+
+            # Process any tool calls in the response (handles parallel tool calling)
+            if mcp_facade is not None and mcp_facade.has_tool_calls(completion_response):
+                tool_call_turns += 1
+
+                if tool_call_turns > mcp_facade.max_tool_call_turns:
+                    # Gracefully refuse tool calls when budget is exhausted
+                    messages.extend(mcp_facade.refuse_completion_response(completion_response))
+                else:
+                    messages.extend(mcp_facade.process_completion_response(completion_response))
+
+                # Update checkpoint so restarts don't repeat tool calls
+                restart_checkpoint = deepcopy(messages)
+                checkpoint_tool_call_turns = tool_call_turns
+
+                continue  # Back to top
+
+            # No tool calls remaining to process
             response = completion_response.choices[0].message.content or ""
             reasoning_trace = getattr(completion_response.choices[0].message, "reasoning_content", None)
             messages.append(ChatMessage.as_assistant(content=response, reasoning_content=reasoning_trace or None))
@@ -226,7 +290,8 @@ class ModelFacade:
                 elif curr_num_restarts < max_conversation_restarts:
                     curr_num_correction_steps = 0
                     curr_num_restarts += 1
-                    messages = deepcopy(starting_messages)
+                    messages = deepcopy(restart_checkpoint)
+                    tool_call_turns = checkpoint_tool_call_turns
 
                 else:
                     raise GenerationValidationFailureError(

data_designer/engine/models/factory.py

@@ -10,6 +10,7 @@ from data_designer.engine.model_provider import ModelProviderRegistry
 from data_designer.engine.secret_resolver import SecretResolver
 
 if TYPE_CHECKING:
+    from data_designer.engine.mcp.registry import MCPRegistry
     from data_designer.engine.models.registry import ModelRegistry
 
 
@@ -18,12 +19,23 @@ def create_model_registry(
     model_configs: list[ModelConfig] | None = None,
     secret_resolver: SecretResolver,
     model_provider_registry: ModelProviderRegistry,
+    mcp_registry: MCPRegistry | None = None,
 ) -> ModelRegistry:
     """Factory function for creating a ModelRegistry instance.
 
     Heavy dependencies (litellm, httpx) are deferred until this function is called.
     This is a factory function pattern - imports inside factories are idiomatic Python
     for lazy initialization.
+
+    Args:
+        model_configs: Optional list of model configurations to register.
+        secret_resolver: Resolver for secrets referenced in provider configs.
+        model_provider_registry: Registry of model provider configurations.
+        mcp_registry: Optional MCP registry for tool operations. When provided,
+            ModelFacades can look up MCPFacades by tool_alias for tool-enabled generation.
+
+    Returns:
+        A configured ModelRegistry instance.
     """
     from data_designer.engine.models.facade import ModelFacade
     from data_designer.engine.models.litellm_overrides import apply_litellm_patches
@@ -32,7 +44,12 @@ def create_model_registry(
     apply_litellm_patches()
 
     def model_facade_factory(model_config, secret_resolver, model_provider_registry):
-        return ModelFacade(model_config, secret_resolver, model_provider_registry)
+        return ModelFacade(
+            model_config,
+            secret_resolver,
+            model_provider_registry,
+            mcp_registry=mcp_registry,
+        )
 
     return ModelRegistry(
         model_configs=model_configs,

data_designer/engine/models/utils.py

@@ -36,11 +36,14 @@ class ChatMessage:
     def to_dict(self) -> dict[str, Any]:
         """Convert the message to a dictionary format for API calls.
 
+        Content is normalized to a list of ChatML-style blocks to keep a
+        consistent schema across traces and API payloads.
+
         Returns:
             A dictionary containing the message fields. Only includes non-empty
             optional fields to keep the output clean.
         """
-        result: dict[str, Any] = {"role": self.role, "content": self.content}
+        result: dict[str, Any] = {"role": self.role, "content": _normalize_content_blocks(self.content)}
         if self.reasoning_content:
             result["reasoning_content"] = self.reasoning_content
         if self.tool_calls:
@@ -99,3 +102,27 @@ def prompt_to_messages(
     if system_prompt:
         return [ChatMessage.as_system(system_prompt), ChatMessage.as_user(user_content)]
     return [ChatMessage.as_user(user_content)]
+
+
+def _normalize_content_blocks(content: Any) -> list[dict[str, Any]]:
+    if isinstance(content, list):
+        return [_normalize_content_block(block) for block in content]
+    if content is None:
+        return []
+    return [_text_block(content)]
+
+
+def _normalize_content_block(block: Any) -> dict[str, Any]:
+    if isinstance(block, dict) and "type" in block:
+        return block
+    return _text_block(block)
+
+
+def _text_block(value: Any) -> dict[str, Any]:
+    if value is None:
+        text_value = ""
+    elif isinstance(value, str):
+        text_value = value
+    else:
+        text_value = str(value)
+    return {"type": "text", "text": text_value}

data_designer/engine/resources/resource_provider.py

@@ -5,15 +5,21 @@ from __future__ import annotations
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.dataset_metadata import DatasetMetadata
+from data_designer.config.mcp import MCPProviderT, ToolConfig
 from data_designer.config.models import ModelConfig
 from data_designer.config.run_config import RunConfig
 from data_designer.config.seed_source import SeedSource
 from data_designer.config.utils.type_helpers import StrEnum
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
-from data_designer.engine.model_provider import ModelProviderRegistry
+from data_designer.engine.mcp.factory import create_mcp_registry
+from data_designer.engine.mcp.registry import MCPRegistry
+from data_designer.engine.model_provider import (
+    ModelProviderRegistry,
+    resolve_mcp_provider_registry,
+)
 from data_designer.engine.models.factory import create_model_registry
 from data_designer.engine.models.registry import ModelRegistry
-from data_designer.engine.resources.managed_storage import ManagedBlobStorage, init_managed_blob_storage
+from data_designer.engine.resources.managed_storage import ManagedBlobStorage
 from data_designer.engine.resources.seed_reader import SeedReader, SeedReaderRegistry
 from data_designer.engine.secret_resolver import SecretResolver
 
@@ -28,6 +34,7 @@ class ResourceProvider(ConfigBase):
     artifact_storage: ArtifactStorage
     blob_storage: ManagedBlobStorage | None = None
     model_registry: ModelRegistry | None = None
+    mcp_registry: MCPRegistry | None = None
     run_config: RunConfig = RunConfig()
     seed_reader: SeedReader | None = None
 
@@ -43,6 +50,31 @@ class ResourceProvider(ConfigBase):
         return DatasetMetadata(seed_column_names=seed_column_names)
 
 
+def _validate_tool_configs_against_providers(
+    tool_configs: list[ToolConfig],
+    mcp_providers: list[MCPProviderT],
+) -> None:
+    """Validate that all providers referenced in tool configs exist.
+
+    Args:
+        tool_configs: List of tool configurations to validate.
+        mcp_providers: List of available MCP provider configurations.
+
+    Raises:
+        ValueError: If a tool config references a provider that doesn't exist.
+    """
+    available_providers = {p.name for p in mcp_providers}
+
+    for tc in tool_configs:
+        missing_providers = [p for p in tc.providers if p not in available_providers]
+        if missing_providers:
+            available_list = sorted(available_providers) if available_providers else ["(none configured)"]
+            raise ValueError(
+                f"ToolConfig '{tc.tool_alias}' references provider(s) {missing_providers!r} "
+                f"which are not registered. Available providers: {available_list}"
+            )
+
+
 def create_resource_provider(
     *,
     artifact_storage: ArtifactStorage,
@@ -53,9 +85,31 @@ def create_resource_provider(
     blob_storage: ManagedBlobStorage | None = None,
     seed_dataset_source: SeedSource | None = None,
     run_config: RunConfig | None = None,
+    mcp_providers: list[MCPProviderT] | None = None,
+    tool_configs: list[ToolConfig] | None = None,
 ) -> ResourceProvider:
     """Factory function for creating a ResourceProvider instance.
+
     This function triggers lazy loading of heavy dependencies like litellm.
+    The creation order is:
+    1. MCPProviderRegistry (can be empty)
+    2. MCPRegistry with tool_configs
+    3. ModelRegistry with mcp_registry
+
+    Args:
+        artifact_storage: Storage for build artifacts.
+        model_configs: List of model configurations.
+        secret_resolver: Resolver for secrets.
+        model_provider_registry: Registry of model providers.
+        seed_reader_registry: Registry of seed readers.
+        blob_storage: Optional blob storage for large files.
+        seed_dataset_source: Optional source for seed datasets.
+        run_config: Optional runtime configuration.
+        mcp_providers: Optional list of MCP provider configurations.
+        tool_configs: Optional list of tool configurations.
+
+    Returns:
+        A configured ResourceProvider instance.
     """
     seed_reader = None
     if seed_dataset_source:
@@ -64,14 +118,29 @@ def create_resource_provider(
             secret_resolver,
         )
 
+    # Create MCPProviderRegistry first (can be empty)
+    mcp_provider_registry = resolve_mcp_provider_registry(mcp_providers)
+
+    # Create MCPRegistry with tool configs (only if tool_configs provided)
+    # Tool validation is performed during dataset builder health checks.
+    mcp_registry = None
+    if tool_configs:
+        mcp_registry = create_mcp_registry(
+            tool_configs=tool_configs,
+            secret_resolver=secret_resolver,
+            mcp_provider_registry=mcp_provider_registry,
+        )
+
     return ResourceProvider(
         artifact_storage=artifact_storage,
         model_registry=create_model_registry(
             model_configs=model_configs,
             secret_resolver=secret_resolver,
             model_provider_registry=model_provider_registry,
+            mcp_registry=mcp_registry,
         ),
-        blob_storage=blob_storage or init_managed_blob_storage(),
+        blob_storage=blob_storage,
+        mcp_registry=mcp_registry,
         seed_reader=seed_reader,
         run_config=run_config or RunConfig(),
     )

data_designer/engine/testing/fixtures.py

@@ -0,0 +1,233 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pytest fixtures for engine testing.
+
+Located in src/ so it can be packaged and shared across subpackages via pytest_plugins.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider, ToolConfig
+from data_designer.engine.mcp.facade import MCPFacade
+from data_designer.engine.model_provider import MCPProviderRegistry
+from data_designer.engine.secret_resolver import SecretResolver
+from data_designer.engine.testing.stubs import StubHuggingFaceSeedReader
+
+# =============================================================================
+# Fake LLM response classes (used by completion response fixtures)
+# =============================================================================
+
+
+class _FakeMessage:
+    """Fake message class for mocking LLM completion responses."""
+
+    def __init__(
+        self,
+        content: str | None,
+        tool_calls: list[dict] | None = None,
+        reasoning_content: str | None = None,
+    ) -> None:
+        self.content = content
+        self.tool_calls = tool_calls
+        self.reasoning_content = reasoning_content
+
+
+class _FakeChoice:
+    """Fake choice class for mocking LLM completion responses."""
+
+    def __init__(self, message: _FakeMessage) -> None:
+        self.message = message
+
+
+class _FakeResponse:
+    """Fake response class for mocking LLM completion responses."""
+
+    def __init__(self, message: _FakeMessage) -> None:
+        self.choices = [_FakeChoice(message)]
+
+
+# =============================================================================
+# Seed reader fixtures
+# =============================================================================
+
+
+@pytest.fixture
+def stub_seed_reader() -> StubHuggingFaceSeedReader:
+    """Stub seed reader for testing seed dataset functionality."""
+    return StubHuggingFaceSeedReader()
+
+
+# =============================================================================
+# MCP Provider fixtures
+# =============================================================================
+
+
+@pytest.fixture
+def stub_mcp_provider_registry() -> MCPProviderRegistry:
+    """Create a stub MCP provider registry with test providers."""
+    return MCPProviderRegistry(
+        providers=[
+            LocalStdioMCPProvider(name="tools", command="python"),
+            LocalStdioMCPProvider(name="secondary", command="python"),
+        ]
+    )
+
+
+@pytest.fixture
+def stub_mcp_provider_registry_single() -> MCPProviderRegistry:
+    """Create a stub MCP provider registry with a single provider."""
+    return MCPProviderRegistry(providers=[LocalStdioMCPProvider(name="tools", command="python")])
+
+
+@pytest.fixture
+def stub_secret_resolver() -> MagicMock:
+    """Create a stub secret resolver for testing."""
+    resolver = MagicMock(spec=SecretResolver)
+    resolver.resolve.side_effect = lambda x: x  # Return the input as-is
+    return resolver
+
+
+@pytest.fixture
+def stub_stdio_provider() -> LocalStdioMCPProvider:
+    """Create a stub stdio MCP provider for testing."""
+    return LocalStdioMCPProvider(
+        name="test-stdio",
+        command="python",
+        args=["-m", "test_server"],
+        env={"TEST_VAR": "value"},
+    )
+
+
+@pytest.fixture
+def stub_sse_provider() -> MCPProvider:
+    """Create a stub SSE MCP provider for testing."""
+    return MCPProvider(
+        name="test-sse",
+        endpoint="http://localhost:8080/sse",
+        api_key="test-key",
+    )
+
+
+# =============================================================================
+# Tool config fixtures
+# =============================================================================
+
+
+@pytest.fixture
+def stub_tool_config() -> ToolConfig:
+    """Create a basic tool configuration for testing."""
+    return ToolConfig(
+        tool_alias="test-tools",
+        providers=["tools"],
+        max_tool_call_turns=3,
+        timeout_sec=30.0,
+    )
+
+
+@pytest.fixture
+def stub_tool_config_with_allow_list() -> ToolConfig:
+    """Create a tool configuration with an allow list."""
+    return ToolConfig(
+        tool_alias="test-tools",
+        providers=["tools"],
+        allow_tools=["lookup", "search"],
+        max_tool_call_turns=3,
+    )
+
+
+# =============================================================================
+# Facade fixtures
+# =============================================================================
+
+
+@pytest.fixture
+def stub_mcp_facade(
+    stub_tool_config: ToolConfig, stub_secret_resolver: MagicMock, stub_mcp_provider_registry: MCPProviderRegistry
+) -> MCPFacade:
+    """Create a stub MCPFacade for testing."""
+    return MCPFacade(
+        tool_config=stub_tool_config,
+        secret_resolver=stub_secret_resolver,
+        mcp_provider_registry=stub_mcp_provider_registry,
+    )
+
+
+@pytest.fixture
+def stub_mcp_facade_factory() -> Any:
+    """Create a stub MCP facade factory for testing."""
+
+    def factory(
+        tool_config: ToolConfig, secret_resolver: SecretResolver, provider_registry: MCPProviderRegistry
+    ) -> MCPFacade:
+        return MCPFacade(
+            tool_config=tool_config, secret_resolver=secret_resolver, mcp_provider_registry=provider_registry
+        )
+
+    return factory
+
+
+# =============================================================================
+# Completion response fixtures
+# =============================================================================
+
+
+@pytest.fixture
+def mock_completion_response_no_tools() -> _FakeResponse:
+    """Mock LLM response with no tool calls."""
+    return _FakeResponse(_FakeMessage(content="Hello, I can help with that."))
+
+
+@pytest.fixture
+def mock_completion_response_single_tool() -> _FakeResponse:
+    """Mock LLM response with single tool call."""
+    tool_call = {
+        "id": "call-1",
+        "type": "function",
+        "function": {"name": "lookup", "arguments": '{"query": "test"}'},
+    }
+    return _FakeResponse(_FakeMessage(content="Let me look that up.", tool_calls=[tool_call]))
+
+
+@pytest.fixture
+def mock_completion_response_parallel_tools() -> _FakeResponse:
+    """Mock LLM response with multiple parallel tool calls."""
+    tool_calls = [
+        {"id": "call-1", "type": "function", "function": {"name": "lookup", "arguments": '{"query": "first"}'}},
+        {"id": "call-2", "type": "function", "function": {"name": "search", "arguments": '{"term": "second"}'}},
+        {"id": "call-3", "type": "function", "function": {"name": "fetch", "arguments": '{"url": "example.com"}'}},
+    ]
+    return _FakeResponse(_FakeMessage(content="Executing multiple tools.", tool_calls=tool_calls))
+
+
+@pytest.fixture
+def mock_completion_response_with_reasoning() -> _FakeResponse:
+    """Mock LLM response with reasoning_content."""
+    return _FakeResponse(
+        _FakeMessage(
+            content="  Final answer with extra spaces.  ",
+            reasoning_content="  Thinking about the problem...  ",
+        )
+    )
+
+
+@pytest.fixture
+def mock_completion_response_tool_with_reasoning() -> _FakeResponse:
+    """Mock LLM response with tool calls and reasoning_content."""
+    tool_call = {
+        "id": "call-1",
+        "type": "function",
+        "function": {"name": "lookup", "arguments": '{"query": "test"}'},
+    }
+    return _FakeResponse(
+        _FakeMessage(
+            content="  Looking it up...  ",
+            tool_calls=[tool_call],
+            reasoning_content="  I should use the lookup tool.  ",
+        )
+    )