autogen/test/oai/test_client_stream.py

import json
from typing import Any, Dict, List, Literal, Optional, Union
from unittest.mock import MagicMock
import pytest
from autogen import OpenAIWrapper, config_list_from_json, config_list_openai_aoai
import sys
import os

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from conftest import skip_openai  # noqa: E402

try:
    from openai import OpenAI
except ImportError:
    skip = True
else:
    skip = False or skip_openai

    # raises exception if openai>=1 is installed and something is wrong with imports
    # otherwise the test will be skipped
    from openai.types.chat.chat_completion_chunk import (
        ChoiceDeltaFunctionCall,
        ChoiceDeltaToolCall,
        ChoiceDeltaToolCallFunction,
    )
    from openai.types.chat.chat_completion import ChatCompletionMessage  # type: ignore [attr-defined]

KEY_LOC = "notebook"
OAI_CONFIG_LIST = "OAI_CONFIG_LIST"


@pytest.mark.skipif(skip, reason="openai>=1 not installed")
def test_aoai_chat_completion_stream() -> None:
    config_list = config_list_from_json(
        env_or_file=OAI_CONFIG_LIST,
        file_location=KEY_LOC,
        filter_dict={"api_type": ["azure"], "model": ["gpt-3.5-turbo", "gpt-35-turbo"]},
    )
    client = OpenAIWrapper(config_list=config_list)
    response = client.create(messages=[{"role": "user", "content": "2+2="}], stream=True)
    print(response)
    print(client.extract_text_or_completion_object(response))


@pytest.mark.skipif(skip, reason="openai>=1 not installed")
def test_chat_completion_stream() -> None:
    config_list = config_list_from_json(
        env_or_file=OAI_CONFIG_LIST,
        file_location=KEY_LOC,
        filter_dict={"model": ["gpt-3.5-turbo", "gpt-35-turbo"]},
    )
    client = OpenAIWrapper(config_list=config_list)
    response = client.create(messages=[{"role": "user", "content": "1+1="}], stream=True)
    print(response)
    print(client.extract_text_or_completion_object(response))


# no need for OpenAI, works with any model
def test__update_dict_from_chunk() -> None:
    # dictionaries and lists are not supported
    mock = MagicMock()
    empty_collections: List[Union[List[Any], Dict[str, Any]]] = [{}, []]
    for c in empty_collections:
        mock.c = c
        with pytest.raises(NotImplementedError):
            OpenAIWrapper._update_dict_from_chunk(mock, {}, "c")

    org_d: Dict[str, Any] = {}
    for i, v in enumerate([0, 1, False, True, 0.0, 1.0]):
        field = "abcedfghijklmnopqrstuvwxyz"[i]
        setattr(mock, field, v)

        d = org_d.copy()
        OpenAIWrapper._update_dict_from_chunk(mock, d, field)

        org_d[field] = v
        assert d == org_d

    mock.s = "beginning"
    OpenAIWrapper._update_dict_from_chunk(mock, d, "s")
    assert d["s"] == "beginning"

    mock.s = " and"
    OpenAIWrapper._update_dict_from_chunk(mock, d, "s")
    assert d["s"] == "beginning and"

    mock.s = " end"
    OpenAIWrapper._update_dict_from_chunk(mock, d, "s")
    assert d["s"] == "beginning and end"


@pytest.mark.skipif(skip, reason="openai>=1 not installed")
def test__update_function_call_from_chunk() -> None:
    function_call_chunks = [
        ChoiceDeltaFunctionCall(arguments=None, name="get_current_weather"),
        ChoiceDeltaFunctionCall(arguments='{"', name=None),
        ChoiceDeltaFunctionCall(arguments="location", name=None),
        ChoiceDeltaFunctionCall(arguments='":"', name=None),
        ChoiceDeltaFunctionCall(arguments="San", name=None),
        ChoiceDeltaFunctionCall(arguments=" Francisco", name=None),
        ChoiceDeltaFunctionCall(arguments='"}', name=None),
    ]
    expected = {"name": "get_current_weather", "arguments": '{"location":"San Francisco"}'}

    full_function_call = None
    completion_tokens = 0
    for function_call_chunk in function_call_chunks:
        # print(f"{function_call_chunk=}")
        full_function_call, completion_tokens = OpenAIWrapper._update_function_call_from_chunk(
            function_call_chunk=function_call_chunk,
            full_function_call=full_function_call,
            completion_tokens=completion_tokens,
        )

    print(f"{full_function_call=}")
    print(f"{completion_tokens=}")

    assert full_function_call == expected
    assert completion_tokens == len(function_call_chunks)

    ChatCompletionMessage(role="assistant", function_call=full_function_call, content=None)


@pytest.mark.skipif(skip, reason="openai>=1 not installed")
def test__update_tool_calls_from_chunk() -> None:
    tool_calls_chunks = [
        ChoiceDeltaToolCall(
            index=0,
            id="call_D2HOWGMekmkxXu9Ix3DUqJRv",
            function=ChoiceDeltaToolCallFunction(arguments="", name="get_current_weather"),
            type="function",
        ),
        ChoiceDeltaToolCall(
            index=0, id=None, function=ChoiceDeltaToolCallFunction(arguments='{"lo', name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=0, id=None, function=ChoiceDeltaToolCallFunction(arguments="catio", name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=0, id=None, function=ChoiceDeltaToolCallFunction(arguments='n": "S', name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=0, id=None, function=ChoiceDeltaToolCallFunction(arguments="an F", name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=0, id=None, function=ChoiceDeltaToolCallFunction(arguments="ranci", name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=0, id=None, function=ChoiceDeltaToolCallFunction(arguments="sco, C", name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=0, id=None, function=ChoiceDeltaToolCallFunction(arguments='A"}', name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=1,
            id="call_22HgJep4nwoKU3UOr96xaLmd",
            function=ChoiceDeltaToolCallFunction(arguments="", name="get_current_weather"),
            type="function",
        ),
        ChoiceDeltaToolCall(
            index=1, id=None, function=ChoiceDeltaToolCallFunction(arguments='{"lo', name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=1, id=None, function=ChoiceDeltaToolCallFunction(arguments="catio", name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=1, id=None, function=ChoiceDeltaToolCallFunction(arguments='n": "N', name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=1, id=None, function=ChoiceDeltaToolCallFunction(arguments="ew Y", name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=1, id=None, function=ChoiceDeltaToolCallFunction(arguments="ork, ", name=None), type=None
        ),
        ChoiceDeltaToolCall(
            index=1, id=None, function=ChoiceDeltaToolCallFunction(arguments='NY"}', name=None), type=None
        ),
    ]

    full_tool_calls: List[Optional[Dict[str, Any]]] = [None, None]
    completion_tokens = 0
    for tool_calls_chunk in tool_calls_chunks:
        index = tool_calls_chunk.index
        full_tool_calls[index], completion_tokens = OpenAIWrapper._update_tool_calls_from_chunk(
            tool_calls_chunk=tool_calls_chunk,
            full_tool_call=full_tool_calls[index],
            completion_tokens=completion_tokens,
        )

    print(f"{full_tool_calls=}")
    print(f"{completion_tokens=}")

    ChatCompletionMessage(role="assistant", tool_calls=full_tool_calls, content=None)


# todo: remove when OpenAI removes functions from the API
@pytest.mark.skipif(skip, reason="openai>=1 not installed")
def test_chat_functions_stream() -> None:
    config_list = config_list_from_json(
        env_or_file=OAI_CONFIG_LIST,
        file_location=KEY_LOC,
        filter_dict={"model": ["gpt-3.5-turbo", "gpt-35-turbo"]},
    )
    functions = [
        {
            "name": "get_current_weather",
            "description": "Get the current weather",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city and state, e.g. San Francisco, CA",
                    },
                },
                "required": ["location"],
            },
        },
    ]
    client = OpenAIWrapper(config_list=config_list)
    response = client.create(
        messages=[{"role": "user", "content": "What's the weather like today in San Francisco?"}],
        functions=functions,
        stream=True,
    )
    print(response)
    print(client.extract_text_or_completion_object(response))


# test for tool support instead of the deprecated function calls
@pytest.mark.skipif(skip, reason="openai>=1 not installed")
def test_chat_tools_stream() -> None:
    config_list = config_list_from_json(
        env_or_file=OAI_CONFIG_LIST,
        file_location=KEY_LOC,
        filter_dict={"model": ["gpt-3.5-turbo", "gpt-35-turbo"]},
    )
    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_current_weather",
                "description": "Get the current weather",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The city and state, e.g. San Francisco, CA",
                        },
                    },
                    "required": ["location"],
                },
            },
        },
    ]
    print(f"{config_list=}")
    client = OpenAIWrapper(config_list=config_list)
    response = client.create(
        # the intention is to trigger two tool invocations as a response to a single message
        messages=[{"role": "user", "content": "What's the weather like today in San Francisco and New York?"}],
        tools=tools,
        stream=True,
    )
    print(f"{response=}")
    print(f"{type(response)=}")
    print(f"{client.extract_text_or_completion_object(response)=}")
    # check response
    choices = response.choices
    assert isinstance(choices, list)
    assert len(choices) == 1
    choice = choices[0]
    assert choice.finish_reason == "tool_calls"
    message = choice.message
    tool_calls = message.tool_calls
    assert isinstance(tool_calls, list)
    assert len(tool_calls) == 2
    arguments = [tool_call.function.arguments for tool_call in tool_calls]
    locations = [json.loads(argument)["location"] for argument in arguments]
    print(f"{locations=}")
    assert any(["San Francisco" in location for location in locations])
    assert any(["New York" in location for location in locations])


@pytest.mark.skipif(skip, reason="openai>=1 not installed")
def test_completion_stream() -> None:
    config_list = config_list_openai_aoai(KEY_LOC)
    client = OpenAIWrapper(config_list=config_list)
    # Azure can't have dot in model/deployment name
    model = "gpt-35-turbo-instruct" if config_list[0].get("api_type") == "azure" else "gpt-3.5-turbo-instruct"
    response = client.create(prompt="1+1=", model=model, stream=True)
    print(response)
    print(client.extract_text_or_completion_object(response))


if __name__ == "__main__":
    test_aoai_chat_completion_stream()
    test_chat_completion_stream()
    test_chat_functions_stream()
    test_completion_stream()