From eae458929a50bb798223b1d70ed8ed62f65afc9c Mon Sep 17 00:00:00 2001 From: Alexander Barannikov Date: Wed, 11 Dec 2024 18:37:06 +0000 Subject: [PATCH 1/7] OPIK-569 [SDK] Implement integration with aisuite [wip] --- .../src/opik/integrations/aisuite/__init__.py | 4 + .../integrations/aisuite/aisuite_decorator.py | 120 ++++++++ .../opik/integrations/aisuite/opik_tracker.py | 43 +++ .../library_integration/aisuite/__init__.py | 0 .../aisuite/requirements.txt | 1 + .../aisuite/test_aisuite.py | 262 ++++++++++++++++++ 6 files changed, 430 insertions(+) create mode 100644 sdks/python/src/opik/integrations/aisuite/__init__.py create mode 100644 sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py create mode 100644 sdks/python/src/opik/integrations/aisuite/opik_tracker.py create mode 100644 sdks/python/tests/library_integration/aisuite/__init__.py create mode 100644 sdks/python/tests/library_integration/aisuite/requirements.txt create mode 100644 sdks/python/tests/library_integration/aisuite/test_aisuite.py diff --git a/sdks/python/src/opik/integrations/aisuite/__init__.py b/sdks/python/src/opik/integrations/aisuite/__init__.py new file mode 100644 index 0000000000..74c9026332 --- /dev/null +++ b/sdks/python/src/opik/integrations/aisuite/__init__.py @@ -0,0 +1,4 @@ +from .opik_tracker import track_aisuite + + +__all__ = ["track_aisuite"] diff --git a/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py b/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py new file mode 100644 index 0000000000..73f21edd37 --- /dev/null +++ b/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py @@ -0,0 +1,120 @@ +import logging +from typing import ( + List, + Any, + Dict, + Optional, + Callable, + Tuple, + Union, + Iterator, + AsyncIterator, +) + +from aisuite.framework import ChatCompletionResponse + +from opik import dict_utils +from opik.decorator import base_track_decorator, arguments_helpers +# from . import stream_patchers + +import aisuite +import openai +from openai.types.chat import chat_completion, chat_completion_chunk +from openai import _types as _openai_types + +LOGGER = logging.getLogger(__name__) + +# CreateCallResult = Union[chat_completion.ChatCompletion, List[Any]] + +# KWARGS_KEYS_TO_LOG_AS_INPUTS = ["messages", "function_call"] +KWARGS_KEYS_TO_LOG_AS_INPUTS = ["messages"] +RESPONSE_KEYS_TO_LOG_AS_OUTPUT = ["choices"] + + +class AISuiteTrackDecorator(base_track_decorator.BaseTrackDecorator): + """ + An implementation of BaseTrackDecorator designed specifically for tracking + calls of AISuite's `chat.completion.create` and `chat.completions.parse` functions. + + Besides special processing for input arguments and response content, it + overrides _generators_handler() method to work correctly with + openai.Stream and openai.AsyncStream objects. + """ + + # def __init__(self) -> None: + # super().__init__() + # self.provider = "openai" + + def _start_span_inputs_preprocessor( + self, + func: Callable, + track_options: arguments_helpers.TrackOptions, + args: Optional[Tuple], + kwargs: Optional[Dict[str, Any]], + ) -> arguments_helpers.StartSpanParameters: + assert ( + kwargs is not None + ), "Expected kwargs to be not None in chat.completion.create(**kwargs)" + + name = track_options.name if track_options.name is not None else func.__name__ + metadata = track_options.metadata if track_options.metadata is not None else {} + + input, new_metadata = dict_utils.split_dict_by_keys( + kwargs, keys=KWARGS_KEYS_TO_LOG_AS_INPUTS + ) + metadata = dict_utils.deepmerge(metadata, new_metadata) + metadata.update( + { + "created_from": "aisuite", + "type": "aisuite_chat", + } + ) + + tags = ["aisuite"] + + result = arguments_helpers.StartSpanParameters( + name=name, + input=input, + type=track_options.type, + tags=tags, + metadata=metadata, + project_name=track_options.project_name, + model=kwargs.get("model", None), + # provider=self.provider, + ) + + return result + + def _end_span_inputs_preprocessor( + self, output: Any, capture_output: bool + ) -> arguments_helpers.EndSpanParameters: + assert isinstance( + output, + ( + chat_completion.ChatCompletion, + ChatCompletionResponse, + ) + ) + + result_dict = output.model_dump(mode="json") + output, metadata = dict_utils.split_dict_by_keys(result_dict, ["choices"]) + usage = result_dict["usage"] + model = result_dict["model"] + + result = arguments_helpers.EndSpanParameters( + output=output, + usage=usage, + metadata=metadata, + model=model, + # provider=self.provider, + ) + + return result + + def _generators_handler( + self, + output: Any, + capture_output: bool, + generations_aggregator: Optional[Callable[[List[Any]], Any]], + ) -> None: + return None diff --git a/sdks/python/src/opik/integrations/aisuite/opik_tracker.py b/sdks/python/src/opik/integrations/aisuite/opik_tracker.py new file mode 100644 index 0000000000..f814b26389 --- /dev/null +++ b/sdks/python/src/opik/integrations/aisuite/opik_tracker.py @@ -0,0 +1,43 @@ +from typing import Optional + +import aisuite + +from . import aisuite_decorator + + +def track_aisuite( + aisuite_client: aisuite.Client, + project_name: Optional[str] = None, +) -> aisuite.Client: + """Adds Opik tracking to an AISuite client. + + Tracks calls to: + * `aisuite_client.chat.completions.create()`, + + Can be used within other Opik-tracked functions. + + Args: + aisuite_client: An instance of AISuite client. + project_name: The name of the project to log data. + + Returns: + The modified AISuite client with Opik tracking enabled. + """ + if hasattr(aisuite_client, "opik_tracked"): + return aisuite_client + + aisuite_client.opik_tracked = True + + decorator_factory = aisuite_decorator.AISuiteTrackDecorator() + + completions_create_decorator = decorator_factory.track( + type="llm", + name="chat_completion_create", + project_name=project_name, + ) + + aisuite_client.chat.completions.create = completions_create_decorator( + aisuite_client.chat.completions.create + ) + + return aisuite_client diff --git a/sdks/python/tests/library_integration/aisuite/__init__.py b/sdks/python/tests/library_integration/aisuite/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/sdks/python/tests/library_integration/aisuite/requirements.txt b/sdks/python/tests/library_integration/aisuite/requirements.txt new file mode 100644 index 0000000000..aaab7cf072 --- /dev/null +++ b/sdks/python/tests/library_integration/aisuite/requirements.txt @@ -0,0 +1 @@ +aisuite diff --git a/sdks/python/tests/library_integration/aisuite/test_aisuite.py b/sdks/python/tests/library_integration/aisuite/test_aisuite.py new file mode 100644 index 0000000000..db2d5d864d --- /dev/null +++ b/sdks/python/tests/library_integration/aisuite/test_aisuite.py @@ -0,0 +1,262 @@ +import asyncio +from typing import Any, Dict, List + +import aisuite +import pytest +from pydantic import BaseModel + +import opik +from opik.config import OPIK_PROJECT_DEFAULT_NAME +from opik.integrations.aisuite import track_aisuite +from ...testlib import ( + ANY_BUT_NONE, + ANY_DICT, + ANY_LIST, + ANY_STRING, + SpanModel, + TraceModel, + assert_dict_has_keys, + assert_equal, +) + + +pytestmark = pytest.mark.usefixtures("ensure_openai_configured") + + +def _assert_metadata_contains_required_keys(metadata: Dict[str, Any]): + REQUIRED_METADATA_KEYS = [ + "usage", + "model", + "max_tokens", + "created_from", + "type", + "id", + "created", + "object", + ] + assert_dict_has_keys(metadata, REQUIRED_METADATA_KEYS) + + +@pytest.mark.parametrize( + "project_name, expected_project_name", + [ + # (None, OPIK_PROJECT_DEFAULT_NAME), + ("aisuite-integration-test", "aisuite-integration-test"), + ], +) +def test_aisuite_client_chat_completions_create__happyflow( + fake_backend, project_name, expected_project_name +): + client = aisuite.Client() + wrapped_client = track_aisuite( + aisuite_client=client, + project_name=project_name, + ) + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Tell a fact"}, + ] + + _ = wrapped_client.chat.completions.create( + model="openai:gpt-3.5-turbo", + messages=messages, + max_tokens=10, + ) + + opik.flush_tracker() + + EXPECTED_TRACE_TREE = TraceModel( + id=ANY_BUT_NONE, + name="chat_completion_create", + input={"messages": messages}, + output={"choices": ANY_BUT_NONE}, + tags=["openai"], + metadata=ANY_DICT, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=expected_project_name, + spans=[ + SpanModel( + id=ANY_BUT_NONE, + type="llm", + name="chat_completion_create", + input={"messages": messages}, + output={"choices": ANY_BUT_NONE}, + tags=["openai"], + metadata=ANY_DICT, + usage={ + "prompt_tokens": ANY_BUT_NONE, + "completion_tokens": ANY_BUT_NONE, + "total_tokens": ANY_BUT_NONE, + }, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=expected_project_name, + spans=[], + model=ANY_STRING(startswith="gpt-3.5-turbo"), + provider="openai", + ) + ], + ) + + assert len(fake_backend.trace_trees) == 1 + trace_tree = fake_backend.trace_trees[0] + + # assert_equal(EXPECTED_TRACE_TREE, trace_tree) + # + # llm_span_metadata = trace_tree.spans[0].metadata + # _assert_metadata_contains_required_keys(llm_span_metadata) + + +@pytest.mark.skip +def test_openai_client_chat_completions_create__create_raises_an_error__span_and_trace_finished_gracefully__error_info_is_logged( + fake_backend, +): + client = aisuite.Client() + wrapped_client = track_aisuite(aisuite_client=client) + + with pytest.raises(openai.OpenAIError): + _ = wrapped_client.chat.completions.create( + messages=None, + model=None, + ) + + opik.flush_tracker() + + EXPECTED_TRACE_TREE = TraceModel( + id=ANY_BUT_NONE, + name="chat_completion_create", + input={"messages": None}, + output=None, + tags=["openai"], + metadata={ + "created_from": "openai", + "type": "openai_chat", + "model": None, + }, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=ANY_BUT_NONE, + error_info={ + "exception_type": ANY_STRING(), + "message": ANY_STRING(), + "traceback": ANY_STRING(), + }, + spans=[ + SpanModel( + id=ANY_BUT_NONE, + type="llm", + name="chat_completion_create", + input={"messages": None}, + output=None, + tags=["openai"], + metadata={ + "created_from": "openai", + "type": "openai_chat", + "model": None, + }, + usage=None, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=ANY_BUT_NONE, + model=None, + provider="openai", + error_info={ + "exception_type": ANY_STRING(), + "message": ANY_STRING(), + "traceback": ANY_STRING(), + }, + spans=[], + ) + ], + ) + + assert len(fake_backend.trace_trees) == 1 + + trace_tree = fake_backend.trace_trees[0] + assert_equal(EXPECTED_TRACE_TREE, trace_tree) + + +@pytest.mark.skip +def test_aisuite_client_chat_completions_create__openai_call_made_in_another_tracked_function__openai_span_attached_to_existing_trace( + fake_backend, +): + project_name = "aisuite-integration-test" + + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Tell a fact"}, + ] + + @opik.track(project_name=project_name) + def f(): + client = aisuite.Client() + wrapped_client = track_aisuite( + aisuite_client=client, + # we are trying to log span into another project, but parent's project name will be used + project_name="aisuite-integration-test-nested-level", + ) + + _ = wrapped_client.chat.completions.create( + model="gpt-3.5-turbo", + messages=messages, + max_tokens=10, + ) + + f() + + opik.flush_tracker() + + EXPECTED_TRACE_TREE = TraceModel( + id=ANY_BUT_NONE, + name="f", + input={}, + output=None, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=project_name, + spans=[ + SpanModel( + id=ANY_BUT_NONE, + name="f", + input={}, + output=None, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=project_name, + model=None, + provider=None, + spans=[ + SpanModel( + id=ANY_BUT_NONE, + type="llm", + name="chat_completion_create", + input={"messages": messages}, + output={"choices": ANY_BUT_NONE}, + tags=["openai"], + metadata=ANY_DICT, + usage={ + "prompt_tokens": ANY_BUT_NONE, + "completion_tokens": ANY_BUT_NONE, + "total_tokens": ANY_BUT_NONE, + }, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=project_name, + spans=[], + model=ANY_STRING(startswith="gpt-3.5-turbo"), + provider="openai", + ) + ], + ) + ], + ) + + assert len(fake_backend.trace_trees) == 1 + + trace_tree = fake_backend.trace_trees[0] + + assert_equal(EXPECTED_TRACE_TREE, trace_tree) + + llm_span_metadata = trace_tree.spans[0].spans[0].metadata + _assert_metadata_contains_required_keys(llm_span_metadata) From 5cc9350aabe2cdd633b66709222e03b621be8c51 Mon Sep 17 00:00:00 2001 From: Alexander Barannikov Date: Fri, 13 Dec 2024 11:43:41 +0000 Subject: [PATCH 2/7] RC 1 --- .../integrations/aisuite/aisuite_decorator.py | 85 +++++++----- .../aisuite/test_aisuite.py | 123 +++++++++++++----- 2 files changed, 140 insertions(+), 68 deletions(-) diff --git a/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py b/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py index 73f21edd37..f6b87051fc 100644 --- a/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py +++ b/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py @@ -1,34 +1,21 @@ import logging -from typing import ( - List, - Any, - Dict, - Optional, - Callable, - Tuple, - Union, - Iterator, - AsyncIterator, -) +from typing import (Any, Callable, Dict, List, Optional, Tuple) from aisuite.framework import ChatCompletionResponse +from openai.types.chat import chat_completion from opik import dict_utils -from opik.decorator import base_track_decorator, arguments_helpers -# from . import stream_patchers +from opik.decorator import arguments_helpers, base_track_decorator -import aisuite -import openai -from openai.types.chat import chat_completion, chat_completion_chunk -from openai import _types as _openai_types +# from . import stream_patchers LOGGER = logging.getLogger(__name__) -# CreateCallResult = Union[chat_completion.ChatCompletion, List[Any]] - # KWARGS_KEYS_TO_LOG_AS_INPUTS = ["messages", "function_call"] KWARGS_KEYS_TO_LOG_AS_INPUTS = ["messages"] -RESPONSE_KEYS_TO_LOG_AS_OUTPUT = ["choices"] + + +# RESPONSE_KEYS_TO_LOG_AS_OUTPUT = ["choices"] class AISuiteTrackDecorator(base_track_decorator.BaseTrackDecorator): @@ -41,10 +28,6 @@ class AISuiteTrackDecorator(base_track_decorator.BaseTrackDecorator): openai.Stream and openai.AsyncStream objects. """ - # def __init__(self) -> None: - # super().__init__() - # self.provider = "openai" - def _start_span_inputs_preprocessor( self, func: Callable, @@ -72,6 +55,8 @@ def _start_span_inputs_preprocessor( tags = ["aisuite"] + model, provider = self._get_provider_info(func, **kwargs) + result = arguments_helpers.StartSpanParameters( name=name, input=input, @@ -79,34 +64,68 @@ def _start_span_inputs_preprocessor( tags=tags, metadata=metadata, project_name=track_options.project_name, - model=kwargs.get("model", None), - # provider=self.provider, + model=model, + provider=provider, ) return result + def _get_provider_info(self, func: Callable, **kwargs) -> Tuple[Optional[str], Optional[str]]: + provider: Optional[str] = None + model: Optional[str] = kwargs.get("model", None) + + if model is not None and ":" in model: + provider, model = model.split(":", 1) + + if provider != "openai": + return model, provider + elif base_url_provider := func.__self__.client.providers.get("openai"): + base_url = base_url_provider.client.base_url + if base_url.host != "api.openai.com": + provider = base_url.host + + return model, provider + def _end_span_inputs_preprocessor( self, output: Any, capture_output: bool ) -> arguments_helpers.EndSpanParameters: assert isinstance( output, ( - chat_completion.ChatCompletion, - ChatCompletionResponse, + chat_completion.ChatCompletion, # openai + ChatCompletionResponse, # non-openai ) ) - result_dict = output.model_dump(mode="json") - output, metadata = dict_utils.split_dict_by_keys(result_dict, ["choices"]) - usage = result_dict["usage"] - model = result_dict["model"] + metadata = None + usage = None + model = None + + # provider == openai + if isinstance(output, chat_completion.ChatCompletion): + result_dict = output.model_dump(mode="json") + output, metadata = dict_utils.split_dict_by_keys(result_dict, ["choices"]) + usage = result_dict["usage"] + model = result_dict["model"] + + # provider == non-openai + elif isinstance(output, ChatCompletionResponse): + choices = [] + + for choice in output.choices: + choices.append( + { + "message": {"content": choice.message.content}, + } + ) + + output = {"choices": choices} result = arguments_helpers.EndSpanParameters( output=output, usage=usage, metadata=metadata, model=model, - # provider=self.provider, ) return result diff --git a/sdks/python/tests/library_integration/aisuite/test_aisuite.py b/sdks/python/tests/library_integration/aisuite/test_aisuite.py index db2d5d864d..9f2d675f07 100644 --- a/sdks/python/tests/library_integration/aisuite/test_aisuite.py +++ b/sdks/python/tests/library_integration/aisuite/test_aisuite.py @@ -2,6 +2,7 @@ from typing import Any, Dict, List import aisuite +import openai import pytest from pydantic import BaseModel @@ -22,6 +23,8 @@ pytestmark = pytest.mark.usefixtures("ensure_openai_configured") +PROJECT_NAME = "aisuite-integration-test" + def _assert_metadata_contains_required_keys(metadata: Dict[str, Any]): REQUIRED_METADATA_KEYS = [ @@ -37,20 +40,11 @@ def _assert_metadata_contains_required_keys(metadata: Dict[str, Any]): assert_dict_has_keys(metadata, REQUIRED_METADATA_KEYS) -@pytest.mark.parametrize( - "project_name, expected_project_name", - [ - # (None, OPIK_PROJECT_DEFAULT_NAME), - ("aisuite-integration-test", "aisuite-integration-test"), - ], -) -def test_aisuite_client_chat_completions_create__happyflow( - fake_backend, project_name, expected_project_name -): +def test_aisuite__openai_provider__client_chat_completions_create__happyflow(fake_backend): client = aisuite.Client() wrapped_client = track_aisuite( aisuite_client=client, - project_name=project_name, + project_name=PROJECT_NAME, ) messages = [ {"role": "system", "content": "You are a helpful assistant."}, @@ -70,11 +64,11 @@ def test_aisuite_client_chat_completions_create__happyflow( name="chat_completion_create", input={"messages": messages}, output={"choices": ANY_BUT_NONE}, - tags=["openai"], + tags=["aisuite"], metadata=ANY_DICT, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, - project_name=expected_project_name, + project_name=PROJECT_NAME, spans=[ SpanModel( id=ANY_BUT_NONE, @@ -82,7 +76,7 @@ def test_aisuite_client_chat_completions_create__happyflow( name="chat_completion_create", input={"messages": messages}, output={"choices": ANY_BUT_NONE}, - tags=["openai"], + tags=["aisuite"], metadata=ANY_DICT, usage={ "prompt_tokens": ANY_BUT_NONE, @@ -91,7 +85,7 @@ def test_aisuite_client_chat_completions_create__happyflow( }, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, - project_name=expected_project_name, + project_name=PROJECT_NAME, spans=[], model=ANY_STRING(startswith="gpt-3.5-turbo"), provider="openai", @@ -102,23 +96,82 @@ def test_aisuite_client_chat_completions_create__happyflow( assert len(fake_backend.trace_trees) == 1 trace_tree = fake_backend.trace_trees[0] - # assert_equal(EXPECTED_TRACE_TREE, trace_tree) - # - # llm_span_metadata = trace_tree.spans[0].metadata - # _assert_metadata_contains_required_keys(llm_span_metadata) + assert_equal(EXPECTED_TRACE_TREE, trace_tree) + + llm_span_metadata = trace_tree.spans[0].metadata + _assert_metadata_contains_required_keys(llm_span_metadata) -@pytest.mark.skip -def test_openai_client_chat_completions_create__create_raises_an_error__span_and_trace_finished_gracefully__error_info_is_logged( +def test_aisuite__nonopenai_provider__client_chat_completions_create__happyflow(fake_backend): + client = aisuite.Client() + wrapped_client = track_aisuite( + aisuite_client=client, + project_name=PROJECT_NAME, + ) + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Tell a fact"}, + ] + + _ = wrapped_client.chat.completions.create( + # model="anthropic:claude-3-5-sonnet-20240620", + model="anthropic:claude-3-5-sonnet-latest", + messages=messages, + max_tokens=10, + ) + + opik.flush_tracker() + + EXPECTED_TRACE_TREE = TraceModel( + id=ANY_BUT_NONE, + name="chat_completion_create", + input={"messages": messages}, + output={"choices": ANY_BUT_NONE}, + tags=["aisuite"], + metadata=ANY_DICT, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=PROJECT_NAME, + spans=[ + SpanModel( + id=ANY_BUT_NONE, + type="llm", + name="chat_completion_create", + input={"messages": messages}, + output={"choices": ANY_BUT_NONE}, + tags=["aisuite"], + metadata=ANY_DICT, + usage=None, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=PROJECT_NAME, + spans=[], + model=ANY_STRING(startswith="claude-3-5-sonnet"), + provider="anthropic", + ) + ], + ) + + assert len(fake_backend.trace_trees) == 1 + trace_tree = fake_backend.trace_trees[0] + + assert_equal(EXPECTED_TRACE_TREE, trace_tree) + + +def test_aisuite_client_chat_completions_create__create_raises_an_error__span_and_trace_finished_gracefully__error_info_is_logged( fake_backend, ): client = aisuite.Client() - wrapped_client = track_aisuite(aisuite_client=client) + wrapped_client = track_aisuite( + aisuite_client=client, + project_name=PROJECT_NAME, + ) - with pytest.raises(openai.OpenAIError): + # with pytest.raises(openai.OpenAIError): + with pytest.raises(openai.BadRequestError): _ = wrapped_client.chat.completions.create( messages=None, - model=None, + model="openai:gpt-3.5-turbo", ) opik.flush_tracker() @@ -128,15 +181,15 @@ def test_openai_client_chat_completions_create__create_raises_an_error__span_and name="chat_completion_create", input={"messages": None}, output=None, - tags=["openai"], + tags=["aisuite"], metadata={ - "created_from": "openai", - "type": "openai_chat", - "model": None, + "created_from": "aisuite", + "type": "aisuite_chat", + "model": 'openai:gpt-3.5-turbo', }, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, - project_name=ANY_BUT_NONE, + project_name=PROJECT_NAME, error_info={ "exception_type": ANY_STRING(), "message": ANY_STRING(), @@ -149,17 +202,17 @@ def test_openai_client_chat_completions_create__create_raises_an_error__span_and name="chat_completion_create", input={"messages": None}, output=None, - tags=["openai"], + tags=["aisuite"], metadata={ - "created_from": "openai", - "type": "openai_chat", - "model": None, + "created_from": "aisuite", + "type": "aisuite_chat", + "model": 'openai:gpt-3.5-turbo', }, usage=None, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, - project_name=ANY_BUT_NONE, - model=None, + project_name=PROJECT_NAME, + model=ANY_STRING(startswith="gpt-3.5-turbo"), provider="openai", error_info={ "exception_type": ANY_STRING(), From 8a54adc97b3c6c123dd60703e22e5cf453050f36 Mon Sep 17 00:00:00 2001 From: Alexander Barannikov Date: Fri, 13 Dec 2024 11:46:24 +0000 Subject: [PATCH 3/7] RC 2 --- .../library_integration/aisuite/test_aisuite.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/sdks/python/tests/library_integration/aisuite/test_aisuite.py b/sdks/python/tests/library_integration/aisuite/test_aisuite.py index 9f2d675f07..e79a1de71a 100644 --- a/sdks/python/tests/library_integration/aisuite/test_aisuite.py +++ b/sdks/python/tests/library_integration/aisuite/test_aisuite.py @@ -230,28 +230,26 @@ def test_aisuite_client_chat_completions_create__create_raises_an_error__span_an assert_equal(EXPECTED_TRACE_TREE, trace_tree) -@pytest.mark.skip def test_aisuite_client_chat_completions_create__openai_call_made_in_another_tracked_function__openai_span_attached_to_existing_trace( fake_backend, ): - project_name = "aisuite-integration-test" messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Tell a fact"}, ] - @opik.track(project_name=project_name) + @opik.track(project_name=PROJECT_NAME) def f(): client = aisuite.Client() wrapped_client = track_aisuite( aisuite_client=client, # we are trying to log span into another project, but parent's project name will be used - project_name="aisuite-integration-test-nested-level", + project_name=f"{PROJECT_NAME}-nested-level", ) _ = wrapped_client.chat.completions.create( - model="gpt-3.5-turbo", + model="openai:gpt-3.5-turbo", messages=messages, max_tokens=10, ) @@ -267,7 +265,7 @@ def f(): output=None, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, - project_name=project_name, + project_name=PROJECT_NAME, spans=[ SpanModel( id=ANY_BUT_NONE, @@ -276,7 +274,7 @@ def f(): output=None, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, - project_name=project_name, + project_name=PROJECT_NAME, model=None, provider=None, spans=[ @@ -286,7 +284,7 @@ def f(): name="chat_completion_create", input={"messages": messages}, output={"choices": ANY_BUT_NONE}, - tags=["openai"], + tags=["aisuite"], metadata=ANY_DICT, usage={ "prompt_tokens": ANY_BUT_NONE, @@ -295,7 +293,7 @@ def f(): }, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, - project_name=project_name, + project_name=PROJECT_NAME, spans=[], model=ANY_STRING(startswith="gpt-3.5-turbo"), provider="openai", From cd5409a876a9d42aeea902a74bbdd2f42bcd830c Mon Sep 17 00:00:00 2001 From: Alexander Barannikov Date: Fri, 13 Dec 2024 11:52:37 +0000 Subject: [PATCH 4/7] RC 3 --- .../integrations/aisuite/aisuite_decorator.py | 19 +++++++++--------- .../aisuite/test_aisuite.py | 20 +++++++++---------- 2 files changed, 19 insertions(+), 20 deletions(-) diff --git a/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py b/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py index f6b87051fc..eb17710fca 100644 --- a/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py +++ b/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py @@ -1,5 +1,5 @@ import logging -from typing import (Any, Callable, Dict, List, Optional, Tuple) +from typing import Any, Callable, Dict, List, Optional, Tuple from aisuite.framework import ChatCompletionResponse from openai.types.chat import chat_completion @@ -7,17 +7,12 @@ from opik import dict_utils from opik.decorator import arguments_helpers, base_track_decorator -# from . import stream_patchers LOGGER = logging.getLogger(__name__) -# KWARGS_KEYS_TO_LOG_AS_INPUTS = ["messages", "function_call"] KWARGS_KEYS_TO_LOG_AS_INPUTS = ["messages"] -# RESPONSE_KEYS_TO_LOG_AS_OUTPUT = ["choices"] - - class AISuiteTrackDecorator(base_track_decorator.BaseTrackDecorator): """ An implementation of BaseTrackDecorator designed specifically for tracking @@ -70,7 +65,11 @@ def _start_span_inputs_preprocessor( return result - def _get_provider_info(self, func: Callable, **kwargs) -> Tuple[Optional[str], Optional[str]]: + def _get_provider_info( + self, + func: Callable, + **kwargs: Any, + ) -> Tuple[Optional[str], Optional[str]]: provider: Optional[str] = None model: Optional[str] = kwargs.get("model", None) @@ -79,7 +78,9 @@ def _get_provider_info(self, func: Callable, **kwargs) -> Tuple[Optional[str], O if provider != "openai": return model, provider - elif base_url_provider := func.__self__.client.providers.get("openai"): + + if hasattr(func, "__self__") and func.__self__.client.providers.get("openai"): + base_url_provider = func.__self__.client.providers.get("openai") base_url = base_url_provider.client.base_url if base_url.host != "api.openai.com": provider = base_url.host @@ -94,7 +95,7 @@ def _end_span_inputs_preprocessor( ( chat_completion.ChatCompletion, # openai ChatCompletionResponse, # non-openai - ) + ), ) metadata = None diff --git a/sdks/python/tests/library_integration/aisuite/test_aisuite.py b/sdks/python/tests/library_integration/aisuite/test_aisuite.py index e79a1de71a..0d8fa2ded4 100644 --- a/sdks/python/tests/library_integration/aisuite/test_aisuite.py +++ b/sdks/python/tests/library_integration/aisuite/test_aisuite.py @@ -1,18 +1,14 @@ -import asyncio -from typing import Any, Dict, List +from typing import Any, Dict import aisuite import openai import pytest -from pydantic import BaseModel import opik -from opik.config import OPIK_PROJECT_DEFAULT_NAME from opik.integrations.aisuite import track_aisuite from ...testlib import ( ANY_BUT_NONE, ANY_DICT, - ANY_LIST, ANY_STRING, SpanModel, TraceModel, @@ -20,7 +16,6 @@ assert_equal, ) - pytestmark = pytest.mark.usefixtures("ensure_openai_configured") PROJECT_NAME = "aisuite-integration-test" @@ -40,7 +35,9 @@ def _assert_metadata_contains_required_keys(metadata: Dict[str, Any]): assert_dict_has_keys(metadata, REQUIRED_METADATA_KEYS) -def test_aisuite__openai_provider__client_chat_completions_create__happyflow(fake_backend): +def test_aisuite__openai_provider__client_chat_completions_create__happyflow( + fake_backend, +): client = aisuite.Client() wrapped_client = track_aisuite( aisuite_client=client, @@ -102,7 +99,9 @@ def test_aisuite__openai_provider__client_chat_completions_create__happyflow(fak _assert_metadata_contains_required_keys(llm_span_metadata) -def test_aisuite__nonopenai_provider__client_chat_completions_create__happyflow(fake_backend): +def test_aisuite__nonopenai_provider__client_chat_completions_create__happyflow( + fake_backend, +): client = aisuite.Client() wrapped_client = track_aisuite( aisuite_client=client, @@ -185,7 +184,7 @@ def test_aisuite_client_chat_completions_create__create_raises_an_error__span_an metadata={ "created_from": "aisuite", "type": "aisuite_chat", - "model": 'openai:gpt-3.5-turbo', + "model": "openai:gpt-3.5-turbo", }, start_time=ANY_BUT_NONE, end_time=ANY_BUT_NONE, @@ -206,7 +205,7 @@ def test_aisuite_client_chat_completions_create__create_raises_an_error__span_an metadata={ "created_from": "aisuite", "type": "aisuite_chat", - "model": 'openai:gpt-3.5-turbo', + "model": "openai:gpt-3.5-turbo", }, usage=None, start_time=ANY_BUT_NONE, @@ -233,7 +232,6 @@ def test_aisuite_client_chat_completions_create__create_raises_an_error__span_an def test_aisuite_client_chat_completions_create__openai_call_made_in_another_tracked_function__openai_span_attached_to_existing_trace( fake_backend, ): - messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Tell a fact"}, From 8bfe3aaf2994fb6ad68d0b02911600f9f279da26 Mon Sep 17 00:00:00 2001 From: Alexander Barannikov Date: Fri, 13 Dec 2024 12:54:39 +0000 Subject: [PATCH 5/7] RC 4 --- .github/workflows/lib-aisuite-tests.yml | 52 +++++++++++++++++++ .../integrations/aisuite/aisuite_decorator.py | 20 +++---- .../aisuite/requirements.txt | 2 +- .../aisuite/test_aisuite.py | 2 - 4 files changed, 61 insertions(+), 15 deletions(-) create mode 100644 .github/workflows/lib-aisuite-tests.yml diff --git a/.github/workflows/lib-aisuite-tests.yml b/.github/workflows/lib-aisuite-tests.yml new file mode 100644 index 0000000000..cae55d57ec --- /dev/null +++ b/.github/workflows/lib-aisuite-tests.yml @@ -0,0 +1,52 @@ +# Workflow to run AISuite tests +# +# Please read inputs to provide correct values. +# +name: SDK Lib AISuite Tests +run-name: "SDK Lib AISuite Tests ${{ github.ref_name }} by @${{ github.actor }}" +env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + OPENAI_ORG_ID: ${{ secrets.OPENAI_ORG_ID }} + ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} +on: + workflow_call: + +jobs: + tests: + name: AISuite Python ${{matrix.python_version}} + runs-on: ubuntu-latest + defaults: + run: + working-directory: sdks/python + + strategy: + fail-fast: true + matrix: + python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"] + + steps: + - name: Check out code + uses: actions/checkout@v4 + + - name: Setup Python ${{matrix.python_version}} + uses: actions/setup-python@v5 + with: + python-version: ${{matrix.python_version}} + + - name: Install opik + run: pip install . + + - name: Install test tools + run: | + cd ./tests + pip install --no-cache-dir --disable-pip-version-check -r test_requirements.txt + + - name: Install lib + run: | + cd ./tests + pip install --no-cache-dir --disable-pip-version-check -r library_integration/aisuite/requirements.txt + + - name: Run tests + run: | + cd ./tests/library_integration/aisuite/ + python -m pytest -vv . \ No newline at end of file diff --git a/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py b/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py index eb17710fca..6d886de15a 100644 --- a/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py +++ b/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py @@ -1,8 +1,8 @@ import logging from typing import Any, Callable, Dict, List, Optional, Tuple -from aisuite.framework import ChatCompletionResponse -from openai.types.chat import chat_completion +import aisuite.framework as aisuite_chat_completion +from openai.types.chat import chat_completion as openai_chat_completion from opik import dict_utils from opik.decorator import arguments_helpers, base_track_decorator @@ -16,11 +16,7 @@ class AISuiteTrackDecorator(base_track_decorator.BaseTrackDecorator): """ An implementation of BaseTrackDecorator designed specifically for tracking - calls of AISuite's `chat.completion.create` and `chat.completions.parse` functions. - - Besides special processing for input arguments and response content, it - overrides _generators_handler() method to work correctly with - openai.Stream and openai.AsyncStream objects. + calls of AISuite's `chat.completion.create` """ def _start_span_inputs_preprocessor( @@ -93,8 +89,8 @@ def _end_span_inputs_preprocessor( assert isinstance( output, ( - chat_completion.ChatCompletion, # openai - ChatCompletionResponse, # non-openai + openai_chat_completion.ChatCompletion, # openai + aisuite_chat_completion.ChatCompletionResponse, # non-openai ), ) @@ -103,14 +99,14 @@ def _end_span_inputs_preprocessor( model = None # provider == openai - if isinstance(output, chat_completion.ChatCompletion): + if isinstance(output, openai_chat_completion.ChatCompletion): result_dict = output.model_dump(mode="json") output, metadata = dict_utils.split_dict_by_keys(result_dict, ["choices"]) usage = result_dict["usage"] model = result_dict["model"] - # provider == non-openai - elif isinstance(output, ChatCompletionResponse): + # provider != openai + elif isinstance(output, aisuite_chat_completion.ChatCompletionResponse): choices = [] for choice in output.choices: diff --git a/sdks/python/tests/library_integration/aisuite/requirements.txt b/sdks/python/tests/library_integration/aisuite/requirements.txt index aaab7cf072..db2368a09d 100644 --- a/sdks/python/tests/library_integration/aisuite/requirements.txt +++ b/sdks/python/tests/library_integration/aisuite/requirements.txt @@ -1 +1 @@ -aisuite +aisuite[anthropic,openai] diff --git a/sdks/python/tests/library_integration/aisuite/test_aisuite.py b/sdks/python/tests/library_integration/aisuite/test_aisuite.py index 0d8fa2ded4..355c32882a 100644 --- a/sdks/python/tests/library_integration/aisuite/test_aisuite.py +++ b/sdks/python/tests/library_integration/aisuite/test_aisuite.py @@ -113,7 +113,6 @@ def test_aisuite__nonopenai_provider__client_chat_completions_create__happyflow( ] _ = wrapped_client.chat.completions.create( - # model="anthropic:claude-3-5-sonnet-20240620", model="anthropic:claude-3-5-sonnet-latest", messages=messages, max_tokens=10, @@ -166,7 +165,6 @@ def test_aisuite_client_chat_completions_create__create_raises_an_error__span_an project_name=PROJECT_NAME, ) - # with pytest.raises(openai.OpenAIError): with pytest.raises(openai.BadRequestError): _ = wrapped_client.chat.completions.create( messages=None, From 2367adf9ad78e8cac9b04df5045b4bc3f674aa57 Mon Sep 17 00:00:00 2001 From: Alexander Barannikov Date: Fri, 13 Dec 2024 12:59:40 +0000 Subject: [PATCH 6/7] add integration tests to CI --- .github/workflows/lib-integration-tests-runner.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/lib-integration-tests-runner.yml b/.github/workflows/lib-integration-tests-runner.yml index e2c2c1a367..de81e5b3fc 100644 --- a/.github/workflows/lib-integration-tests-runner.yml +++ b/.github/workflows/lib-integration-tests-runner.yml @@ -67,3 +67,9 @@ jobs: uses: ./.github/workflows/lib-anthropic-tests.yml secrets: inherit + aisuite_tests: + needs: [init_environment] + if: contains(fromJSON('["aisuite", "all"]'), needs.init_environment.outputs.LIBS) + uses: ./.github/workflows/lib-aisuite-tests.yml + secrets: inherit + From 835a0637586e071f316a39d69416493700130f13 Mon Sep 17 00:00:00 2001 From: Alexander Barannikov Date: Fri, 13 Dec 2024 13:04:54 +0000 Subject: [PATCH 7/7] raise minimum python version to 3.10 for integration tests --- .github/workflows/lib-aisuite-tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/lib-aisuite-tests.yml b/.github/workflows/lib-aisuite-tests.yml index cae55d57ec..90e469aa45 100644 --- a/.github/workflows/lib-aisuite-tests.yml +++ b/.github/workflows/lib-aisuite-tests.yml @@ -22,7 +22,7 @@ jobs: strategy: fail-fast: true matrix: - python_version: ["3.8", "3.9", "3.10", "3.11", "3.12"] + python_version: ["3.10", "3.11", "3.12"] steps: - name: Check out code