langtrace-python-sdk 1.3.5__py3-none-any.whl → 1.3.7__py3-none-any.whl

tests/openai/cassettes/test_image_generation.yaml

@@ -0,0 +1,97 @@
+interactions:
+- request:
+    body: '{"prompt": "A charming and adorable baby sea otter. This small, fluffy
+      creature is floating gracefully on its back, with its tiny webbed paws folded
+      cutely over its fuzzy belly. It has big, round, innocent eyes that are brimming
+      with youthful curiosity. As it blissfully floats on the calm, sparkling ocean
+      surface under the glow of the golden sunset, it playfully tosses a shiny seashell
+      from one paw to another, showcasing its playful and distinctively otter-like
+      behavior.", "model": "dall-e-3"}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '498'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+      user-agent:
+      - OpenAI/Python 1.23.2
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - 'false'
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 1.23.2
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.11.5
+    method: POST
+    uri: https://api.openai.com/v1/images/generations
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAAA1yR3W7bRhCF7/UUAwFSb0SKf6IoA0Rh2Q4SO2nTRIjVFEUx5A7JjVa71M5SMhPk
+        3QvSaVH3ZjHAmT3znZlvE4BpaQkdiekVhOswTuNgna4WgyDQ4fQK/pgAAHwbX4CppbNkEn+11hxb
+        N72C6S21snSAUDZoj1LXgFoACmOxUAQFFj0wIRjnyEKlDLqhyWiQjqHA8jDUriEwJaH+iYE7W2FJ
+        PuwaAj6iUj8+N8iAUKmuqnooDbpxlJO6hwsVBQlo8cJQGSVIgDmTfZ5BSvU+vBlKWS/Amk4LoJ54
+        AZVUigRcpGugN51rqk5B2VlpWLp+AbUyFyisrBs3eIxEhIDPHiP1iFYreTySZei0IDsK9YChwWLP
+        YCpAYHJjdO60D4/DRKNpQF78x6hV2FedUj04w0xDYG6GhEzIDSkFUjszfkBpFyCV6tjZ56UOaX8Y
+        jKvptDx1BAU1eJbG+tPFP3fsrBqO1zjX8tVyaVAKVIqwla01gsuO/UKZwi+NJf8itTAX9jW5ZWvl
+        GR0tja297Zd0v396elUf7un3V7rhSO9VoO2yY7Le69Oxtfbw8HD+9Onrm93Ddpc0ew4/LuWx9oL1
+        5VLf9Yz3vz3eZrtTtmp+vd9tG7/V9c/s8iiIEi9IvCjYhatZfB1Hs/g6WX2eM73Q1i+0NrdzPg8N
+        oRdkXpDO2ebF3HIpcqmV1DTULpdHrGnZ6nrOByNFniKiIEFeUhWxl6SbzMOsSr00S8R6naVFkK7n
+        fHBS5JhkZYmr1KNUoJdkCXmYhYm3KbNkk66ioiziofVfynCzi0bAcBZfR8HnOR9eRPi/yHkx58PL
+        ELLOt7+oWbR9e7rb7LuPNzfqJny3uzezaPv44cOXlXvb3d29pvf7p9vTZluvg1l8Ox1v/X0C8Ofk
+        ++RvAAAA//8DAI8VAGztAwAA
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 877687ee8e730d81-MRS
+      Connection:
+      - keep-alive
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - application/json
+      Date:
+      - Sat, 20 Apr 2024 16:32:45 GMT
+      Server:
+      - cloudflare
+      Set-Cookie:
+      - __cf_bm=kYFV6abEqmWMPFItBN6IHrGsHEkBMh7G1oaVuNVZR1A-1713630765-1.0.1.1-bQc9xmqCV.rTa7TAHJuI6.CIQ15UhKTn1Zyf4YTgWlYYOJIoLtD48FLCul3eEd3f6v.Gcd7Jo_CKQjYqx..dEA;
+        path=/; expires=Sat, 20-Apr-24 17:02:45 GMT; domain=.api.openai.com; HttpOnly;
+        Secure; SameSite=None
+      - _cfuvid=oNNudHOOYv_4yEJDVbToTqSHaHs.q4LVjXgT0QmIVmc-1713630765448-0.0.1.1-604800000;
+        path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
+      Transfer-Encoding:
+      - chunked
+      access-control-allow-origin:
+      - '*'
+      alt-svc:
+      - h3=":443"; ma=86400
+      openai-organization:
+      - scale3-1
+      openai-processing-ms:
+      - '12231'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=15724800; includeSubDomains
+      x-request-id:
+      - req_fcf0eebdfca1850bb1d76b490ec251ff
+    status:
+      code: 200
+      message: OK
+version: 1

tests/openai/conftest.py

@@ -0,0 +1,45 @@
+"""Unit tests configuration module."""
+
+import pytest
+import os
+from openai import OpenAI, AsyncOpenAI
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry import trace
+from langtrace_python_sdk.instrumentation.openai.instrumentation import (
+    OpenAIInstrumentation,
+)
+
+
+@pytest.fixture(autouse=True)
+def environment():
+    if not os.getenv("OPENAI_API_KEY"):
+        os.environ["OPENAI_API_KEY"] = "test_api_key"
+
+
+@pytest.fixture
+def openai_client():
+    return OpenAI()
+
+
+@pytest.fixture
+def async_openai_client():
+    return AsyncOpenAI()
+
+
+@pytest.fixture(scope="module")
+def vcr_config():
+    return {"filter_headers": ["authorization", "api-key"]}
+
+
+@pytest.fixture(scope="session")
+def exporter():
+    exporter = InMemorySpanExporter()
+    processor = SimpleSpanProcessor(exporter)
+
+    provider = TracerProvider()
+    provider.add_span_processor(processor)
+    trace.set_tracer_provider(provider)
+    OpenAIInstrumentation().instrument()
+    return exporter

tests/openai/test_chat_completion.py

@@ -0,0 +1,142 @@
+import pytest
+import importlib
+import json
+from langtrace_python_sdk.constants.instrumentation.openai import APIS
+
+
+@pytest.mark.vcr()
+def test_chat_completion(exporter, openai_client):
+    llm_model_value = "gpt-4"
+    messages_value = [{"role": "user", "content": "Say this is a test three times"}]
+
+    kwargs = {
+        "model": llm_model_value,
+        "messages": messages_value,
+        "stream": False,
+    }
+
+    openai_client.chat.completions.create(**kwargs)
+    spans = exporter.get_finished_spans()
+    completion_span = spans[-1]
+    assert completion_span.name == "openai.chat.completions.create"
+
+    attributes = completion_span.attributes
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == "OpenAI"
+    assert attributes.get("langtrace.service.type") == "llm"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "openai"
+    )
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("url.full") == "https://api.openai.com/v1/"
+    assert attributes.get("llm.api") == APIS["CHAT_COMPLETION"]["ENDPOINT"]
+    assert attributes.get("llm.model") == "gpt-4-0613"
+    assert attributes.get("llm.prompts") == json.dumps(messages_value)
+    assert attributes.get("llm.stream") is False
+
+    tokens = json.loads(attributes.get("llm.token.counts"))
+    output_tokens = tokens.get("output_tokens")
+    prompt_tokens = tokens.get("input_tokens")
+    total_tokens = tokens.get("total_tokens")
+
+    assert output_tokens and prompt_tokens and total_tokens
+    assert output_tokens + prompt_tokens == total_tokens
+
+
+@pytest.mark.vcr()
+def test_chat_completion_streaming(exporter, openai_client):
+    llm_model_value = "gpt-4"
+    messages_value = [{"role": "user", "content": "Say this is a test three times"}]
+
+    kwargs = {
+        "model": llm_model_value,
+        "messages": messages_value,
+        "stream": True,
+    }
+
+    response = openai_client.chat.completions.create(**kwargs)
+    chunk_count = 0
+    for _ in response:
+        chunk_count += 1
+
+    spans = exporter.get_finished_spans()
+    streaming_span = spans[-1]
+
+    assert streaming_span.name == "openai.chat.completions.create"
+    attributes = streaming_span.attributes
+
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == "OpenAI"
+    assert attributes.get("langtrace.service.type") == "llm"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "openai"
+    )
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("url.full") == "https://api.openai.com/v1/"
+    assert attributes.get("llm.api") == APIS["CHAT_COMPLETION"]["ENDPOINT"]
+    assert attributes.get("llm.model") == "gpt-4-0613"
+    assert attributes.get("llm.prompts") == json.dumps(messages_value)
+    assert attributes.get("llm.stream") is True
+
+    events = streaming_span.events
+    assert len(events) - 2 == chunk_count  # -2 for start and end events
+
+    # check token usage attributes for stream
+    tokens = json.loads(attributes.get("llm.token.counts"))
+
+    output_tokens = tokens.get("output_tokens")
+    prompt_tokens = tokens.get("input_tokens")
+    total_tokens = tokens.get("total_tokens")
+
+    assert output_tokens and prompt_tokens and total_tokens
+    assert output_tokens + prompt_tokens == total_tokens
+
+
+@pytest.mark.vcr()
+@pytest.mark.asyncio()
+async def test_async_chat_completion_streaming(exporter, async_openai_client):
+    llm_model_value = "gpt-4"
+    messages_value = [{"role": "user", "content": "Say this is a test three times"}]
+
+    kwargs = {
+        "model": llm_model_value,
+        "messages": messages_value,
+        "stream": True,
+    }
+
+    response = await async_openai_client.chat.completions.create(**kwargs)
+    chunk_count = 0
+    async for _ in response:
+        chunk_count += 1
+
+    spans = exporter.get_finished_spans()
+    streaming_span = spans[-1]
+
+    assert streaming_span.name == "openai.chat.completions.create"
+    attributes = streaming_span.attributes
+
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == "OpenAI"
+    assert attributes.get("langtrace.service.type") == "llm"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "openai"
+    )
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("url.full") == "https://api.openai.com/v1/"
+    assert attributes.get("llm.api") == APIS["CHAT_COMPLETION"]["ENDPOINT"]
+    assert attributes.get("llm.model") == "gpt-4-0613"
+    assert attributes.get("llm.prompts") == json.dumps(messages_value)
+    assert attributes.get("llm.stream") is True
+
+    events = streaming_span.events
+    assert len(events) - 2 == chunk_count  # -2 for start and end events
+
+    # check token usage attributes for stream
+    tokens = json.loads(attributes.get("llm.token.counts"))
+
+    output_tokens = tokens.get("output_tokens")
+    prompt_tokens = tokens.get("input_tokens")
+    total_tokens = tokens.get("total_tokens")
+
+    assert output_tokens and prompt_tokens and total_tokens
+    assert output_tokens + prompt_tokens == total_tokens

tests/openai/test_image_generation.py

@@ -0,0 +1,77 @@
+import pytest
+import json
+import importlib
+from langtrace_python_sdk.constants.instrumentation.openai import APIS
+
+
+@pytest.mark.vcr()
+def test_image_generation(openai_client, exporter):
+    llm_model_value = "dall-e-3"
+    prompt = "A charming and adorable baby sea otter. This small, fluffy creature is floating gracefully on its back, with its tiny webbed paws folded cutely over its fuzzy belly. It has big, round, innocent eyes that are brimming with youthful curiosity. As it blissfully floats on the calm, sparkling ocean surface under the glow of the golden sunset, it playfully tosses a shiny seashell from one paw to another, showcasing its playful and distinctively otter-like behavior."
+
+    kwargs = {
+        "model": llm_model_value,
+        "prompt": prompt,
+    }
+
+    response = openai_client.images.generate(**kwargs)
+    spans = exporter.get_finished_spans()
+    image_generation_span = spans[-1]
+    assert image_generation_span.name == "openai.images.generate"
+
+    attributes = image_generation_span.attributes
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == "OpenAI"
+    assert attributes.get("langtrace.service.type") == "llm"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "openai"
+    )
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("url.full") == "https://api.openai.com/v1/"
+    assert attributes.get("llm.api") == APIS["IMAGES_GENERATION"]["ENDPOINT"]
+    assert attributes.get("llm.model") == llm_model_value
+    assert attributes.get("llm.prompts") == json.dumps([prompt])
+
+    langtrace_responses = json.loads(attributes.get("llm.responses"))
+    for langtrace_response in langtrace_responses:
+        assert response.data[0].url == langtrace_response.get("url")
+        assert response.data[0].revised_prompt == langtrace_response.get(
+            "revised_prompt"
+        )
+
+
+@pytest.mark.vcr()
+@pytest.mark.asyncio()
+async def test_async_image_generation(async_openai_client, exporter):
+    llm_model_value = "dall-e-3"
+    prompt = "A charming and adorable baby sea otter. This small, fluffy creature is floating gracefully on its back, with its tiny webbed paws folded cutely over its fuzzy belly. It has big, round, innocent eyes that are brimming with youthful curiosity. As it blissfully floats on the calm, sparkling ocean surface under the glow of the golden sunset, it playfully tosses a shiny seashell from one paw to another, showcasing its playful and distinctively otter-like behavior."
+
+    kwargs = {
+        "model": llm_model_value,
+        "prompt": prompt,
+    }
+
+    response = await async_openai_client.images.generate(**kwargs)
+    spans = exporter.get_finished_spans()
+    image_generation_span = spans[-1]
+    assert image_generation_span.name == "openai.images.generate"
+
+    attributes = image_generation_span.attributes
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == "OpenAI"
+    assert attributes.get("langtrace.service.type") == "llm"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "openai"
+    )
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("url.full") == "https://api.openai.com/v1/"
+    assert attributes.get("llm.api") == APIS["IMAGES_GENERATION"]["ENDPOINT"]
+    assert attributes.get("llm.model") == llm_model_value
+    assert attributes.get("llm.prompts") == json.dumps([prompt])
+
+    langtrace_responses = json.loads(attributes.get("llm.responses"))
+    for langtrace_response in langtrace_responses:
+        assert response.data[0].url == langtrace_response.get("url")
+        assert response.data[0].revised_prompt == langtrace_response.get(
+            "revised_prompt"
+        )

tests/pinecone/test_pinecone.py

@@ -0,0 +1,72 @@
+from unittest.mock import MagicMock, patch, call
+from langtrace_python_sdk.instrumentation.pinecone.patch import generic_patch
+from opentelemetry.trace import SpanKind
+import importlib.metadata
+import pinecone
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+from langtrace_python_sdk.constants.instrumentation.pinecone import APIS
+import unittest
+import json
+from tests.utils import common_setup
+
+
+class TestPinecone(unittest.TestCase):
+    data = {
+    "status": "success",
+    "message": "Data upserted successfully",
+    "upserted_ids": [1, 2, 3]
+    }
+    
+    def setUp(self):
+        self.pinecone_mock, self.tracer, self.span = common_setup(self.data, 'pinecone.Index.upsert')
+
+
+    def tearDown(self):
+        self.pinecone_mock.stop()
+
+    def test_pinecone(self):
+    
+       # Arrange
+        version = importlib.metadata.version('pinecone-client')
+        method = "UPSERT"
+        vectors = [[1, 2, 3], [4, 5, 6]]
+
+        # Act
+        wrapped_function = generic_patch(pinecone.Index.upsert, method, version, self.tracer)
+        result = wrapped_function(MagicMock(), MagicMock(), (vectors,), {})
+
+        # Assert
+        self.assertTrue(self.tracer.start_as_current_span.called_once_with("pinecone.data.index", kind=SpanKind.CLIENT))
+        api = APIS[method]
+        service_provider = SERVICE_PROVIDERS["PINECONE"]
+        expected_attributes = {
+            'langtrace.sdk.name': 'langtrace-python-sdk',
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "vectordb",
+            "langtrace.service.version": version,
+            "langtrace.version": "1.0.0",
+            "db.system": "pinecone",
+            "db.operation": api["OPERATION"],
+        }
+        self.assertTrue(
+            self.span.set_attribute.has_calls(
+                [call(key, value) for key, value in expected_attributes.items()], any_order=True
+            )
+        )      
+
+        actual_calls = self.span.set_attribute.call_args_list
+
+        for key, value in expected_attributes.items():
+            self.assertIn(call(key, value), actual_calls)
+            
+        self.assertEqual(self.span.set_status.call_count, 1)
+        self.assertTrue(self.span.set_status.has_calls([call(Status(StatusCode.OK))]))
+
+        expected_result = ['status', 'message', 'upserted_ids']
+        result_keys = json.loads(result).keys()
+        self.assertSetEqual(set(expected_result), set(result_keys), "Keys mismatch")
+
+if __name__ == '__main__':
+    unittest.main()

tests/utils.py

@@ -0,0 +1,21 @@
+from unittest.mock import MagicMock, patch
+import json
+   
+def common_setup(data, method_to_mock=None):
+    if method_to_mock:
+        service_mock = patch(method_to_mock)
+        mock_method = service_mock.start()
+        mock_method.return_value = json.dumps(data)
+    else:
+        service_mock = MagicMock()
+        service_mock.return_value = MagicMock(**data)
+       
+
+    tracer = MagicMock()
+    span = MagicMock()
+
+    context_manager_mock = MagicMock()
+    context_manager_mock.__enter__.return_value = span
+    tracer.start_as_current_span.return_value = context_manager_mock
+
+    return service_mock, tracer, span