client_code.py

class OpenAIChatCompletionClient(BaseOpenAIChatCompletionClient, Component[OpenAIClientConfigurationConfigModel]):
    """Chat completion client for OpenAI hosted models.

    To use this client, you must install the `openai` extra:

    .. code-block:: bash

        pip install "autogen-ext[openai]"

    You can also use this client for OpenAI-compatible ChatCompletion endpoints.
    **Using this client for non-OpenAI models is not tested or guaranteed.**

    For non-OpenAI models, please first take a look at our `community extensions <https://microsoft.github.io/autogen/dev/user-guide/extensions-user-guide/index.html>`_
    for additional model clients.

    Args:
        model (str): Which OpenAI model to use.
        api_key (optional, str): The API key to use. **Required if 'OPENAI_API_KEY' is not found in the environment variables.**
        organization (optional, str): The organization ID to use.
        base_url (optional, str): The base URL to use. **Required if the model is not hosted on OpenAI.**
        timeout: (optional, float): The timeout for the request in seconds.
        max_retries (optional, int): The maximum number of retries to attempt.
        model_info (optional, ModelInfo): The capabilities of the model. **Required if the model name is not a valid OpenAI model.**
        frequency_penalty (optional, float):
        logit_bias: (optional, dict[str, int]):
        max_tokens (optional, int):
        n (optional, int):
        presence_penalty (optional, float):
        response_format (optional, Dict[str, Any]): the format of the response. Possible options are:

            .. code-block:: text

                # Text response, this is the default.
                {"type": "text"}

            .. code-block:: text

                # JSON response, make sure to instruct the model to return JSON.
                {"type": "json_object"}

            .. code-block:: text

                # Structured output response, with a pre-defined JSON schema.
                {
                    "type": "json_schema",
                    "json_schema": {
                        "name": "name of the schema, must be an identifier.",
                        "description": "description for the model.",
                        # You can convert a Pydantic (v2) model to JSON schema
                        # using the `model_json_schema()` method.
                        "schema": "<the JSON schema itself>",
                        # Whether to enable strict schema adherence when
                        # generating the output. If set to true, the model will
                        # always follow the exact schema defined in the
                        # `schema` field. Only a subset of JSON Schema is
                        # supported when `strict` is `true`.
                        # To learn more, read
                        # https://platform.openai.com/docs/guides/structured-outputs.
                        "strict": False,  # or True
                    },
                }

            It is recommended to use the `json_output` parameter in
            :meth:`~autogen_ext.models.openai.BaseOpenAIChatCompletionClient.create` or
            :meth:`~autogen_ext.models.openai.BaseOpenAIChatCompletionClient.create_stream`
            methods instead of `response_format` for structured output.
            The `json_output` parameter is more flexible and allows you to
            specify a Pydantic model class directly.

        seed (optional, int):
        stop (optional, str | List[str]):
        temperature (optional, float):
        top_p (optional, float):
        parallel_tool_calls (optional, bool): Whether to allow parallel tool calls. When not set, defaults to server behavior.
        user (optional, str):
        default_headers (optional, dict[str, str]):  Custom headers; useful for authentication or other custom requirements.
        add_name_prefixes (optional, bool): Whether to prepend the `source` value
            to each :class:`~autogen_core.models.UserMessage` content. E.g.,
            "this is content" becomes "Reviewer said: this is content."
            This can be useful for models that do not support the `name` field in
            message. Defaults to False.
        include_name_in_message (optional, bool): Whether to include the `name` field
            in user message parameters sent to the OpenAI API. Defaults to True. Set to False
            for model providers that don't support the `name` field (e.g., Groq).
        stream_options (optional, dict): Additional options for streaming. Currently only `include_usage` is supported.

    Examples:

        The following code snippet shows how to use the client with an OpenAI model:

        .. code-block:: python

            from autogen_ext.models.openai import OpenAIChatCompletionClient
            from autogen_core.models import UserMessage

            openai_client = OpenAIChatCompletionClient(
                model="gpt-4o-2024-08-06",
                # api_key="sk-...", # Optional if you have an OPENAI_API_KEY environment variable set.
            )

            result = await openai_client.create([UserMessage(content="What is the capital of France?", source="user")])  # type: ignore
            print(result)

            # Close the client when done.
            # await openai_client.close()

        To use the client with a non-OpenAI model, you need to provide the base URL of the model and the model info.
        For example, to use Ollama, you can use the following code snippet:

        .. code-block:: python

            from autogen_ext.models.openai import OpenAIChatCompletionClient
            from autogen_core.models import ModelFamily

            custom_model_client = OpenAIChatCompletionClient(
                model="deepseek-r1:1.5b",
                base_url="http://localhost:11434/v1",
                api_key="placeholder",
                model_info={
                    "vision": False,
                    "function_calling": False,
                    "json_output": False,
                    "family": ModelFamily.R1,
                    "structured_output": True,
                },
            )

            # Close the client when done.
            # await custom_model_client.close()

        To use streaming mode, you can use the following code snippet:

        .. code-block:: python

            import asyncio
            from autogen_core.models import UserMessage
            from autogen_ext.models.openai import OpenAIChatCompletionClient


            async def main() -> None:
                # Similar for AzureOpenAIChatCompletionClient.
                model_client = OpenAIChatCompletionClient(model="gpt-4o")  # assuming OPENAI_API_KEY is set in the environment.

                messages = [UserMessage(content="Write a very short story about a dragon.", source="user")]

                # Create a stream.
                stream = model_client.create_stream(messages=messages)

                # Iterate over the stream and print the responses.
                print("Streamed responses:")
                async for response in stream:
                    if isinstance(response, str):
                        # A partial response is a string.
                        print(response, flush=True, end="")
                    else:
                        # The last response is a CreateResult object with the complete message.
                        print("\\n\\n------------\\n")
                        print("The complete response:", flush=True)
                        print(response.content, flush=True)

                # Close the client when done.
                await model_client.close()


            asyncio.run(main())

        To use structured output as well as function calling, you can use the following code snippet:

        .. code-block:: python

            import asyncio
            from typing import Literal

            from autogen_core.models import (
                AssistantMessage,
                FunctionExecutionResult,
                FunctionExecutionResultMessage,
                SystemMessage,
                UserMessage,
            )
            from autogen_core.tools import FunctionTool
            from autogen_ext.models.openai import OpenAIChatCompletionClient
            from pydantic import BaseModel


            # Define the structured output format.
            class AgentResponse(BaseModel):
                thoughts: str
                response: Literal["happy", "sad", "neutral"]


            # Define the function to be called as a tool.
            def sentiment_analysis(text: str) -> str:
                \"\"\"Given a text, return the sentiment.\"\"\"
                return "happy" if "happy" in text else "sad" if "sad" in text else "neutral"


            # Create a FunctionTool instance with `strict=True`,
            # which is required for structured output mode.
            tool = FunctionTool(sentiment_analysis, description="Sentiment Analysis", strict=True)


            async def main() -> None:
                # Create an OpenAIChatCompletionClient instance.
                model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")

                # Generate a response using the tool.
                response1 = await model_client.create(
                    messages=[
                        SystemMessage(content="Analyze input text sentiment using the tool provided."),
                        UserMessage(content="I am happy.", source="user"),
                    ],
                    tools=[tool],
                )
                print(response1.content)
                # Should be a list of tool calls.
                # [FunctionCall(name="sentiment_analysis", arguments={"text": "I am happy."}, ...)]

                assert isinstance(response1.content, list)
                response2 = await model_client.create(
                    messages=[
                        SystemMessage(content="Analyze input text sentiment using the tool provided."),
                        UserMessage(content="I am happy.", source="user"),
                        AssistantMessage(content=response1.content, source="assistant"),
                        FunctionExecutionResultMessage(
                            content=[FunctionExecutionResult(content="happy", call_id=response1.content[0].id, is_error=False, name="sentiment_analysis")]
                        ),
                    ],
                    # Use the structured output format.
                    json_output=AgentResponse,
                )
                print(response2.content)
                # Should be a structured output.
                # {"thoughts": "The user is happy.", "response": "happy"}

                # Close the client when done.
                await model_client.close()

            asyncio.run(main())


        To load the client from a configuration, you can use the `load_component` method:

        .. code-block:: python

            from autogen_core.models import ChatCompletionClient

            config = {
                "provider": "OpenAIChatCompletionClient",
                "config": {"model": "gpt-4o", "api_key": "REPLACE_WITH_YOUR_API_KEY"},
            }

            client = ChatCompletionClient.load_component(config)

        To view the full list of available configuration options, see the :py:class:`OpenAIClientConfigurationConfigModel` class.

    """

    component_type = "model"
    component_config_schema = OpenAIClientConfigurationConfigModel
    component_provider_override = "autogen_ext.models.openai.OpenAIChatCompletionClient"

    def __init__(self, **kwargs: Unpack[OpenAIClientConfiguration]):
        if "model" not in kwargs:
            raise ValueError("model is required for OpenAIChatCompletionClient")

        model_capabilities: Optional[ModelCapabilities] = None  # type: ignore
        self._raw_config: Dict[str, Any] = dict(kwargs).copy()
        copied_args = dict(kwargs).copy()

        if "model_capabilities" in kwargs:
            model_capabilities = kwargs["model_capabilities"]
            del copied_args["model_capabilities"]

        model_info: Optional[ModelInfo] = None
        if "model_info" in kwargs:
            model_info = kwargs["model_info"]
            del copied_args["model_info"]

        add_name_prefixes: bool = False
        if "add_name_prefixes" in kwargs:
            add_name_prefixes = kwargs["add_name_prefixes"]

        include_name_in_message: bool = True
        if "include_name_in_message" in kwargs:
            include_name_in_message = kwargs["include_name_in_message"]

        # Special handling for Gemini model.
        assert "model" in copied_args and isinstance(copied_args["model"], str)
        if copied_args["model"].startswith("gemini-"):
            if "base_url" not in copied_args:
                copied_args["base_url"] = _model_info.GEMINI_OPENAI_BASE_URL
            if "api_key" not in copied_args and "GEMINI_API_KEY" in os.environ:
                copied_args["api_key"] = os.environ["GEMINI_API_KEY"]
        if copied_args["model"].startswith("claude-"):
            if "base_url" not in copied_args:
                copied_args["base_url"] = _model_info.ANTHROPIC_OPENAI_BASE_URL
            if "api_key" not in copied_args and "ANTHROPIC_API_KEY" in os.environ:
                copied_args["api_key"] = os.environ["ANTHROPIC_API_KEY"]
        if copied_args["model"].startswith("Llama-"):
            if "base_url" not in copied_args:
                copied_args["base_url"] = _model_info.LLAMA_API_BASE_URL
            if "api_key" not in copied_args and "LLAMA_API_KEY" in os.environ:
                copied_args["api_key"] = os.environ["LLAMA_API_KEY"]

        client = _openai_client_from_config(copied_args)
        create_args = _create_args_from_config(copied_args)

        super().__init__(
            client=client,
            create_args=create_args,
            model_capabilities=model_capabilities,
            model_info=model_info,
            add_name_prefixes=add_name_prefixes,
            include_name_in_message=include_name_in_message,
        )

    def __getstate__(self) -> Dict[str, Any]:
        state = self.__dict__.copy()
        state["_client"] = None
        return state

    def __setstate__(self, state: Dict[str, Any]) -> None:
        self.__dict__.update(state)
        self._client = _openai_client_from_config(state["_raw_config"])

    def _to_config(self) -> OpenAIClientConfigurationConfigModel:
        copied_config = self._raw_config.copy()
        return OpenAIClientConfigurationConfigModel(**copied_config)

    @classmethod
    def _from_config(cls, config: OpenAIClientConfigurationConfigModel) -> Self:
        copied_config = config.model_copy().model_dump(exclude_none=True)

        # Handle api_key as SecretStr
        if "api_key" in copied_config and isinstance(config.api_key, SecretStr):
            copied_config["api_key"] = config.api_key.get_secret_value()

        return cls(**copied_config)