mistralai 1.7.1__py3-none-any.whl → 1.8.0__py3-none-any.whl

mistralai/extra/run/context.py

@@ -0,0 +1,295 @@
+import asyncio
+import inspect
+import typing
+from contextlib import AsyncExitStack
+from functools import wraps
+from collections.abc import Callable
+
+from dataclasses import dataclass, field
+from typing import Union, Optional
+
+import pydantic
+
+from mistralai.extra import (
+    response_format_from_pydantic_model,
+)
+from mistralai.extra.exceptions import RunException
+from mistralai.extra.mcp.base import MCPClientProtocol
+from mistralai.extra.run.result import RunResult
+from mistralai.types.basemodel import OptionalNullable, BaseModel, UNSET
+from mistralai.models import (
+    ResponseFormat,
+    FunctionCallEntry,
+    Tools,
+    ToolsTypedDict,
+    CompletionArgs,
+    CompletionArgsTypedDict,
+    FunctionResultEntry,
+    ConversationInputs,
+    ConversationInputsTypedDict,
+    FunctionTool,
+    MessageInputEntry,
+    InputEntries,
+    ResponseFormatTypedDict,
+)
+
+from logging import getLogger
+
+from mistralai.extra.run.tools import (
+    create_function_result,
+    RunFunction,
+    create_tool_call,
+    RunTool,
+    RunMCPTool,
+    RunCoroutine,
+)
+
+if typing.TYPE_CHECKING:
+    from mistralai import Beta, OptionalNullable
+
+logger = getLogger(__name__)
+
+
+class AgentRequestKwargs(typing.TypedDict):
+    agent_id: str
+
+
+class ModelRequestKwargs(typing.TypedDict):
+    model: str
+    instructions: OptionalNullable[str]
+    tools: OptionalNullable[Union[list[Tools], list[ToolsTypedDict]]]
+    completion_args: OptionalNullable[Union[CompletionArgs, CompletionArgsTypedDict]]
+
+
+@dataclass
+class RunContext:
+    """A context for running a conversation with an agent or a model.
+
+    The context can be used to execute function calls, connect to MCP server, and keep track of information about
+    the run.
+
+    Args:
+        conversation_id (Options[str]): The unique identifier for the conversation. This is
+          passed if the user wants to continue an existing conversation.
+        model (Options[str]): The model name to be used for the conversation. Can't be used along with 'agent_id'.
+        agent_id (Options[str]): The agent id to be used for the conversation. Can't be used along with 'model'.
+        output_format (Optional[type[BaseModel]]): The output format expected from the conversation. It represents
+          the `response_format` which is part of the `CompletionArgs`.
+        request_count (int): The number of requests made in the current `RunContext`.
+        continue_on_fn_error (bool): Flag to determine if the conversation should continue when function execution
+          resulted in an error.
+    """
+
+    _exit_stack: AsyncExitStack = field(init=False)
+    _callable_tools: dict[str, RunTool] = field(init=False, default_factory=dict)
+    _mcp_clients: list[MCPClientProtocol] = field(init=False, default_factory=list)
+
+    conversation_id: Optional[str] = field(default=None)
+    model: Optional[str] = field(default=None)
+    agent_id: Optional[str] = field(default=None)
+    output_format: Optional[type[BaseModel]] = field(default=None)
+    request_count: int = field(default=0)
+    continue_on_fn_error: bool = field(default=False)
+
+    def __post_init__(self):
+        if self.model and self.agent_id:
+            raise RunException("Only one for model or agent_id should be set")
+        self._exit_stack = AsyncExitStack()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._exit_stack.aclose()
+        for mcp_client in self._mcp_clients:
+            await mcp_client.aclose()
+
+    def register_func(self, func: Callable):
+        """Add a function to the context."""
+        if not inspect.isfunction(func):
+            raise RunException(
+                "Only object of type function can be registered at the moment."
+            )
+
+        if inspect.iscoroutinefunction(func):
+            self._callable_tools[func.__name__] = RunCoroutine(
+                name=func.__name__,
+                awaitable=func,
+                tool=create_tool_call(func),
+            )
+        else:
+            self._callable_tools[func.__name__] = RunFunction(
+                name=func.__name__,
+                callable=func,
+                tool=create_tool_call(func),
+            )
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            logger.info(f"Executing {func.__name__}")
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    async def register_mcp_clients(self, mcp_clients: list[MCPClientProtocol]) -> None:
+        """Registering multiple MCP clients at the same time in the same asyncio.Task."""
+        for mcp_client in mcp_clients:
+            await self.register_mcp_client(mcp_client)
+
+    async def register_mcp_client(self, mcp_client: MCPClientProtocol) -> None:
+        """Add a MCP client to the context."""
+        await mcp_client.initialize(exit_stack=self._exit_stack)
+        tools = await mcp_client.get_tools()
+        for tool in tools:
+            logger.info(
+                f"Adding tool {tool.function.name} from {mcp_client._name or 'mcp client'}"
+            )
+            self._callable_tools[tool.function.name] = RunMCPTool(
+                name=tool.function.name,
+                tool=tool,
+                mcp_client=mcp_client,
+            )
+        self._mcp_clients.append(mcp_client)
+
+    async def execute_function_calls(
+        self, function_calls: list[FunctionCallEntry]
+    ) -> list[FunctionResultEntry]:
+        """Execute function calls and create function results from them."""
+        if not all(
+            function_call.name in self._callable_tools
+            for function_call in function_calls
+        ):
+            logger.warning("Can't execute all functions, stopping run here")
+            return []
+        function_result_tasks = []
+        for function_call in function_calls:
+            function_result_tasks.append(
+                asyncio.create_task(
+                    create_function_result(
+                        function_call=function_call,
+                        run_tool=self._callable_tools[function_call.name],
+                        continue_on_fn_error=self.continue_on_fn_error,
+                    )
+                )
+            )
+        await asyncio.gather(*function_result_tasks)
+        return [task.result() for task in function_result_tasks]
+
+    def get_tools(self) -> list[FunctionTool]:
+        """Get the tools that are part of the context."""
+        callable_tools = [
+            run_functions.tool for run_functions in self._callable_tools.values()
+        ]
+        return callable_tools
+
+    async def prepare_agent_request(self, beta_client: "Beta") -> AgentRequestKwargs:
+        """Prepare an agent request with the functions added to the context.
+
+        Update the agent definition before making the request.
+        """
+        if self.agent_id is None:
+            raise RunException(
+                "Can't prepare an agent request, if no agent_id is provided"
+            )
+        agent = await beta_client.agents.get_async(agent_id=self.agent_id)
+        agent_tools = agent.tools or []
+        updated_tools = []
+        for i in range(len(agent_tools)):
+            tool = agent_tools[i]
+            if tool.type != "function":
+                updated_tools.append(tool)
+            elif tool.function.name in self._callable_tools:
+                # function already exists in the agent, don't add it again
+                continue
+            else:
+                updated_tools.append(tool)
+        updated_tools += self.get_tools()
+        completion_args = (
+            CompletionArgs(response_format=self.response_format)
+            if self.output_format
+            else None
+        )
+        beta_client.agents.update(
+            agent_id=self.agent_id, tools=updated_tools, completion_args=completion_args
+        )
+        return AgentRequestKwargs(agent_id=self.agent_id)
+
+    async def prepare_model_request(
+        self,
+        tools: OptionalNullable[Union[list[Tools], list[ToolsTypedDict]]] = UNSET,
+        completion_args: OptionalNullable[
+            Union[CompletionArgs, CompletionArgsTypedDict]
+        ] = UNSET,
+        instructions: OptionalNullable[str] = None,
+    ) -> ModelRequestKwargs:
+        if self.model is None:
+            raise RunException("Can't prepare a model request, if no model is provided")
+        if not completion_args and self.output_format:
+            completion_args = CompletionArgs(response_format=self.response_format)
+        elif isinstance(completion_args, CompletionArgs) and self.output_format:
+            completion_args.response_format = self.response_format
+        elif isinstance(completion_args, dict) and self.output_format:
+            completion_args["response_format"] = typing.cast(
+                ResponseFormatTypedDict, self.response_format.model_dump()
+            )
+        request_tools = []
+        if isinstance(tools, list):
+            for tool in tools:
+                request_tools.append(typing.cast(Tools, tool))
+        for tool in self.get_tools():
+            request_tools.append(tool)
+        return ModelRequestKwargs(
+            model=self.model,
+            tools=request_tools,
+            instructions=instructions,
+            completion_args=completion_args,
+        )
+
+    @property
+    def response_format(self) -> ResponseFormat:
+        if not self.output_format:
+            raise RunException("No response format exist for the current RunContext.")
+        return response_format_from_pydantic_model(self.output_format)
+
+
+async def _validate_run(
+    *,
+    beta_client: "Beta",
+    run_ctx: RunContext,
+    inputs: Union[ConversationInputs, ConversationInputsTypedDict],
+    instructions: OptionalNullable[str] = UNSET,
+    tools: OptionalNullable[Union[list[Tools], list[ToolsTypedDict]]] = UNSET,
+    completion_args: OptionalNullable[
+        Union[CompletionArgs, CompletionArgsTypedDict]
+    ] = UNSET,
+) -> tuple[
+    Union[AgentRequestKwargs, ModelRequestKwargs], RunResult, list[InputEntries]
+]:
+    input_entries: list[InputEntries] = []
+    if isinstance(inputs, str):
+        input_entries.append(MessageInputEntry(role="user", content=inputs))
+    else:
+        for input in inputs:
+            if isinstance(input, dict):
+                input_entries.append(
+                    pydantic.TypeAdapter(InputEntries).validate_python(input)
+                )
+    run_result = RunResult(
+        input_entries=input_entries,
+        output_model=run_ctx.output_format,
+        conversation_id=run_ctx.conversation_id,
+    )
+    req: Union[AgentRequestKwargs, ModelRequestKwargs]
+    if run_ctx.agent_id:
+        if tools or completion_args:
+            raise RunException("Can't set tools or completion_args when using an agent")
+        req = await run_ctx.prepare_agent_request(beta_client=beta_client)
+    elif run_ctx.model:
+        req = await run_ctx.prepare_model_request(
+            instructions=instructions,
+            tools=tools,
+            completion_args=completion_args,
+        )
+    else:
+        raise RunException("Either agent_id or model must be set in the run context")
+    return req, run_result, input_entries

mistralai/extra/run/result.py

@@ -0,0 +1,212 @@
+import datetime
+import json
+import typing
+from typing import Union, Annotated, Optional, Literal
+from dataclasses import dataclass, field
+from pydantic import Discriminator, Tag, BaseModel
+
+from mistralai.extra.utils.response_format import pydantic_model_from_json
+from mistralai.models import (
+    FunctionResultEntry,
+    FunctionCallEntry,
+    MessageOutputEntry,
+    AgentHandoffEntry,
+    ToolExecutionEntry,
+    MessageInputEntry,
+    AgentHandoffDoneEvent,
+    AgentHandoffStartedEvent,
+    ResponseDoneEvent,
+    ResponseErrorEvent,
+    ResponseStartedEvent,
+    FunctionCallEvent,
+    MessageOutputEvent,
+    ToolExecutionDoneEvent,
+    ToolExecutionStartedEvent,
+    ConversationEventsData,
+    MessageOutputEventContent,
+    MessageOutputEntryContent,
+    TextChunk,
+    MessageOutputContentChunks,
+    SSETypes,
+    InputEntries,
+    ToolFileChunk,
+    ToolReferenceChunk,
+    FunctionCallEntryArguments,
+)
+from mistralai.utils import get_discriminator
+
+RunOutputEntries = typing.Union[
+    MessageOutputEntry,
+    FunctionCallEntry,
+    FunctionResultEntry,
+    AgentHandoffEntry,
+    ToolExecutionEntry,
+]
+
+RunEntries = typing.Union[RunOutputEntries, MessageInputEntry]
+
+
+def as_text(entry: RunOutputEntries) -> str:
+    """Keep only the messages and turn content into textual representation."""
+    text = ""
+    if isinstance(entry, MessageOutputEntry):
+        if isinstance(entry.content, str):
+            text += entry.content
+        else:
+            for chunk in entry.content:
+                if isinstance(chunk, TextChunk):
+                    text += chunk.text
+                elif isinstance(chunk, ToolFileChunk):
+                    text += f"<File id={chunk.file_id} name={chunk.file_name}>"
+                elif isinstance(chunk, ToolReferenceChunk):
+                    text += f"<Reference title={chunk.title}>"
+    return text
+
+
+def reconstitute_message_content(
+    chunks: list[MessageOutputEventContent],
+) -> MessageOutputEntryContent:
+    """Given a list of MessageOutputEventContent, recreate a normalised MessageOutputEntryContent."""
+    if all(isinstance(chunk, str) for chunk in chunks):
+        return "".join(typing.cast(list[str], chunks))
+    content: list[MessageOutputContentChunks] = []
+    for chunk in chunks:
+        if isinstance(chunk, str):
+            chunk = TextChunk(text=chunk)
+        if isinstance(chunk, TextChunk):
+            if len(content) and isinstance(content[-1], TextChunk):
+                content[-1].text += chunk.text
+            else:
+                content.append(chunk)
+        else:
+            content.append(chunk)
+    return content
+
+
+def reconstitute_function_call_args(chunks: list[str]) -> FunctionCallEntryArguments:
+    """Recreates function call arguments from stream"""
+    return typing.cast(FunctionCallEntryArguments, "".join(chunks))
+
+
+def reconstitue_entries(
+    received_event_tracker: dict[int, list[ConversationEventsData]],
+) -> list[RunOutputEntries]:
+    """Given a list of events, recreate the corresponding entries."""
+    run_entries: list[RunOutputEntries] = []
+    for idx, events in sorted(received_event_tracker.items(), key=lambda x: x[0]):
+        first_event = events[0]
+        if isinstance(first_event, MessageOutputEvent):
+            message_events = typing.cast(list[MessageOutputEvent], events)
+            run_entries.append(
+                MessageOutputEntry(
+                    content=reconstitute_message_content(
+                        chunks=[
+                            message_event.content for message_event in message_events
+                        ]
+                    ),
+                    created_at=first_event.created_at,
+                    id=first_event.id,
+                    agent_id=first_event.agent_id,
+                    model=first_event.model,
+                    role=first_event.role,
+                )
+            )
+        elif isinstance(first_event, FunctionCallEvent):
+            function_call_events = typing.cast(list[FunctionCallEvent], events)
+            run_entries.append(
+                FunctionCallEntry(
+                    name=first_event.name,
+                    arguments=reconstitute_function_call_args(
+                        chunks=[
+                            function_call_event.arguments
+                            for function_call_event in function_call_events
+                        ]
+                    ),
+                    created_at=first_event.created_at,
+                    id=first_event.id,
+                    tool_call_id=first_event.tool_call_id,
+                )
+            )
+    return run_entries
+
+
+@dataclass
+class RunFiles:
+    id: str
+    name: str
+    content: bytes
+
+
+@dataclass
+class RunResult:
+    input_entries: list[InputEntries]
+    conversation_id: Optional[str] = field(default=None)
+    output_entries: list[RunOutputEntries] = field(default_factory=list)
+    files: dict[str, RunFiles] = field(default_factory=dict)
+    output_model: Optional[type[BaseModel]] = field(default=None)
+
+    def get_file(self, file_id: str) -> Optional[RunFiles]:
+        return self.files.get(file_id)
+
+    @property
+    def entries(self) -> list[RunEntries]:
+        return [*self.input_entries, *self.output_entries]
+
+    @property
+    def output_as_text(self) -> str:
+        if not self.output_entries:
+            raise ValueError("No output entries were started.")
+        return "\n".join(
+            as_text(entry)
+            for entry in self.output_entries
+            if entry.type == "message.output"
+        )
+
+    @property
+    def output_as_model(self) -> BaseModel:
+        if self.output_model is None:
+            raise ValueError("No output format was not set.")
+        return pydantic_model_from_json(
+            json.loads(self.output_as_text), self.output_model
+        )
+
+
+class FunctionResultEvent(BaseModel):
+    id: Optional[str] = None
+
+    type: Optional[Literal["function.result"]] = "function.result"
+
+    result: str
+
+    tool_call_id: str
+
+    created_at: Optional[datetime.datetime] = datetime.datetime.now(
+        tz=datetime.timezone.utc
+    )
+
+    output_index: Optional[int] = 0
+
+
+RunResultEventsType = typing.Union[SSETypes, Literal["function.result"]]
+
+RunResultEventsData = typing.Annotated[
+    Union[
+        Annotated[AgentHandoffDoneEvent, Tag("agent.handoff.done")],
+        Annotated[AgentHandoffStartedEvent, Tag("agent.handoff.started")],
+        Annotated[ResponseDoneEvent, Tag("conversation.response.done")],
+        Annotated[ResponseErrorEvent, Tag("conversation.response.error")],
+        Annotated[ResponseStartedEvent, Tag("conversation.response.started")],
+        Annotated[FunctionCallEvent, Tag("function.call.delta")],
+        Annotated[MessageOutputEvent, Tag("message.output.delta")],
+        Annotated[ToolExecutionDoneEvent, Tag("tool.execution.done")],
+        Annotated[ToolExecutionStartedEvent, Tag("tool.execution.started")],
+        Annotated[FunctionResultEvent, Tag("function.result")],
+    ],
+    Discriminator(lambda m: get_discriminator(m, "type", "type")),
+]
+
+
+class RunResultEvents(BaseModel):
+    event: RunResultEventsType
+
+    data: RunResultEventsData