mirror of
https://github.com/microsoft/autogen.git
synced 2026-04-20 03:02:16 -04:00
f812a238b1
* add cost calculation * Update autogen/oai/client.py Co-authored-by: Joshua Kim <joshkyh@users.noreply.github.com> * Update autogen/oai/client.py Co-authored-by: Joshua Kim <joshkyh@users.noreply.github.com> * update * add doc --------- Co-authored-by: Joshua Kim <joshkyh@users.noreply.github.com>
72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
import pytest
|
|
from autogen import OpenAIWrapper, config_list_from_json, config_list_openai_aoai
|
|
from test_utils import OAI_CONFIG_LIST, KEY_LOC
|
|
|
|
try:
|
|
from openai import OpenAI
|
|
except ImportError:
|
|
skip = True
|
|
else:
|
|
skip = False
|
|
|
|
|
|
@pytest.mark.skipif(skip, reason="openai>=1 not installed")
|
|
def test_aoai_chat_completion():
|
|
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"]},
|
|
)
|
|
client = OpenAIWrapper(config_list=config_list)
|
|
# for config in config_list:
|
|
# print(config)
|
|
# client = OpenAIWrapper(**config)
|
|
# response = client.create(messages=[{"role": "user", "content": "2+2="}], cache_seed=None)
|
|
response = client.create(messages=[{"role": "user", "content": "2+2="}], cache_seed=None)
|
|
print(response)
|
|
print(client.extract_text_or_function_call(response))
|
|
|
|
|
|
@pytest.mark.skipif(skip, reason="openai>=1 not installed")
|
|
def test_chat_completion():
|
|
config_list = config_list_from_json(
|
|
env_or_file=OAI_CONFIG_LIST,
|
|
file_location=KEY_LOC,
|
|
)
|
|
client = OpenAIWrapper(config_list=config_list)
|
|
response = client.create(messages=[{"role": "user", "content": "1+1="}])
|
|
print(response)
|
|
print(client.extract_text_or_function_call(response))
|
|
|
|
|
|
@pytest.mark.skipif(skip, reason="openai>=1 not installed")
|
|
def test_completion():
|
|
config_list = config_list_openai_aoai(KEY_LOC)
|
|
client = OpenAIWrapper(config_list=config_list)
|
|
response = client.create(prompt="1+1=", model="gpt-3.5-turbo-instruct")
|
|
print(response)
|
|
print(client.extract_text_or_function_call(response))
|
|
|
|
|
|
@pytest.mark.skipif(skip, reason="openai>=1 not installed")
|
|
@pytest.mark.parametrize(
|
|
"cache_seed, model",
|
|
[
|
|
(None, "gpt-3.5-turbo-instruct"),
|
|
(42, "gpt-3.5-turbo-instruct"),
|
|
(None, "text-ada-001"),
|
|
],
|
|
)
|
|
def test_cost(cache_seed, model):
|
|
config_list = config_list_openai_aoai(KEY_LOC)
|
|
client = OpenAIWrapper(config_list=config_list, cache_seed=cache_seed)
|
|
response = client.create(prompt="1+3=", model=model)
|
|
print(response.cost)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_aoai_chat_completion()
|
|
test_chat_completion()
|
|
test_completion()
|
|
test_cost()
|