feat: adding support for basic pgvector tracing

Author: DSdatsme

pyproject.toml

@@ -57,6 +57,8 @@ dev = [
   "google-cloud-aiplatform",
   "mistralai",
   "embedchain",
+  "psycopg",
+  "pgvector"
 ]
 
 test = ["pytest", "pytest-vcr", "pytest-asyncio"]

src/examples/pgvector_example/basic.py

@@ -0,0 +1,41 @@
+from openai import OpenAI
+from pgvector.psycopg import register_vector
+import psycopg
+from langtrace_python_sdk import langtrace
+
+langtrace.init(write_spans_to_console=True)
+
+conn = psycopg.connect(dbname='postgres', autocommit=True, password="mypw", user="postgres", host="localhost")
+conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
+register_vector(conn)
+
+client = OpenAI()
+
+
+def setup_db():
+    conn.execute('DROP TABLE IF EXISTS documents')
+    conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(1536))')
+
+    input = [
+        'The dog is barking',
+        'The cat is purring',
+        'The bear is growling'
+    ]
+
+    response = client.embeddings.create(input=input, model='text-embedding-3-small')
+    embeddings = [v.embedding for v in response.data]
+
+    for content, embedding in zip(input, embeddings):
+        conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding))
+
+
+def basic_pgvector():
+    setup_db()
+    document_id = 1
+    neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall()
+    for neighbor in neighbors:
+        print(neighbor[0])
+
+    print("neighbors", neighbors)
+    return neighbors
+

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -32,6 +32,7 @@
     "MISTRAL": "Mistral",
     "EMBEDCHAIN": "Embedchain",
     "AUTOGEN": "Autogen",
+    "PGVECTOR": "PgVector",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -19,6 +19,7 @@
 from .gemini import GeminiInstrumentation
 from .mistral import MistralInstrumentation
 from .embedchain import EmbedchainInstrumentation
+from .pgvector import PgVectorInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -42,4 +43,5 @@
     "VertexAIInstrumentation",
     "GeminiInstrumentation",
     "MistralInstrumentation",
+    "PgVectorInstrumentation",
 ]

src/langtrace_python_sdk/instrumentation/pgvector/__init__.py

@@ -0,0 +1,5 @@
+from .instrumentation import PgVectorInstrumentation
+
+__all__ = [
+    "PgVectorInstrumentation",
+]

src/langtrace_python_sdk/instrumentation/pgvector/instrumentation.py

@@ -0,0 +1,58 @@
+"""
+Copyright (c) 2024 Scale3 Labs
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import importlib.metadata
+import logging
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from langtrace_python_sdk.constants.instrumentation.pgvector import APIS
+from langtrace_python_sdk.instrumentation.pgvector.patch import (
+    generic_patch,
+)
+
+logging.basicConfig(level=logging.DEBUG)  # Set to DEBUG for detailed logging
+
+
+class PgVectorInstrumentation(BaseInstrumentor):
+    """
+    The PgVectorInstrumentation class represents the instrumentation for the Postgres Vector.
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["pgvector >= 0.3.3", "trace-attributes >= 7.1.0"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version("pgvector")
+
+        print("Instrumenting PgVector")
+        for api_name, api_config in APIS.items():
+            wrap_function_wrapper(
+                api_config["MODULE"],
+                api_config["METHOD"],
+                generic_patch(api_name, version, tracer),
+            )
+
+    def _instrument_module(self, module_name):
+        pass
+
+    def _uninstrument(self, **kwargs):
+        pass

src/langtrace_python_sdk/instrumentation/pgvector/patch.py

@@ -0,0 +1,90 @@
+"""
+Copyright (c) 2024 Scale3 Labs
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import re
+from importlib_metadata import version as v
+from langtrace.trace_attributes import DatabaseSpanAttributes
+from opentelemetry import baggage, trace
+from opentelemetry.trace import SpanKind,Tracer
+from opentelemetry.trace.propagation import set_span_in_context
+from opentelemetry.trace.status import Status, StatusCode
+
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from langtrace_python_sdk.constants.instrumentation.common import (
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    SERVICE_PROVIDERS,
+)
+from langtrace_python_sdk.constants.instrumentation.pgvector import APIS
+from langtrace_python_sdk.utils.llm import get_span_name
+from langtrace_python_sdk.utils.misc import extract_input_params, to_iso_format
+
+
+# Utility function to extract table name
+def extract_table_name(query):
+    # This regex assumes basic SQL queries like SELECT, INSERT INTO, UPDATE, DELETE FROM
+    match = re.search(r'(from|into|update|delete\s+from)\s+(\w+)', query, re.IGNORECASE)
+    if match:
+        return match.group(2)
+    return None
+
+
+def generic_patch(method_name, version, tracer: Tracer):
+
+    def traced_method(wrapped, instance, args, kwargs): 
+        query = args[0]
+        params = args[1] if len(args) > 1 else None 
+        api = APIS[method_name]
+        service_provider = SERVICE_PROVIDERS["PGVECTOR"]
+        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+
+        span_attributes = {
+            "langtrace.sdk.name": "langtrace-python-sdk",
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "vectordb",
+            "langtrace.service.version": version,
+            "langtrace.version": v(LANGTRACE_SDK_NAME),
+            "db.system": "pgvector",
+            "db.operation": api["OPERATION"],
+            "db.collection.name": extract_table_name(query),
+            "db.namespace": instance.connection.info.dbname,
+            "db.query": str(query),
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        attributes = DatabaseSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=get_span_name(method_name),
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
+            for field, value in attributes.model_dump(by_alias=True).items():
+                if value is not None:
+                    span.set_attribute(field, value)
+            try:
+                # Attempt to call the original method
+                result = wrapped(*args, **kwargs)
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as err:
+                # Record the exception in the span
+                span.record_exception(err)
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(err)))
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method

src/langtrace_python_sdk/langtrace.py

@@ -50,6 +50,7 @@
     MistralInstrumentation,
     OllamaInstrumentor,
     OpenAIInstrumentation,
+    PgVectorInstrumentation,
     PineconeInstrumentation,
     QdrantInstrumentation,
     AutogenInstrumentation,
@@ -148,6 +149,7 @@ def init(
         "google-generativeai": GeminiInstrumentation(),
         "mistralai": MistralInstrumentation(),
         "autogen": AutogenInstrumentation(),
+        "pgvector": PgVectorInstrumentation(),
     }
 
     init_instrumentations(disable_instrumentations, all_instrumentations)

src/langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "2.3.20"
+__version__ = "2.4.0"