langtrace-python-sdk 2.0.6__py3-none-any.whl → 2.0.8__py3-none-any.whl

tests/pinecone/test_pinecone.py

@@ -1,81 +1,91 @@
-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
+import pytest
+import importlib
+import json
 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):
+def create_embedding(openai_client):
+    result = openai_client.embeddings.create(
+        model="text-embedding-ada-002",
+        input="Some random text string goes here",
+        encoding_format="float",
+    )
+    return result.data[0].embedding
 
-        # 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)
+@pytest.mark.vcr()
+def test_upsert(openai_client, pinecone_client, exporter):
+    embedding = create_embedding(openai_client)
+    unique_id = "unique_random_id"
+    data_to_upsert = {
+        "id": unique_id,
+        "values": embedding,
+        "metadata": {"random": "random"},
+    }
+    index = pinecone_client.Index("test-index")
+    index.upsert(vectors=[data_to_upsert], namespace="test-namespace")
+    spans = exporter.get_finished_spans()
+    pinecone_span = spans[-1]
 
-        self.assertEqual(self.span.set_status.call_count, 1)
-        self.assertTrue(self.span.set_status.has_calls([call(Status(StatusCode.OK))]))
+    assert pinecone_span.name == APIS["UPSERT"]["METHOD"]
+    attributes = pinecone_span.attributes
 
-        expected_result = ["status", "message", "upserted_ids"]
-        result_keys = json.loads(result).keys()
-        self.assertSetEqual(set(expected_result), set(result_keys), "Keys mismatch")
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == SERVICE_PROVIDERS["PINECONE"]
+    assert attributes.get("langtrace.service.type") == "vectordb"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "pinecone-client"
+    )
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("db.system") == "pinecone"
+    assert attributes.get("db.operation") == APIS["UPSERT"]["OPERATION"]
 
 
-if __name__ == "__main__":
-    unittest.main()
+@pytest.mark.vcr()
+def test_query(openai_client, pinecone_client, exporter):
+    embedding = create_embedding(openai_client)
+    unique_id = "unique_random_id"
+    data_to_upsert = {
+        "id": unique_id,
+        "values": embedding,
+        "metadata": {"random": "random"},
+    }
+    index = pinecone_client.Index("test-index")
+    index.upsert(vectors=[data_to_upsert], namespace="test-namespace")
+    filter = {"random": "random"}
+    res = index.query(
+        vector=embedding,
+        top_k=3,
+        include_values=True,
+        namespace="test-namespace",
+        include_metadata=True,
+        filter=filter,
+    )
+    spans = exporter.get_finished_spans()
+    query_span = spans[-1]
+    assert query_span.name == APIS["QUERY"]["METHOD"]
+    attributes = query_span.attributes
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == SERVICE_PROVIDERS["PINECONE"]
+    assert attributes.get("langtrace.service.type") == "vectordb"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "pinecone-client"
+    )
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("db.system") == "pinecone"
+    assert attributes.get("db.operation") == APIS["QUERY"]["OPERATION"]
+    assert attributes.get("db.query.top_k") == 3
+    assert attributes.get("db.query.namespace") == "test-namespace"
+    assert attributes.get("db.query.include_values") is True
+    assert attributes.get("db.query.include_metadata") is True
+    assert attributes.get("db.query.usage.read_units") == 6
+    assert json.loads(attributes.get("db.query.filter")) == filter
+    res_matches = res.matches
+    events = query_span.events
+    assert len(res_matches) == len(events)
+    for idx, event in enumerate(events):
+        assert event.name == "db.query.match"
+        attrs = event.attributes
+        assert attrs.get("db.query.match.id") == res_matches[idx].id
+        assert attrs.get("db.query.match.score") == res_matches[idx].score

tests/qdrant/conftest.py

@@ -0,0 +1,19 @@
+"""Unit tests configuration module."""
+
+import pytest
+
+
+from langtrace_python_sdk.instrumentation.qdrant.instrumentation import (
+    QdrantInstrumentation,
+)
+from qdrant_client import QdrantClient
+
+
+@pytest.fixture
+def qdrant_client():
+    return QdrantClient(":memory:")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def instrument():
+    QdrantInstrumentation().instrument()

tests/qdrant/test_qdrant.py

@@ -0,0 +1,92 @@
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+from qdrant_client.models import PointStruct, Distance, VectorParams
+import importlib
+
+
+COLLECTION_NAME = "test_collection"
+EMBEDDING_DIM = 4
+
+
+def test_qdrant_upsert(qdrant_client, exporter):
+    qdrant_client.create_collection(
+        collection_name=COLLECTION_NAME,
+        vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.DOT),
+    )
+
+    qdrant_client.upsert(
+        collection_name=COLLECTION_NAME,
+        wait=True,
+        points=[
+            PointStruct(
+                id=1, vector=[0.05, 0.61, 0.76, 0.74], payload={"city": "Berlin"}
+            ),
+            PointStruct(
+                id=2, vector=[0.19, 0.81, 0.75, 0.11], payload={"city": "London"}
+            ),
+            PointStruct(
+                id=3, vector=[0.36, 0.55, 0.47, 0.94], payload={"city": "Moscow"}
+            ),
+            PointStruct(
+                id=4, vector=[0.18, 0.01, 0.85, 0.80], payload={"city": "New York"}
+            ),
+            PointStruct(
+                id=5, vector=[0.24, 0.18, 0.22, 0.44], payload={"city": "Beijing"}
+            ),
+            PointStruct(
+                id=6, vector=[0.35, 0.08, 0.11, 0.44], payload={"city": "Mumbai"}
+            ),
+        ],
+    )
+    spans = exporter.get_finished_spans()
+    upsert_span = spans[0]
+    attributes = upsert_span.attributes
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == SERVICE_PROVIDERS["QDRANT"]
+    assert attributes.get("langtrace.service.type") == "vectordb"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "qdrant_client"
+    )
+    assert attributes.get("db.system") == SERVICE_PROVIDERS["QDRANT"].lower()
+    assert attributes.get("db.operation") == "upsert"
+    assert attributes.get("db.upsert.points_count") == 6
+    assert attributes.get("db.collection.name") == COLLECTION_NAME
+
+
+def test_qdrant_search(qdrant_client, exporter):
+    qdrant_client.create_collection(
+        collection_name=COLLECTION_NAME,
+        vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.DOT),
+    )
+
+    qdrant_client.upsert(
+        collection_name=COLLECTION_NAME,
+        wait=True,
+        points=[
+            PointStruct(
+                id=1, vector=[0.05, 0.61, 0.76, 0.74], payload={"city": "Berlin"}
+            ),
+            PointStruct(
+                id=2, vector=[0.19, 0.81, 0.75, 0.11], payload={"city": "London"}
+            ),
+            PointStruct(
+                id=3, vector=[0.36, 0.55, 0.47, 0.94], payload={"city": "Moscow"}
+            ),
+            PointStruct(
+                id=4, vector=[0.18, 0.01, 0.85, 0.80], payload={"city": "New York"}
+            ),
+            PointStruct(
+                id=5, vector=[0.24, 0.18, 0.22, 0.44], payload={"city": "Beijing"}
+            ),
+            PointStruct(
+                id=6, vector=[0.35, 0.08, 0.11, 0.44], payload={"city": "Mumbai"}
+            ),
+        ],
+    )
+    qdrant_client.search(
+        collection_name=COLLECTION_NAME, query_vector=[0.2, 0.1, 0.9, 0.7], limit=3
+    )
+
+    spans = exporter.get_finished_spans()
+    search_span = spans[-1]
+    attributes = search_span.attributes
+    assert attributes.get("db.operation") == "search"