[OPIK-569] [SDK] Implement integration with aisuite

.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 .

.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
+

sdks/python/src/opik/integrations/aisuite/__init__.py

@@ -0,0 +1,4 @@
+from .opik_tracker import track_aisuite
+
+
+__all__ = ["track_aisuite"]

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

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

sdks/python/tests/library_integration/aisuite/requirements.txt

@@ -0,0 +1 @@
+aisuite[anthropic,openai]