ag2 0.9.1a1__py3-none-any.whl → 0.9.1.post0__py3-none-any.whl

autogen/oai/mistral.py

@@ -0,0 +1,303 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
+# SPDX-License-Identifier: MIT
+"""Create an OpenAI-compatible client using Mistral.AI's API.
+
+Example:
+    ```python
+    llm_config = {
+        "config_list": [
+            {"api_type": "mistral", "model": "open-mixtral-8x22b", "api_key": os.environ.get("MISTRAL_API_KEY")}
+        ]
+    }
+
+    agent = autogen.AssistantAgent("my_agent", llm_config=llm_config)
+    ```
+
+Install Mistral.AI python library using: pip install --upgrade mistralai
+
+Resources:
+- https://docs.mistral.ai/getting-started/quickstart/
+
+NOTE: Requires mistralai package version >= 1.0.1
+"""
+
+import json
+import os
+import time
+import warnings
+from typing import Any, Literal, Optional, Union
+
+from pydantic import Field
+
+from ..import_utils import optional_import_block, require_optional_import
+from ..llm_config import LLMConfigEntry, register_llm_config
+from .client_utils import should_hide_tools, validate_parameter
+from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
+
+with optional_import_block():
+    # Mistral libraries
+    # pip install mistralai
+    from mistralai import (
+        AssistantMessage,
+        Function,
+        FunctionCall,
+        Mistral,
+        SystemMessage,
+        ToolCall,
+        ToolMessage,
+        UserMessage,
+    )
+
+
+@register_llm_config
+class MistralLLMConfigEntry(LLMConfigEntry):
+    api_type: Literal["mistral"] = "mistral"
+    temperature: float = Field(default=0.7)
+    top_p: Optional[float] = None
+    max_tokens: Optional[int] = Field(default=None, ge=0)
+    safe_prompt: bool = False
+    random_seed: Optional[int] = None
+    stream: bool = False
+    hide_tools: Literal["if_all_run", "if_any_run", "never"] = "never"
+    tool_choice: Optional[Literal["none", "auto", "any"]] = None
+
+    def create_client(self):
+        raise NotImplementedError("MistralLLMConfigEntry.create_client is not implemented.")
+
+
+@require_optional_import("mistralai", "mistral")
+class MistralAIClient:
+    """Client for Mistral.AI's API."""
+
+    def __init__(self, **kwargs):
+        """Requires api_key or environment variable to be set
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to the Mistral client.
+        """
+        # Ensure we have the api_key upon instantiation
+        self.api_key = kwargs.get("api_key")
+        if not self.api_key:
+            self.api_key = os.getenv("MISTRAL_API_KEY", None)
+
+        assert self.api_key, (
+            "Please specify the 'api_key' in your config list entry for Mistral or set the MISTRAL_API_KEY env variable."
+        )
+
+        if "response_format" in kwargs and kwargs["response_format"] is not None:
+            warnings.warn("response_format is not supported for Mistral.AI, it will be ignored.", UserWarning)
+
+        self._client = Mistral(api_key=self.api_key)
+
+    def message_retrieval(self, response: ChatCompletion) -> Union[list[str], list[ChatCompletionMessage]]:
+        """Retrieve the messages from the response."""
+        return [choice.message for choice in response.choices]
+
+    def cost(self, response) -> float:
+        return response.cost
+
+    @require_optional_import("mistralai", "mistral")
+    def parse_params(self, params: dict[str, Any]) -> dict[str, Any]:
+        """Loads the parameters for Mistral.AI API from the passed in parameters and returns a validated set. Checks types, ranges, and sets defaults"""
+        mistral_params = {}
+
+        # 1. Validate models
+        mistral_params["model"] = params.get("model")
+        assert mistral_params["model"], (
+            "Please specify the 'model' in your config list entry to nominate the Mistral.ai model to use."
+        )
+
+        # 2. Validate allowed Mistral.AI parameters
+        mistral_params["temperature"] = validate_parameter(params, "temperature", (int, float), True, 0.7, None, None)
+        mistral_params["top_p"] = validate_parameter(params, "top_p", (int, float), True, None, None, None)
+        mistral_params["max_tokens"] = validate_parameter(params, "max_tokens", int, True, None, (0, None), None)
+        mistral_params["safe_prompt"] = validate_parameter(
+            params, "safe_prompt", bool, False, False, None, [True, False]
+        )
+        mistral_params["random_seed"] = validate_parameter(params, "random_seed", int, True, None, False, None)
+        mistral_params["tool_choice"] = validate_parameter(
+            params, "tool_choice", str, False, None, None, ["none", "auto", "any"]
+        )
+
+        # TODO
+        if params.get("stream", False):
+            warnings.warn(
+                "Streaming is not currently supported, streaming will be disabled.",
+                UserWarning,
+            )
+
+        # 3. Convert messages to Mistral format
+        mistral_messages = []
+        tool_call_ids = {}  # tool call ids to function name mapping
+        for message in params["messages"]:
+            if message["role"] == "assistant" and "tool_calls" in message and message["tool_calls"] is not None:
+                # Convert OAI ToolCall to Mistral ToolCall
+                mistral_messages_tools = []
+                for toolcall in message["tool_calls"]:
+                    mistral_messages_tools.append(
+                        ToolCall(
+                            id=toolcall["id"],
+                            function=FunctionCall(
+                                name=toolcall["function"]["name"],
+                                arguments=json.loads(toolcall["function"]["arguments"]),
+                            ),
+                        )
+                    )
+
+                mistral_messages.append(AssistantMessage(content="", tool_calls=mistral_messages_tools))
+
+                # Map tool call id to the function name
+                for tool_call in message["tool_calls"]:
+                    tool_call_ids[tool_call["id"]] = tool_call["function"]["name"]
+
+            elif message["role"] == "system":
+                if len(mistral_messages) > 0 and mistral_messages[-1].role == "assistant":
+                    # System messages can't appear after an Assistant message, so use a UserMessage
+                    mistral_messages.append(UserMessage(content=message["content"]))
+                else:
+                    mistral_messages.append(SystemMessage(content=message["content"]))
+            elif message["role"] == "assistant":
+                mistral_messages.append(AssistantMessage(content=message["content"]))
+            elif message["role"] == "user":
+                mistral_messages.append(UserMessage(content=message["content"]))
+
+            elif message["role"] == "tool":
+                # Indicates the result of a tool call, the name is the function name called
+                mistral_messages.append(
+                    ToolMessage(
+                        name=tool_call_ids[message["tool_call_id"]],
+                        content=message["content"],
+                        tool_call_id=message["tool_call_id"],
+                    )
+                )
+            else:
+                warnings.warn(f"Unknown message role {message['role']}", UserWarning)
+
+        # 4. Last message needs to be user or tool, if not, add a "please continue" message
+        if not isinstance(mistral_messages[-1], UserMessage) and not isinstance(mistral_messages[-1], ToolMessage):
+            mistral_messages.append(UserMessage(content="Please continue."))
+
+        mistral_params["messages"] = mistral_messages
+
+        # 5. Add tools to the call if we have them and aren't hiding them
+        if "tools" in params:
+            hide_tools = validate_parameter(
+                params, "hide_tools", str, False, "never", None, ["if_all_run", "if_any_run", "never"]
+            )
+            if not should_hide_tools(params["messages"], params["tools"], hide_tools):
+                mistral_params["tools"] = tool_def_to_mistral(params["tools"])
+
+        return mistral_params
+
+    @require_optional_import("mistralai", "mistral")
+    def create(self, params: dict[str, Any]) -> ChatCompletion:
+        # 1. Parse parameters to Mistral.AI API's parameters
+        mistral_params = self.parse_params(params)
+
+        # 2. Call Mistral.AI API
+        mistral_response = self._client.chat.complete(**mistral_params)
+        # TODO: Handle streaming
+
+        # 3. Convert Mistral response to OAI compatible format
+        if mistral_response.choices[0].finish_reason == "tool_calls":
+            mistral_finish = "tool_calls"
+            tool_calls = []
+            for tool_call in mistral_response.choices[0].message.tool_calls:
+                tool_calls.append(
+                    ChatCompletionMessageToolCall(
+                        id=tool_call.id,
+                        function={"name": tool_call.function.name, "arguments": tool_call.function.arguments},
+                        type="function",
+                    )
+                )
+        else:
+            mistral_finish = "stop"
+            tool_calls = None
+
+        message = ChatCompletionMessage(
+            role="assistant",
+            content=mistral_response.choices[0].message.content,
+            function_call=None,
+            tool_calls=tool_calls,
+        )
+        choices = [Choice(finish_reason=mistral_finish, index=0, message=message)]
+
+        response_oai = ChatCompletion(
+            id=mistral_response.id,
+            model=mistral_response.model,
+            created=int(time.time()),
+            object="chat.completion",
+            choices=choices,
+            usage=CompletionUsage(
+                prompt_tokens=mistral_response.usage.prompt_tokens,
+                completion_tokens=mistral_response.usage.completion_tokens,
+                total_tokens=mistral_response.usage.prompt_tokens + mistral_response.usage.completion_tokens,
+            ),
+            cost=calculate_mistral_cost(
+                mistral_response.usage.prompt_tokens, mistral_response.usage.completion_tokens, mistral_response.model
+            ),
+        )
+
+        return response_oai
+
+    @staticmethod
+    def get_usage(response: ChatCompletion) -> dict:
+        return {
+            "prompt_tokens": response.usage.prompt_tokens if response.usage is not None else 0,
+            "completion_tokens": response.usage.completion_tokens if response.usage is not None else 0,
+            "total_tokens": (
+                response.usage.prompt_tokens + response.usage.completion_tokens if response.usage is not None else 0
+            ),
+            "cost": response.cost if hasattr(response, "cost") else 0,
+            "model": response.model,
+        }
+
+
+@require_optional_import("mistralai", "mistral")
+def tool_def_to_mistral(tool_definitions: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Converts AG2 tool definition to a mistral tool format"""
+    mistral_tools = []
+
+    for autogen_tool in tool_definitions:
+        mistral_tool = {
+            "type": "function",
+            "function": Function(
+                name=autogen_tool["function"]["name"],
+                description=autogen_tool["function"]["description"],
+                parameters=autogen_tool["function"]["parameters"],
+            ),
+        }
+
+        mistral_tools.append(mistral_tool)
+
+    return mistral_tools
+
+
+def calculate_mistral_cost(input_tokens: int, output_tokens: int, model_name: str) -> float:
+    """Calculate the cost of the mistral response."""
+    # Prices per 1 thousand tokens
+    # https://mistral.ai/technology/
+    model_cost_map = {
+        "open-mistral-7b": {"input": 0.00025, "output": 0.00025},
+        "open-mixtral-8x7b": {"input": 0.0007, "output": 0.0007},
+        "open-mixtral-8x22b": {"input": 0.002, "output": 0.006},
+        "mistral-small-latest": {"input": 0.001, "output": 0.003},
+        "mistral-medium-latest": {"input": 0.00275, "output": 0.0081},
+        "mistral-large-latest": {"input": 0.0003, "output": 0.0003},
+        "mistral-large-2407": {"input": 0.0003, "output": 0.0003},
+        "open-mistral-nemo-2407": {"input": 0.0003, "output": 0.0003},
+        "codestral-2405": {"input": 0.001, "output": 0.003},
+    }
+
+    # Ensure we have the model they are using and return the total cost
+    if model_name in model_cost_map:
+        costs = model_cost_map[model_name]
+
+        return (input_tokens * costs["input"] / 1000) + (output_tokens * costs["output"] / 1000)
+    else:
+        warnings.warn(f"Cost calculation is not implemented for model {model_name}, will return $0.", UserWarning)
+        return 0

autogen/oai/oai_models/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .chat_completion import ChatCompletionExtended as ChatCompletion
+from .chat_completion import Choice
+from .chat_completion_message import ChatCompletionMessage
+from .chat_completion_message_tool_call import ChatCompletionMessageToolCall
+from .completion_usage import CompletionUsage
+
+__all__ = ["ChatCompletion", "ChatCompletionMessage", "ChatCompletionMessageToolCall", "Choice", "CompletionUsage"]

autogen/oai/oai_models/_models.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Taken over from https://github.com/openai/openai-python/blob/main/src/openai/_models.py
+
+import pydantic
+import pydantic.generics
+from pydantic import ConfigDict
+from typing_extensions import ClassVar
+
+__all__ = ["BaseModel"]
+
+
+class BaseModel(pydantic.BaseModel):
+    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")

autogen/oai/oai_models/chat_completion.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Taken over from https://github.com/openai/openai-python/blob/3e69750d47df4f0759d4a28ddc68e4b38756d9ca/src/openai/types/chat/chat_completion.py
+
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from typing import Any, Callable, List, Optional
+
+from typing_extensions import Literal
+
+from ._models import BaseModel
+from .chat_completion_message import ChatCompletionMessage
+from .chat_completion_token_logprob import ChatCompletionTokenLogprob
+from .completion_usage import CompletionUsage
+
+__all__ = ["ChatCompletion", "Choice", "ChoiceLogprobs"]
+
+
+class ChoiceLogprobs(BaseModel):
+    content: Optional[List[ChatCompletionTokenLogprob]] = None
+    """A list of message content tokens with log probability information."""
+
+    refusal: Optional[List[ChatCompletionTokenLogprob]] = None
+    """A list of message refusal tokens with log probability information."""
+
+
+class Choice(BaseModel):
+    finish_reason: Literal["stop", "length", "tool_calls", "content_filter", "function_call"]
+    """The reason the model stopped generating tokens.
+
+    This will be `stop` if the model hit a natural stop point or a provided stop
+    sequence, `length` if the maximum number of tokens specified in the request was
+    reached, `content_filter` if content was omitted due to a flag from our content
+    filters, `tool_calls` if the model called a tool, or `function_call`
+    (deprecated) if the model called a function.
+    """
+
+    index: int
+    """The index of the choice in the list of choices."""
+
+    logprobs: Optional[ChoiceLogprobs] = None
+    """Log probability information for the choice."""
+
+    message: ChatCompletionMessage
+    """A chat completion message generated by the model."""
+
+
+class ChatCompletion(BaseModel):
+    id: str
+    """A unique identifier for the chat completion."""
+
+    choices: List[Choice]
+    """A list of chat completion choices.
+
+    Can be more than one if `n` is greater than 1.
+    """
+
+    created: int
+    """The Unix timestamp (in seconds) of when the chat completion was created."""
+
+    model: str
+    """The model used for the chat completion."""
+
+    object: Literal["chat.completion"]
+    """The object type, which is always `chat.completion`."""
+
+    service_tier: Optional[Literal["auto", "default", "flex"]] = None
+    """The service tier used for processing the request."""
+
+    system_fingerprint: Optional[str] = None
+    """This fingerprint represents the backend configuration that the model runs with.
+
+    Can be used in conjunction with the `seed` request parameter to understand when
+    backend changes have been made that might impact determinism.
+    """
+
+    usage: Optional[CompletionUsage] = None
+    """Usage statistics for the completion request."""
+
+
+class ChatCompletionExtended(ChatCompletion):
+    message_retrieval_function: Optional[Callable[[Any, "ChatCompletion"], list[ChatCompletionMessage]]] = None
+    config_id: Optional[str] = None
+    pass_filter: Optional[Callable[..., bool]] = None
+    cost: Optional[float] = None

autogen/oai/oai_models/chat_completion_audio.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Taken over from https://github.com/openai/openai-python/blob/3e69750d47df4f0759d4a28ddc68e4b38756d9ca/src/openai/types/chat/chat_completion_audio.py
+
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+
+from ._models import BaseModel
+
+__all__ = ["ChatCompletionAudio"]
+
+
+class ChatCompletionAudio(BaseModel):
+    id: str
+    """Unique identifier for this audio response."""
+
+    data: str
+    """
+    Base64 encoded audio bytes generated by the model, in the format specified in
+    the request.
+    """
+
+    expires_at: int
+    """
+    The Unix timestamp (in seconds) for when this audio response will no longer be
+    accessible on the server for use in multi-turn conversations.
+    """
+
+    transcript: str
+    """Transcript of the audio generated by the model."""

autogen/oai/oai_models/chat_completion_message.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Taken over from https://github.com/openai/openai-python/blob/16a10604fbd0d82c1382b84b417a1d6a2d33a7f1/src/openai/types/chat/chat_completion_message.py
+
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from typing import List, Optional
+
+from typing_extensions import Literal
+
+from ._models import BaseModel
+from .chat_completion_audio import ChatCompletionAudio
+from .chat_completion_message_tool_call import ChatCompletionMessageToolCall
+
+__all__ = ["Annotation", "AnnotationURLCitation", "ChatCompletionMessage", "FunctionCall"]
+
+
+class AnnotationURLCitation(BaseModel):
+    end_index: int
+    """The index of the last character of the URL citation in the message."""
+
+    start_index: int
+    """The index of the first character of the URL citation in the message."""
+
+    title: str
+    """The title of the web resource."""
+
+    url: str
+    """The URL of the web resource."""
+
+
+class Annotation(BaseModel):
+    type: Literal["url_citation"]
+    """The type of the URL citation. Always `url_citation`."""
+
+    url_citation: AnnotationURLCitation
+    """A URL citation when using web search."""
+
+
+class FunctionCall(BaseModel):
+    arguments: str
+    """
+    The arguments to call the function with, as generated by the model in JSON
+    format. Note that the model does not always generate valid JSON, and may
+    hallucinate parameters not defined by your function schema. Validate the
+    arguments in your code before calling your function.
+    """
+
+    name: str
+    """The name of the function to call."""
+
+
+class ChatCompletionMessage(BaseModel):
+    content: Optional[str] = None
+    """The contents of the message."""
+
+    refusal: Optional[str] = None
+    """The refusal message generated by the model."""
+
+    role: Literal["assistant"]
+    """The role of the author of this message."""
+
+    annotations: Optional[List[Annotation]] = None
+    """
+    Annotations for the message, when applicable, as when using the
+    [web search tool](https://platform.openai.com/docs/guides/tools-web-search?api-mode=chat).
+    """
+
+    audio: Optional[ChatCompletionAudio] = None
+    """
+    If the audio output modality is requested, this object contains data about the
+    audio response from the model.
+    [Learn more](https://platform.openai.com/docs/guides/audio).
+    """
+
+    function_call: Optional[FunctionCall] = None
+    """Deprecated and replaced by `tool_calls`.
+
+    The name and arguments of a function that should be called, as generated by the
+    model.
+    """
+
+    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
+    """The tool calls generated by the model, such as function calls."""

autogen/oai/oai_models/chat_completion_message_tool_call.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Taken over from https://github.com/openai/openai-python/blob/3e69750d47df4f0759d4a28ddc68e4b38756d9ca/src/openai/types/chat/chat_completion_message_tool_call.py
+
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from typing_extensions import Literal
+
+from ._models import BaseModel
+
+__all__ = ["ChatCompletionMessageToolCall", "Function"]
+
+
+class Function(BaseModel):
+    arguments: str
+    """
+    The arguments to call the function with, as generated by the model in JSON
+    format. Note that the model does not always generate valid JSON, and may
+    hallucinate parameters not defined by your function schema. Validate the
+    arguments in your code before calling your function.
+    """
+
+    name: str
+    """The name of the function to call."""
+
+
+class ChatCompletionMessageToolCall(BaseModel):
+    id: str
+    """The ID of the tool call."""
+
+    function: Function
+    """The function that the model called."""
+
+    type: Literal["function"]
+    """The type of the tool. Currently, only `function` is supported."""

autogen/oai/oai_models/chat_completion_token_logprob.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Taken over from https://github.com/openai/openai-python/blob/3e69750d47df4f0759d4a28ddc68e4b38756d9ca/src/openai/types/chat/chat_completion_token_logprob.py
+
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from typing import List, Optional
+
+from ._models import BaseModel
+
+__all__ = ["ChatCompletionTokenLogprob", "TopLogprob"]
+
+
+class TopLogprob(BaseModel):
+    token: str
+    """The token."""
+
+    bytes: Optional[List[int]] = None
+    """A list of integers representing the UTF-8 bytes representation of the token.
+
+    Useful in instances where characters are represented by multiple tokens and
+    their byte representations must be combined to generate the correct text
+    representation. Can be `null` if there is no bytes representation for the token.
+    """
+
+    logprob: float
+    """The log probability of this token, if it is within the top 20 most likely
+    tokens.
+
+    Otherwise, the value `-9999.0` is used to signify that the token is very
+    unlikely.
+    """
+
+
+class ChatCompletionTokenLogprob(BaseModel):
+    token: str
+    """The token."""
+
+    bytes: Optional[List[int]] = None
+    """A list of integers representing the UTF-8 bytes representation of the token.
+
+    Useful in instances where characters are represented by multiple tokens and
+    their byte representations must be combined to generate the correct text
+    representation. Can be `null` if there is no bytes representation for the token.
+    """
+
+    logprob: float
+    """The log probability of this token, if it is within the top 20 most likely
+    tokens.
+
+    Otherwise, the value `-9999.0` is used to signify that the token is very
+    unlikely.
+    """
+
+    top_logprobs: List[TopLogprob]
+    """List of the most likely tokens and their log probability, at this token
+    position.
+
+    In rare cases, there may be fewer than the number of requested `top_logprobs`
+    returned.
+    """