langtrace-python-sdk 1.0.9__tar.gz → 1.0.10__tar.gz

{langtrace-python-sdk-1.0.9/src/langtrace_python_sdk.egg-info → langtrace-python-sdk-1.0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langtrace-python-sdk
-Version: 1.0.9
+Version: 1.0.10
 Summary: LangTrace - Python SDK
 Home-page: https://github.com/Scale3-Labs/langtrace-python-sdk
 Author: Ali Waleed

langtrace-python-sdk-1.0.10/README.md

@@ -0,0 +1,22 @@
+# langtrace-python-sdk
+
+export PYTHONPATH="/Users/karthikkalyanaraman/work/langtrace/python-sdk:$PYTHONPATH"
+
+## Steps to run
+
+1. From your root directory, create a virtualenv for installing your dependencies
+    ```
+    python -m venv pysdk
+    ```
+2. Activate the virtualenv
+    ```
+    source pysdk/bin/activate
+    ```
+3. Install the dependencies
+   ```
+    pip install -r requirements.txt
+    ```
+4. Run the example and see the traces on the terminal
+   ```
+   python src/entrypoint.py
+   ```

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/setup.py

@@ -4,7 +4,7 @@ with open('requirements.txt') as f:
 
 setup(
     name='langtrace-python-sdk',  # Choose a unique name for PyPI
-    version='1.0.9',
+    version='1.0.10',
     author='Ali Waleed',
     author_email='ali@scale3labs.com',
     description='LangTrace - Python SDK',

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/examples/chroma_example/basic.py

@@ -2,7 +2,7 @@ import chromadb
 from chromadb.utils import embedding_functions
 from dotenv import find_dotenv, load_dotenv
 
-from instrumentation.setup import setup_instrumentation
+from examples.setup import setup_instrumentation
 from instrumentation.with_root_span import with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/examples/langchain_example/basic.py

@@ -7,7 +7,7 @@ from langchain_core.prompts.chat import ChatPromptTemplate
 from langchain_core.runnables import RunnablePassthrough
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 
-from instrumentation.setup import setup_instrumentation
+from examples.setup import setup_instrumentation
 from instrumentation.with_root_span import with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/examples/langchain_example/tool.py

@@ -6,7 +6,7 @@ from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain_core.tools import Tool
 from langchain_openai import ChatOpenAI
 
-from instrumentation.setup import setup_instrumentation
+from examples.setup import setup_instrumentation
 from instrumentation.with_root_span import with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())
@@ -45,7 +45,7 @@ async def aget_prime(n: int, primes: dict = primes) -> str:
     return str(primes.get(int(n)))
 
 
-# @with_langtrace_root_span()
+@with_langtrace_root_span()
 def tool_example():
 
     tools = [

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/examples/llamaindex_example/basic.py

@@ -1,7 +1,7 @@
 from dotenv import find_dotenv, load_dotenv
 from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
 
-from instrumentation.setup import setup_instrumentation
+from examples.setup import setup_instrumentation
 from instrumentation.with_root_span import with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/examples/openai/chat_completion.py

@@ -1,7 +1,7 @@
 from dotenv import find_dotenv, load_dotenv
 from openai import OpenAI
 
-from instrumentation.setup import setup_instrumentation
+from examples.setup import setup_instrumentation
 from instrumentation.with_root_span import with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())
@@ -13,14 +13,27 @@ client = OpenAI()
 
 @with_langtrace_root_span()
 def chat_completion():
-    stream = client.chat.completions.create(
+    response = client.chat.completions.create(
         model="gpt-4",
         messages=[{"role": "user", "content": "Say this is a test three times"}],
-        stream=False,
+        stream=True,
     )
+    # print(stream)
     # stream = client.chat.completions.create(
     #     model="gpt-4",
     #     messages=[{"role": "user", "content": "Say this is a test three times"}, {"role": "assistant", "content": "This is a test. This is a test. This is a test"},
     #               {"role": "user", "content": "Say this is a mock 4 times"}],
     #     stream=False,
     # )
+
+    result = []
+    for chunk in response:
+        if chunk.choices[0].delta.function_call is not None:
+            content = [
+                choice.delta.function_call.arguments if choice.delta.function_call and
+                choice.delta.function_call.arguments else ""
+                for choice in chunk.choices]
+            result.append(
+                content[0] if len(content) > 0 else "")
+
+    print("".join(result))

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/examples/openai/embeddings_create.py

@@ -1,7 +1,7 @@
 from dotenv import find_dotenv, load_dotenv
 from openai import OpenAI
 
-from instrumentation.setup import setup_instrumentation
+from examples.setup import setup_instrumentation
 from instrumentation.with_root_span import with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/examples/openai/images_generate.py

@@ -1,7 +1,7 @@
 from dotenv import find_dotenv, load_dotenv
 from openai import OpenAI
 
-from instrumentation.setup import setup_instrumentation
+from examples.setup import setup_instrumentation
 from instrumentation.with_root_span import with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/examples/pinecone_example/basic.py

@@ -2,7 +2,7 @@ from dotenv import find_dotenv, load_dotenv
 from openai import OpenAI
 from pinecone import Pinecone
 
-from instrumentation.setup import setup_instrumentation
+from examples.setup import setup_instrumentation
 from instrumentation.with_root_span import with_langtrace_root_span
 
 _ = load_dotenv(find_dotenv())

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/instrumentation/chroma/instrumentation.py

@@ -1,3 +1,6 @@
+"""
+Instrumentation for ChromaDB
+"""
 import importlib.metadata
 from typing import Collection
 
@@ -5,11 +8,14 @@ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper
 
-from instrumentation.chroma.lib.apis import APIS
+from instrumentation.chroma.apis import APIS
 from instrumentation.chroma.patch import collection_patch
 
 
 class ChromaInstrumentation(BaseInstrumentor):
+    """
+    The ChromaInstrumentation class represents the ChromaDB instrumentation
+    """
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return ["chromadb >= 0.4.23"]
@@ -19,7 +25,7 @@ class ChromaInstrumentation(BaseInstrumentor):
         tracer = get_tracer(__name__, "", tracer_provider)
         version = importlib.metadata.version('chromadb')
 
-        for operation, details in APIS.items():
+        for operation, _ in APIS.items():
             wrap_function_wrapper(
                 'chromadb.api.models.Collection',
                 f'Collection.{operation.lower()}',
@@ -27,7 +33,7 @@ class ChromaInstrumentation(BaseInstrumentor):
             )
 
     def _instrument_module(self, module_name):
-        print(module_name)
+        pass
 
     def _uninstrument(self, **kwargs):
-        print(kwargs)
+        pass

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/instrumentation/chroma/patch.py

@@ -1,12 +1,18 @@
+"""
+This module contains the patching logic for the Chroma client.
+"""
 from langtrace.trace_attributes import DatabaseSpanAttributes
-from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
+from instrumentation.chroma.apis import APIS
 from instrumentation.constants import SERVICE_PROVIDERS
-from instrumentation.chroma.lib.apis import APIS
 
 
 def collection_patch(method, version, tracer):
+    """
+    A generic patch method that wraps a function with a span
+    """
     def traced_method(wrapped, instance, args, kwargs):
         api = APIS[method]
         service_provider = SERVICE_PROVIDERS['CHROMA']

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/instrumentation/constants.py

@@ -1,16 +1,11 @@
-
-TRACE_NAMESPACES = {
-    "OPENAI": "Langtrace OpenAI SDK",
-    "LANGCHAIN": "Langtrace Langchain SDK",
-    "PINECONE": "Langtrace Pinecone SDK",
-    "LLAMAINDEX": "Langtrace LlamaIndex SDK",
-    "CHROMA": "Langtrace Chroma SDK",
-}
-
+"""
+This file contains the constants used in the project.
+"""
 SERVICE_PROVIDERS = {
     "OPENAI": "OpenAI",
     "AZURE": "Azure",
     "LANGCHAIN": "Langchain",
+    "LANGCHAIN_CORE": "Langchain Core",
     "LANGCHAIN_COMMUNITY": "Langchain Community",
     "PINECONE": "Pinecone",
     "LLAMAINDEX": "LlamaIndex",

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/instrumentation/langchain_core/instrumentation.py

@@ -71,15 +71,21 @@ class LangchainCoreInstrumentation(BaseInstrumentor):
         version = importlib.metadata.version('langchain-core')
 
         exclude_methods = ['get_name', 'get_output_schema',
-                           'get_input_schema', 'get_graph', 'to_json']
-        exclude_classes = ['BaseChatPromptTemplate']
+                           'get_input_schema', 'get_graph', 'to_json',
+                           'to_json_not_implemented', 'bind',  'dict',
+                           'format', 'format_messages', 'format_prompt']
+        exclude_classes = ['BaseChatPromptTemplate', 'Runnable', 'RunnableBinding',
+                           'RunnableBindingBase', 'RunnableEach', 'RunnableEachBase',
+                           'RunnableGenerator', 'RunnablePick', 'RunnableMap',
+                           'RunnableSerializable']
+
         modules_to_patch = [
             ('langchain_core.retrievers', 'retriever',
              generic_patch, True, True),
-            ('langchain_core.prompts.chat', 'chatprompt',
-             generic_patch, True, False),
+            ('langchain_core.prompts.chat', 'prompt',
+             generic_patch, True, True),
             ('langchain_core.runnables.base',
-             'runnableparallel', runnable_patch, True, True),
+             'runnable', runnable_patch, True, True),
             ('langchain_core.runnables.passthrough',
              'runnablepassthrough', runnable_patch, True, True),
             ('langchain_core.output_parsers.string',

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/instrumentation/langchain_core/patch.py

@@ -22,7 +22,7 @@ def generic_patch(method_name, task, tracer, version, trace_output=True, trace_i
     """
 
     def traced_method(wrapped, instance, args, kwargs):
-        service_provider = SERVICE_PROVIDERS['LANGCHAIN']
+        service_provider = SERVICE_PROVIDERS['LANGCHAIN_CORE']
         span_attributes = {
             'langtrace.service.name': service_provider,
             'langtrace.service.type': 'framework',
@@ -32,7 +32,18 @@ def generic_patch(method_name, task, tracer, version, trace_output=True, trace_i
         }
 
         if len(args) > 0 and trace_input:
-            span_attributes['langchain.inputs'] = to_json_string(args)
+            inputs = {}
+            for arg in args:
+                if isinstance(arg, dict):
+                    for key, value in arg.items():
+                        if isinstance(value, list):
+                            for item in value:
+                                inputs[key] = item.__class__.__name__
+                        elif isinstance(value, str):
+                            inputs[key] = value
+                elif isinstance(arg, str):
+                    inputs['input'] = arg
+            span_attributes['langchain.inputs'] = to_json_string(inputs)
 
         attributes = FrameworkSpanAttributes(**span_attributes)
 
@@ -73,7 +84,7 @@ def runnable_patch(method_name, task, tracer, version, trace_output=True, trace_
     trace_input: Whether to trace the input of the patched methods.
     """
     def traced_method(wrapped, instance, args, kwargs):
-        service_provider = SERVICE_PROVIDERS['LANGCHAIN']
+        service_provider = SERVICE_PROVIDERS['LANGCHAIN_CORE']
         span_attributes = {
             'langtrace.service.name': service_provider,
             'langtrace.service.type': 'framework',
@@ -84,12 +95,17 @@ def runnable_patch(method_name, task, tracer, version, trace_output=True, trace_
 
         if trace_input:
             inputs = {}
-            args_list = []
             if len(args) > 0:
-                for value in args:
-                    if isinstance(value, str):
-                        args_list.append(value)
-            inputs['args'] = args_list
+                for arg in args:
+                    if isinstance(arg, dict):
+                        for key, value in arg.items():
+                            if isinstance(value, list):
+                                for item in value:
+                                    inputs[key] = item.__class__.__name__
+                            elif isinstance(value, str):
+                                inputs[key] = value
+                    elif isinstance(arg, str):
+                        inputs['input'] = arg
 
             for field, value in instance.steps.items() if hasattr(instance, "steps") and \
                     isinstance(instance.steps, dict) else {}:

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/instrumentation/llamaindex/instrumentation.py

@@ -1,3 +1,6 @@
+"""
+The LlamaindexInstrumentation class represents the LlamaIndex instrumentation
+"""
 import importlib.metadata
 from typing import Collection
 
@@ -8,12 +11,11 @@ from wrapt import wrap_function_wrapper
 
 from instrumentation.llamaindex.patch import generic_patch
 
-MODULES = [
-    "llama_index.core.query_pipeline.query",
-]
-
 
 class LlamaindexInstrumentation(BaseInstrumentor):
+    """
+    The LlamaindexInstrumentation class represents the LlamaIndex instrumentation
+    """
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return ["llama-index >= 0.10.0"]

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/instrumentation/llamaindex/patch.py

@@ -1,11 +1,16 @@
+"""
+This module contains a generic patch method that wraps a function with a span.
+"""
 from langtrace.trace_attributes import FrameworkSpanAttributes
-from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
 from instrumentation.constants import SERVICE_PROVIDERS
 
 
 def generic_patch(method, task, tracer, version):
+    """
+    A generic patch method that wraps a function with a span"""
     def traced_method(wrapped, instance, args, kwargs):
         service_provider = SERVICE_PROVIDERS['LLAMAINDEX']
         span_attributes = {

{langtrace-python-sdk-1.0.9/src/instrumentation/openai/lib → langtrace-python-sdk-1.0.10/src/instrumentation/openai}/apis.py

@@ -1,3 +1,6 @@
+"""
+APIs to instrument OpenAI.
+"""
 from langtrace.trace_attributes import OpenAIMethods
 
 APIS = {

{langtrace-python-sdk-1.0.9/src/instrumentation/openai/lib → langtrace-python-sdk-1.0.10/src/instrumentation/openai}/constants.py

@@ -1,3 +1,6 @@
+"""
+Constants for OpenAI API"""
+
 OPENAI_COST_TABLE = {
     "gpt-4-0125-preview": {
         "input": 0.01,
@@ -28,3 +31,13 @@ OPENAI_COST_TABLE = {
         "output": 0.002,
     },
 }
+
+# TODO: Add more models
+# https://github.com/dqbd/tiktoken/blob/74c147e19584a3a1acea0c8e0da4d39415cd33e0/wasm/src/lib.rs#L328
+TIKTOKEN_MODEL_MAPPING = {
+    "gpt-4": "cl100k_base",
+    "gpt-4-32k": "cl100k_base",
+    "gpt-4-0125-preview": "cl100k_base",
+    "gpt-4-1106-preview": "cl100k_base",
+    "gpt-4-1106-vision-preview": "cl100k_base",
+}

langtrace-python-sdk-1.0.10/src/instrumentation/openai/patch.py

@@ -0,0 +1,268 @@
+"""
+This module contains the patching logic for the OpenAI library."""
+import json
+
+from langtrace.trace_attributes import Event, LLMSpanAttributes
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+
+from instrumentation.constants import SERVICE_PROVIDERS
+from instrumentation.openai.apis import APIS
+from instrumentation.openai.token_estimation import (calculate_prompt_tokens,
+                                                     estimate_tokens)
+
+
+def images_generate(original_method, version, tracer):
+    """
+    Wrap the `generate` method of the `Images` class to trace it.
+    """
+    def traced_method(wrapped, instance, args, kwargs):
+        base_url = str(instance._client._base_url) if hasattr(
+            instance, '_client') and hasattr(instance._client, '_base_url') else ""
+        service_provider = SERVICE_PROVIDERS['OPENAI']
+        span_attributes = {
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "llm",
+            "langtrace.service.version": version,
+            "langtrace.version": "1.0.0",
+            "url.full": base_url,
+            "llm.api": APIS["IMAGES_GENERATION"]["ENDPOINT"],
+            "llm.model": kwargs.get('model'),
+            "llm.stream": kwargs.get('stream'),
+            "llm.prompts": json.dumps([kwargs.get('prompt', [])])
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(APIS["IMAGES_GENERATION"]["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(*args, **kwargs)
+                if kwargs.get('stream') is False or kwargs.get('stream') is None:
+                    data = result.data[0] if hasattr(
+                        result, 'data') and len(result.data) > 0 else {}
+                    response = [{
+                        "url": data.url if hasattr(data, 'url') else "",
+                        "revised_prompt": data.revised_prompt if
+                        hasattr(data, 'revised_prompt') else "",
+                    }]
+                    span.set_attribute(
+                        "llm.responses", json.dumps(response))
+
+                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
+
+
+def chat_completions_create(original_method, version, tracer):
+    """Wrap the `create` method of the `ChatCompletion` class to trace it."""
+    def traced_method(wrapped, instance, args, kwargs):
+        base_url = str(instance._client._base_url) if hasattr(
+            instance, '_client') and hasattr(instance._client, '_base_url') else ""
+        service_provider = SERVICE_PROVIDERS['OPENAI']
+        span_attributes = {
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "llm",
+            "langtrace.service.version": version,
+            "langtrace.version": "1.0.0",
+            "url.full": base_url,
+            "llm.api": APIS["CHAT_COMPLETION"]["ENDPOINT"],
+            "llm.model": kwargs.get('model'),
+            "llm.prompts": json.dumps(kwargs.get('messages', [])),
+            "llm.stream": kwargs.get('stream'),
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+
+        if kwargs.get('temperature') is not None:
+            attributes.llm_temperature = kwargs.get('temperature')
+        if kwargs.get('top_p') is not None:
+            attributes.llm_top_p = kwargs.get('top_p')
+        if kwargs.get('user') is not None:
+            attributes.llm_user = kwargs.get('user')
+        if kwargs.get('functions') is not None:
+            attributes.llm_function_prompts = json.dumps(
+                kwargs.get('functions'))
+
+        # TODO(Karthik): Gotta figure out how to handle streaming with context
+        # with tracer.start_as_current_span(APIS["CHAT_COMPLETION"]["METHOD"],
+        #                                   kind=SpanKind.CLIENT) as span:
+        span = tracer.start_span(
+            APIS["CHAT_COMPLETION"]["METHOD"], kind=SpanKind.CLIENT)
+        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(*args, **kwargs)
+            if kwargs.get('stream') is False:
+                if hasattr(result, 'choices') and result.choices is not None:
+                    responses = [
+                        {
+                            "message": choice.message.content if choice.message and
+                            choice.message.content else choice.message.function_call.arguments
+                            if choice.message and
+                            choice.message.function_call.arguments else "",
+                            **({"content_filter_results": choice["content_filter_results"]}
+                               if "content_filter_results" in choice else {})
+                        }
+                        for choice in result.choices
+                    ]
+                    span.set_attribute(
+                        "llm.responses", json.dumps(responses))
+                else:
+                    responses = []
+                    span.set_attribute(
+                        "llm.responses", json.dumps(responses))
+                if hasattr(result, 'system_fingerprint') and \
+                        result.system_fingerprint is not None:
+                    span.set_attribute(
+                        "llm.system.fingerprint", result.system_fingerprint)
+                # Get the usage
+                if hasattr(result, 'usage') and result.usage is not None:
+                    usage = result.usage
+                    if usage is not None:
+                        usage_dict = {
+                            "prompt_tokens": result.usage.prompt_tokens,
+                            "completion_tokens": usage.completion_tokens,
+                            "total_tokens": usage.total_tokens
+                        }
+                        span.set_attribute(
+                            "llm.token.counts", json.dumps(usage_dict))
+                span.set_status(StatusCode.OK)
+                span.end()
+                return result
+            else:
+                prompt_tokens = calculate_prompt_tokens(json.dumps(
+                    kwargs.get('messages', {})[0]), kwargs.get('model'))
+                return handle_streaming_response(result, span, prompt_tokens,
+                                                 function_call=kwargs.get(
+                                                     'functions')
+                                                 is not None)
+        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
+            span.end()
+            raise
+
+    def handle_streaming_response(result, span, prompt_tokens, function_call=False):
+        """Process and yield streaming response chunks."""
+        result_content = []
+        span.add_event(Event.STREAM_START.value)
+        completion_tokens = 0
+        try:
+            for chunk in result:
+                if hasattr(chunk, 'choices') and chunk.choices is not None:
+                    token_counts = [
+                        estimate_tokens(choice.delta.content) if choice.delta
+                        and choice.delta.content
+                        else estimate_tokens(choice.delta.function_call.arguments)
+                        if choice.delta.function_call and
+                        choice.delta.function_call.arguments else 0
+                        for choice in chunk.choices
+                    ]
+                    completion_tokens += sum(token_counts)
+                    content = [
+                        choice.delta.content if choice.delta and choice.delta.content
+                        else choice.delta.function_call.arguments if choice.delta.function_call and
+                        choice.delta.function_call.arguments else ""
+                        for choice in chunk.choices
+                    ]
+                else:
+                    content = []
+                span.add_event(Event.STREAM_OUTPUT.value, {
+                    "response": "".join(content)
+                })
+                result_content.append(
+                    content[0] if len(content) > 0 else "")
+                yield chunk
+        finally:
+
+            # Finalize span after processing all chunks
+            span.add_event(Event.STREAM_END.value)
+            span.set_attribute("llm.token.counts", json.dumps({
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens
+            }))
+            if function_call is False:
+                span.set_attribute("llm.responses", json.dumps(
+                    {"message": {"role": "assistant", "content": "".join(result_content)}}))
+            else:
+                span.set_attribute("llm.responses", json.dumps(
+                    {"message": {"role": "assistant", "function_call": "".join(result_content)}}))
+            span.set_status(StatusCode.OK)
+            span.end()
+
+    # return the wrapped method
+    return traced_method
+
+
+def embeddings_create(original_method, version, tracer):
+    """
+    Wrap the `create` method of the `Embeddings` class to trace it.
+    """
+    def traced_method(wrapped, instance, args, kwargs):
+        base_url = str(instance._client._base_url) if hasattr(
+            instance, '_client') and hasattr(instance._client, '_base_url') else ""
+
+        service_provider = SERVICE_PROVIDERS['OPENAI']
+        span_attributes = {
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "llm",
+            "langtrace.service.version": version,
+            "langtrace.version": "1.0.0",
+            "url.full": base_url,
+            "llm.api": APIS["EMBEDDINGS_CREATE"]["ENDPOINT"],
+            "llm.model": kwargs.get('model'),
+            "llm.prompts": "",
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+        kwargs.get('encoding_format')
+
+        if kwargs.get('encoding_format') is not None:
+            attributes.llm_encoding_format = kwargs.get('encoding_format')
+        if kwargs.get('dimensions') is not None:
+            attributes["llm.dimensions"] = kwargs.get('dimensions')
+        if kwargs.get('user') is not None:
+            attributes["llm.user"] = kwargs.get('user')
+
+        with tracer.start_as_current_span(APIS["EMBEDDINGS_CREATE"]["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(*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

langtrace-python-sdk-1.0.10/src/instrumentation/openai/token_estimation.py

@@ -0,0 +1,48 @@
+"""
+This module contains functions to estimate the number of tokens in a prompt and 
+to calculate the price of a model based on its usage.
+"""
+
+from tiktoken import get_encoding
+
+from instrumentation.openai.constants import (OPENAI_COST_TABLE,
+                                              TIKTOKEN_MODEL_MAPPING)
+
+
+def estimate_tokens(prompt):
+    """
+    Estimate the number of tokens in a prompt."""
+    if prompt and len(prompt) > 0:
+        # Simplified token estimation: count the words.
+        return len([word for word in prompt.split() if word])
+    return 0
+
+
+def estimate_tokens_using_tiktoken(prompt, model):
+    """
+    Estimate the number of tokens in a prompt using tiktoken."""
+    encoding = get_encoding(model)
+    tokens = encoding.encode(prompt)
+    return len(tokens)
+
+
+def calculate_prompt_tokens(prompt_content, model):
+    """
+    Calculate the number of tokens in a prompt. If the model is supported by tiktoken, use it for the estimation."""
+    try:
+        tiktoken_model = TIKTOKEN_MODEL_MAPPING[model]
+        return estimate_tokens_using_tiktoken(prompt_content, tiktoken_model)
+    except Exception:
+        return estimate_tokens(prompt_content)  # Fallback method
+
+
+def calculate_price_from_usage(model, usage):
+    """
+    Calculate the price of a model based on its usage."""
+    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

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/instrumentation/pinecone/instrumentation.py

@@ -1,3 +1,7 @@
+"""
+Pinecone instrumentation
+"""
+
 import importlib.metadata
 from typing import Collection
 
@@ -7,11 +11,13 @@ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper
 
-from instrumentation.pinecone.lib.apis import APIS
+from instrumentation.pinecone.apis import APIS
 from instrumentation.pinecone.patch import generic_patch
 
 
 class PineconeInstrumentation(BaseInstrumentor):
+    """
+    The PineconeInstrumentation class represents the Pinecone instrumentation"""
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return ["pinecone-client >= 3.1.0"]
@@ -40,4 +46,4 @@ class PineconeInstrumentation(BaseInstrumentor):
             )
 
     def _uninstrument(self, **kwargs):
-        print(kwargs)
+        pass

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/instrumentation/pinecone/patch.py

@@ -1,14 +1,16 @@
-import json
-
+"""
+This module contains the patching logic for the Pinecone client."""
 from langtrace.trace_attributes import DatabaseSpanAttributes
-from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
 from instrumentation.constants import SERVICE_PROVIDERS
-from instrumentation.pinecone.lib.apis import APIS
+from instrumentation.pinecone.apis import APIS
 
 
 def generic_patch(original_method, method, version, tracer):
+    """
+    A generic patch method that wraps a function with a span"""
     def traced_method(wrapped, instance, args, kwargs):
         api = APIS[method]
         service_provider = SERVICE_PROVIDERS['PINECONE']

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10/src/langtrace_python_sdk.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langtrace-python-sdk
-Version: 1.0.9
+Version: 1.0.10
 Summary: LangTrace - Python SDK
 Home-page: https://github.com/Scale3-Labs/langtrace-python-sdk
 Author: Ali Waleed

{langtrace-python-sdk-1.0.9 → langtrace-python-sdk-1.0.10}/src/langtrace_python_sdk.egg-info/SOURCES.txt

@@ -17,14 +17,12 @@ src/examples/pinecone_example/__init__.py
 src/examples/pinecone_example/basic.py
 src/instrumentation/__init__.py
 src/instrumentation/constants.py
-src/instrumentation/setup.py
-src/instrumentation/utils.py
 src/instrumentation/with_root_span.py
 src/instrumentation/chroma/__init__.py
+src/instrumentation/chroma/apis.py
 src/instrumentation/chroma/instrumentation.py
 src/instrumentation/chroma/patch.py
 src/instrumentation/chroma/lib/__init__.py
-src/instrumentation/chroma/lib/apis.py
 src/instrumentation/langchain/__init__.py
 src/instrumentation/langchain/instrumentation.py
 src/instrumentation/langchain/patch.py
@@ -38,16 +36,17 @@ src/instrumentation/llamaindex/__init__.py
 src/instrumentation/llamaindex/instrumentation.py
 src/instrumentation/llamaindex/patch.py
 src/instrumentation/openai/__init__.py
+src/instrumentation/openai/apis.py
+src/instrumentation/openai/constants.py
 src/instrumentation/openai/instrumentation.py
 src/instrumentation/openai/patch.py
+src/instrumentation/openai/token_estimation.py
 src/instrumentation/openai/lib/__init__.py
-src/instrumentation/openai/lib/apis.py
-src/instrumentation/openai/lib/constants.py
 src/instrumentation/pinecone/__init__.py
+src/instrumentation/pinecone/apis.py
 src/instrumentation/pinecone/instrumentation.py
 src/instrumentation/pinecone/patch.py
 src/instrumentation/pinecone/lib/__init__.py
-src/instrumentation/pinecone/lib/apis.py
 src/langtrace_python_sdk.egg-info/PKG-INFO
 src/langtrace_python_sdk.egg-info/SOURCES.txt
 src/langtrace_python_sdk.egg-info/dependency_links.txt

langtrace-python-sdk-1.0.9/README.md

@@ -1,3 +0,0 @@
-# langtrace-python-sdk
-
-export PYTHONPATH="/Users/karthikkalyanaraman/work/langtrace/python-sdk:$PYTHONPATH"

langtrace-python-sdk-1.0.9/src/instrumentation/openai/patch.py

@@ -1,209 +0,0 @@
-import json
-
-from langtrace.trace_attributes import Event, LLMSpanAttributes
-from opentelemetry.trace import SpanKind, StatusCode
-from opentelemetry.trace.status import Status, StatusCode
-
-from instrumentation.constants import SERVICE_PROVIDERS
-from instrumentation.openai.lib.apis import APIS
-
-
-def images_generate(original_method, version, tracer):
-    def traced_method(wrapped, instance, args, kwargs):
-        base_url = str(instance._client._base_url) if hasattr(
-            instance, '_client') and hasattr(instance._client, '_base_url') else ""
-        service_provider = SERVICE_PROVIDERS['OPENAI']
-        span_attributes = {
-            "langtrace.service.name": service_provider,
-            "langtrace.service.type": "llm",
-            "langtrace.service.version": version,
-            "langtrace.version": "1.0.0",
-            "url.full": base_url,
-            "llm.api": APIS["IMAGES_GENERATION"]["ENDPOINT"],
-            "llm.model": kwargs.get('model'),
-            "llm.stream": kwargs.get('stream'),
-            "llm.prompts": json.dumps([kwargs.get('prompt', [])])
-        }
-
-        attributes = LLMSpanAttributes(**span_attributes)
-
-        with tracer.start_as_current_span(APIS["IMAGES_GENERATION"]["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(*args, **kwargs)
-                if kwargs.get('stream') is False or kwargs.get('stream') is None:
-                    data = result.data[0] if hasattr(
-                        result, 'data') and len(result.data) > 0 else {}
-                    response = [{
-                        "url": data.url if hasattr(data, 'url') else "",
-                        "revised_prompt": data.revised_prompt if hasattr(data, 'revised_prompt') else "",
-                    }]
-                    span.set_attribute(
-                        "llm.responses", json.dumps(response))
-
-                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
-
-
-def chat_completions_create(original_method, version, tracer):
-    def traced_method(wrapped, instance, args, kwargs):
-        base_url = str(instance._client._base_url) if hasattr(
-            instance, '_client') and hasattr(instance._client, '_base_url') else ""
-        service_provider = SERVICE_PROVIDERS['OPENAI']
-        span_attributes = {
-            "langtrace.service.name": service_provider,
-            "langtrace.service.type": "llm",
-            "langtrace.service.version": version,
-            "langtrace.version": "1.0.0",
-            "url.full": base_url,
-            "llm.api": APIS["CHAT_COMPLETION"]["ENDPOINT"],
-            "llm.model": kwargs.get('model'),
-            "llm.prompts": json.dumps(kwargs.get('messages', [])),
-            "llm.stream": kwargs.get('stream'),
-        }
-
-        attributes = LLMSpanAttributes(**span_attributes)
-
-        if kwargs.get('temperature') is not None:
-            attributes.llm_temperature = kwargs.get('temperature')
-        if kwargs.get('top_p') is not None:
-            attributes.llm_top_p = kwargs.get('top_p')
-        if kwargs.get('user') is not None:
-            attributes.llm_user = kwargs.get('user')
-
-        with tracer.start_as_current_span(APIS["CHAT_COMPLETION"]["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(*args, **kwargs)
-                if kwargs.get('stream') is False:
-                    if hasattr(result, 'choices') and result.choices is not None:
-                        responses = [
-                            {
-                                "message": choice.message.content if choice.message and choice.message.content else "",
-                                **({"content_filter_results": choice["content_filter_results"]} if "content_filter_results" in choice else {})
-                            }
-                            for choice in result.choices
-                        ]
-                    else:
-                        responses = []
-                    span.set_attribute("llm.responses", json.dumps(responses))
-
-                    if hasattr(result, 'system_fingerprint') and result.system_fingerprint is not None:
-                        span.set_attribute(
-                            "llm.system.fingerprint", result.system_fingerprint)
-
-                    # Get the usage
-                    if hasattr(result, 'usage') and result.usage is not None:
-                        usage = result.usage
-                        if usage is not None:
-                            usage_dict = {
-                                "prompt_tokens": result.usage.prompt_tokens,
-                                "completion_tokens": usage.completion_tokens,
-                                "total_tokens": usage.total_tokens
-                            }
-                            span.set_attribute(
-                                "llm.token.counts", json.dumps(usage_dict))
-
-                    span.set_status(StatusCode.OK)
-                    return result
-                else:
-                    result_content = []
-                    span.add_event(Event.STREAM_START.value)
-
-                    for chunk in result:
-                        # Assuming `chunk` has a structure similar to what OpenAI might return,
-                        # adjust the access accordingly based on actual response structure.
-                        if hasattr(chunk, 'choices') and chunk.choices is not None:
-                            content = [
-                                choice.delta.content if choice.delta and choice.delta.content else ""
-                                for choice in chunk.choices
-                            ]
-                        else:
-                            content = []
-                        span.add_event(Event.STREAM_OUTPUT.value, {
-                            "response": "".join(content)
-                        })
-                        result_content.append(
-                            content[0] if len(content) > 0 else "")
-                    span.add_event(Event.STREAM_END.value)
-                    span.set_attribute("llm.responses", json.dumps(
-                        {"message": {"role": "assistant", "content": "".join(result_content)}}))
-
-            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 the wrapped method
-    return traced_method
-
-
-def embeddings_create(original_method, version, tracer):
-    def traced_method(wrapped, instance, args, kwargs):
-        base_url = str(instance._client._base_url) if hasattr(
-            instance, '_client') and hasattr(instance._client, '_base_url') else ""
-
-        service_provider = SERVICE_PROVIDERS['OPENAI']
-        span_attributes = {
-            "langtrace.service.name": service_provider,
-            "langtrace.service.type": "llm",
-            "langtrace.service.version": version,
-            "langtrace.version": "1.0.0",
-            "url.full": base_url,
-            "llm.api": APIS["EMBEDDINGS_CREATE"]["ENDPOINT"],
-            "llm.model": kwargs.get('model'),
-            "llm.prompts": "",
-        }
-
-        attributes = LLMSpanAttributes(**span_attributes)
-        kwargs.get('encoding_format')
-
-        if kwargs.get('encoding_format') is not None:
-            attributes.llm_encoding_format = kwargs.get('encoding_format')
-        if kwargs.get('dimensions') is not None:
-            attributes["llm.dimensions"] = kwargs.get('dimensions')
-        if kwargs.get('user') is not None:
-            attributes["llm.user"] = kwargs.get('user')
-
-        with tracer.start_as_current_span(APIS["EMBEDDINGS_CREATE"]["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(*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

langtrace-python-sdk-1.0.9/src/instrumentation/setup.py

@@ -1,50 +0,0 @@
-
-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")

langtrace-python-sdk-1.0.9/src/instrumentation/utils.py

@@ -1,27 +0,0 @@
-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