data-designer-engine 0.4.0rc2__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,
@@ -18,13 +19,22 @@ from data_designer.engine.models.errors import (
 from data_designer.engine.models.litellm_overrides import CustomRouter, LiteLLMRouterDefaultKwargs
 from data_designer.engine.models.parsers.errors import ParserException
 from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenUsageStats
-from data_designer.engine.models.utils import prompt_to_messages, str_to_message
+from data_designer.engine.models.utils import ChatMessage, prompt_to_messages
 from data_designer.engine.secret_resolver import SecretResolver
 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()
@@ -67,16 +80,17 @@ class ModelFacade:
         return self._usage_stats
 
     def completion(
-        self, messages: list[dict[str, str]], skip_usage_tracking: bool = False, **kwargs
+        self, messages: list[ChatMessage], skip_usage_tracking: bool = False, **kwargs
     ) -> litellm.ModelResponse:
+        message_payloads = [message.to_dict() for message in messages]
         logger.debug(
             f"Prompting model {self.model_name!r}...",
-            extra={"model": self.model_name, "messages": messages},
+            extra={"model": self.model_name, "messages": message_payloads},
         )
         response = None
         kwargs = self.consolidate_kwargs(**kwargs)
         try:
-            response = self._router.completion(model=self.model_name, messages=messages, **kwargs)
+            response = self._router.completion(model=self.model_name, messages=message_payloads, **kwargs)
             logger.debug(
                 f"Received completion from model {self.model_name!r}",
                 extra={
@@ -103,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
@@ -141,15 +166,16 @@ 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,
         purpose: str | None = None,
         **kwargs,
-    ) -> tuple[Any, str | None]:
+    ) -> tuple[Any, list[ChatMessage]]:
         """Generate a parsed output with correction steps.
 
         This generation call will attempt to generate an output which is
@@ -169,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
@@ -182,37 +211,67 @@ class ModelFacade:
                 It is expected to be used by the @catch_llm_exceptions decorator.
             **kwargs: Additional arguments to pass to the model.
 
+        Returns:
+            A tuple containing:
+                - The parsed output object from the parser.
+                - 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
-        curr_generation_attempt = 0
-        max_generation_attempts = (max_correction_steps + 1) * (max_conversation_restarts + 1)
 
-        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 = 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:
-            curr_generation_attempt += 1
-            logger.debug(
-                f"Starting generation attempt {curr_generation_attempt} of {max_generation_attempts} attempts."
+            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,
             )
 
-            completion_response = self.completion(messages, skip_usage_tracking=skip_usage_tracking, **kwargs)
-            response = completion_response.choices[0].message.content or ""
-            reasoning_trace = getattr(completion_response.choices[0].message, "reasoning_content", None)
+            # 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
 
-            if reasoning_trace:
-                ## There are generally some extra newlines with how these get parsed.
-                response = response.strip()
-                reasoning_trace = reasoning_trace.strip()
+                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))
             curr_num_correction_steps += 1
 
             try:
@@ -223,21 +282,24 @@ class ModelFacade:
                     raise GenerationValidationFailureError(
                         "Unsuccessful generation attempt. No retries were attempted."
                     ) from exc
+
                 if curr_num_correction_steps <= max_correction_steps:
-                    ## Add turns to loop-back errors for correction
-                    messages += [
-                        str_to_message(content=response, role="assistant"),
-                        str_to_message(content=str(get_exception_primary_cause(exc)), role="user"),
-                    ]
+                    # Add user message with error for correction
+                    messages.append(ChatMessage.as_user(content=str(get_exception_primary_cause(exc))))
+
                 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(
-                        f"Unsuccessful generation attempt despite {max_generation_attempts} attempts."
+                        f"Unsuccessful generation despite {max_correction_steps} correction steps "
+                        f"and {max_conversation_restarts} conversation restarts."
                     ) from exc
-        return output_obj, reasoning_trace
+
+        return output_obj, messages
 
     def _get_litellm_deployment(self, model_config: ModelConfig) -> litellm.DeploymentTypedDict:
         provider = self._model_provider_registry.get_provider(model_config.provider)

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

@@ -3,7 +3,84 @@
 
 from __future__ import annotations
 
-from typing import Any
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+@dataclass
+class ChatMessage:
+    """A chat message in an LLM conversation.
+
+    This dataclass represents messages exchanged in a conversation with an LLM,
+    supporting various message types including user prompts, assistant responses,
+    system instructions, and tool interactions.
+
+    Attributes:
+        role: The role of the message sender. One of 'user', 'assistant', 'system', or 'tool'.
+        content: The message content. Can be a string or a list of content blocks
+            for multimodal messages (e.g., text + images).
+        reasoning_content: Optional reasoning/thinking content from the assistant,
+            typically from extended thinking or chain-of-thought models.
+        tool_calls: Optional list of tool calls requested by the assistant.
+            Each tool call contains 'id', 'type', and 'function' keys.
+        tool_call_id: Optional ID linking a tool response to its corresponding
+            tool call. Required for messages with role='tool'.
+    """
+
+    role: Literal["user", "assistant", "system", "tool"]
+    content: str | list[dict[str, Any]] = ""
+    reasoning_content: str | None = None
+    tool_calls: list[dict[str, Any]] = field(default_factory=list)
+    tool_call_id: str | None = None
+
+    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": _normalize_content_blocks(self.content)}
+        if self.reasoning_content:
+            result["reasoning_content"] = self.reasoning_content
+        if self.tool_calls:
+            result["tool_calls"] = self.tool_calls
+        if self.tool_call_id:
+            result["tool_call_id"] = self.tool_call_id
+        return result
+
+    @classmethod
+    def as_user(cls, content: str | list[dict[str, Any]]) -> ChatMessage:
+        """Create a user message."""
+        return cls(role="user", content=content)
+
+    @classmethod
+    def as_assistant(
+        cls,
+        content: str = "",
+        reasoning_content: str | None = None,
+        tool_calls: list[dict[str, Any]] | None = None,
+    ) -> ChatMessage:
+        """Create an assistant message."""
+        return cls(
+            role="assistant",
+            content=content,
+            reasoning_content=reasoning_content,
+            tool_calls=tool_calls or [],
+        )
+
+    @classmethod
+    def as_system(cls, content: str) -> ChatMessage:
+        """Create a system message."""
+        return cls(role="system", content=content)
+
+    @classmethod
+    def as_tool(cls, content: str, tool_call_id: str) -> ChatMessage:
+        """Create a tool response message."""
+        return cls(role="tool", content=content, tool_call_id=tool_call_id)
 
 
 def prompt_to_messages(
@@ -11,28 +88,41 @@ def prompt_to_messages(
     user_prompt: str,
     system_prompt: str | None = None,
     multi_modal_context: list[dict[str, Any]] | None = None,
-) -> list[dict[str, str | list[dict]]]:
-    """Convert a user and system prompt into Messages format.
+) -> list[ChatMessage]:
+    """Convert a user and system prompt into ChatMessage list.
 
     Args:
         user_prompt (str): A user prompt.
         system_prompt (str, optional): An optional system prompt.
     """
-    user_content = user_prompt
-    if multi_modal_context and len(multi_modal_context) > 0:
-        user_content = []
-        for context in multi_modal_context:
-            user_content.append(context)
-        user_content.append({"type": "text", "text": user_prompt})
-    return (
-        [
-            str_to_message(content=system_prompt, role="system"),
-            str_to_message(content=user_content, role="user"),
-        ]
-        if system_prompt
-        else [str_to_message(content=user_content, role="user")]
-    )
-
-
-def str_to_message(content: str | list[dict], role: str = "user") -> dict[str, str | list[dict]]:
-    return {"content": content, "role": role}
+    user_content: str | list[dict[str, Any]] = user_prompt
+    if multi_modal_context:
+        user_content = [*multi_modal_context, {"type": "text", "text": user_prompt}]
+
+    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(),
     )