diff --git a/.github/workflows/lib-aisuite-tests.yml b/.github/workflows/lib-aisuite-tests.yml new file mode 100644 index 0000000000..90e469aa45 --- /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.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/.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 + 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..6d886de15a --- /dev/null +++ b/sdks/python/src/opik/integrations/aisuite/aisuite_decorator.py @@ -0,0 +1,136 @@ +import logging +from typing import Any, Callable, Dict, List, Optional, Tuple + +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 + + +LOGGER = logging.getLogger(__name__) + +KWARGS_KEYS_TO_LOG_AS_INPUTS = ["messages"] + + +class AISuiteTrackDecorator(base_track_decorator.BaseTrackDecorator): + """ + An implementation of BaseTrackDecorator designed specifically for tracking + calls of AISuite's `chat.completion.create` + """ + + 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"] + + model, provider = self._get_provider_info(func, **kwargs) + + result = arguments_helpers.StartSpanParameters( + name=name, + input=input, + type=track_options.type, + tags=tags, + metadata=metadata, + project_name=track_options.project_name, + model=model, + provider=provider, + ) + + return result + + 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) + + if model is not None and ":" in model: + provider, model = model.split(":", 1) + + if provider != "openai": + return model, provider + + 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 + + return model, provider + + def _end_span_inputs_preprocessor( + self, output: Any, capture_output: bool + ) -> arguments_helpers.EndSpanParameters: + assert isinstance( + output, + ( + openai_chat_completion.ChatCompletion, # openai + aisuite_chat_completion.ChatCompletionResponse, # non-openai + ), + ) + + metadata = None + usage = None + model = None + + # provider == openai + 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 != openai + elif isinstance(output, aisuite_chat_completion.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, + ) + + 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..db2368a09d --- /dev/null +++ b/sdks/python/tests/library_integration/aisuite/requirements.txt @@ -0,0 +1 @@ +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 new file mode 100644 index 0000000000..355c32882a --- /dev/null +++ b/sdks/python/tests/library_integration/aisuite/test_aisuite.py @@ -0,0 +1,309 @@ +from typing import Any, Dict + +import aisuite +import openai +import pytest + +import opik +from opik.integrations.aisuite import track_aisuite +from ...testlib import ( + ANY_BUT_NONE, + ANY_DICT, + ANY_STRING, + SpanModel, + TraceModel, + assert_dict_has_keys, + assert_equal, +) + +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 = [ + "usage", + "model", + "max_tokens", + "created_from", + "type", + "id", + "created", + "object", + ] + assert_dict_has_keys(metadata, REQUIRED_METADATA_KEYS) + + +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, + ) + 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=["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={ + "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].metadata + _assert_metadata_contains_required_keys(llm_span_metadata) + + +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-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, + project_name=PROJECT_NAME, + ) + + with pytest.raises(openai.BadRequestError): + _ = wrapped_client.chat.completions.create( + messages=None, + model="openai:gpt-3.5-turbo", + ) + + opik.flush_tracker() + + EXPECTED_TRACE_TREE = TraceModel( + id=ANY_BUT_NONE, + name="chat_completion_create", + input={"messages": None}, + output=None, + tags=["aisuite"], + metadata={ + "created_from": "aisuite", + "type": "aisuite_chat", + "model": "openai:gpt-3.5-turbo", + }, + start_time=ANY_BUT_NONE, + end_time=ANY_BUT_NONE, + project_name=PROJECT_NAME, + 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=["aisuite"], + metadata={ + "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=PROJECT_NAME, + model=ANY_STRING(startswith="gpt-3.5-turbo"), + 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) + + +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"}, + ] + + @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=f"{PROJECT_NAME}-nested-level", + ) + + _ = wrapped_client.chat.completions.create( + model="openai: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=["aisuite"], + 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)