langtrace-python-sdk 1.0.9__py3-none-any.whl

instrumentation/pinecone/patch.py

@@ -0,0 +1,45 @@
+import json
+
+from langtrace.trace_attributes import DatabaseSpanAttributes
+from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace.status import Status, StatusCode
+
+from instrumentation.constants import SERVICE_PROVIDERS
+from instrumentation.pinecone.lib.apis import APIS
+
+
+def generic_patch(original_method, method, version, tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        api = APIS[method]
+        service_provider = SERVICE_PROVIDERS['PINECONE']
+        span_attributes = {
+            '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'],
+        }
+
+        attributes = DatabaseSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(api['METHOD'], kind=SpanKind.CLIENT) 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 = original_method(instance, *args, **kwargs)
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as e:
+                # Record the exception in the span
+                span.record_exception(e)
+
+                # Set the span status to indicate an error
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Reraise the exception to ensure it's not swallowed
+                raise
+
+    return traced_method

instrumentation/setup.py

@@ -0,0 +1,50 @@
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import (ConsoleSpanExporter,
+                                            SimpleSpanProcessor)
+
+from instrumentation.chroma.instrumentation import ChromaInstrumentation
+from instrumentation.langchain.instrumentation import LangchainInstrumentation
+from instrumentation.langchain_community.instrumentation import \
+    LangchainCommunityInstrumentation
+from instrumentation.langchain_core.instrumentation import \
+    LangchainCoreInstrumentation
+from instrumentation.llamaindex.instrumentation import \
+    LlamaindexInstrumentation
+from instrumentation.openai.instrumentation import OpenAIInstrumentation
+from instrumentation.pinecone.instrumentation import PineconeInstrumentation
+
+
+def setup_instrumentation():
+
+    # Set up OpenTelemetry tracing
+    tracer_provider = TracerProvider()
+
+    # Use the ConsoleSpanExporter to print traces to the console
+    console_exporter = ConsoleSpanExporter()
+    tracer_provider.add_span_processor(SimpleSpanProcessor(console_exporter))
+
+    # Initialize tracer
+    trace.set_tracer_provider(tracer_provider)
+
+    # Initialize and enable your custom OpenAI instrumentation
+    # Create an instance of OpenAIInstrumentation
+    openai_instrumentation = OpenAIInstrumentation()
+    pinecone_instrumentation = PineconeInstrumentation()
+    llamaindex_instrumentation = LlamaindexInstrumentation()
+    chroma_instrumentation = ChromaInstrumentation()
+    langchain_instrumentation = LangchainInstrumentation()
+    langchain_core_instrumentation = LangchainCoreInstrumentation()
+    langchain_community_instrumentation = LangchainCommunityInstrumentation()
+
+    # Call the instrument method with some arguments
+    openai_instrumentation.instrument()
+    pinecone_instrumentation.instrument()
+    llamaindex_instrumentation.instrument()
+    chroma_instrumentation.instrument()
+    langchain_instrumentation.instrument()
+    langchain_core_instrumentation.instrument()
+    langchain_community_instrumentation.instrument()
+
+    print("setup complete")

instrumentation/utils.py

@@ -0,0 +1,27 @@
+from tiktoken import TiktokenEncoding, get_encoding
+from .constants import TIKTOKEN_MODEL_MAPPING, OPENAI_COST_TABLE
+
+def estimate_tokens(prompt: str) -> int:
+    if prompt and len(prompt) > 0:
+        # Simplified token estimation: count the words.
+        return len(prompt.split())
+    return 0
+
+def estimate_tokens_using_tiktoken(prompt: str, model: TiktokenEncoding) -> int:
+    encoding = get_encoding(model)
+    tokens = encoding.encode(prompt)
+    return len(tokens)
+
+def calculate_prompt_tokens(prompt_content: str, model: str) -> int:
+    try:
+        tiktoken_model = TIKTOKEN_MODEL_MAPPING[model]
+        return estimate_tokens_using_tiktoken(prompt_content, tiktoken_model)
+    except KeyError:
+        return estimate_tokens(prompt_content)  # Fallback method
+
+def calculate_price_from_usage(model: str, usage: dict) -> float:
+    cost_table = OPENAI_COST_TABLE.get(model)
+    if cost_table:
+        return ((cost_table['input'] * usage['prompt_tokens'] +
+                 cost_table['output'] * usage['completion_tokens']) / 1000)
+    return 0

instrumentation/with_root_span.py

@@ -0,0 +1,28 @@
+import asyncio
+from functools import wraps
+
+from opentelemetry import trace
+from opentelemetry.trace import SpanKind
+
+
+def with_langtrace_root_span(name="LangtraceRootSpan", kind=SpanKind.INTERNAL):
+    def decorator(func):
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = trace.get_tracer(__name__)
+            operation_name = name if name else func.__name__
+            with tracer.start_as_current_span(operation_name, kind=kind):
+                return func(*args, **kwargs)
+
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = trace.get_tracer(__name__)
+            operation_name = name if name else func.__name__
+            with tracer.start_as_current_span(operation_name, kind=kind):
+                return await func(*args, **kwargs)
+
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+    return decorator