langtrace-python-sdk 3.3.22__tar.gz → 3.3.24__tar.gz

{langtrace_python_sdk-3.3.22 → langtrace_python_sdk-3.3.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langtrace-python-sdk
-Version: 3.3.22
+Version: 3.3.24
 Summary: Python SDK for LangTrace
 Project-URL: Homepage, https://github.com/Scale3-Labs/langtrace-python-sdk
 Author-email: Scale3 Labs <engineering@scale3labs.com>

langtrace_python_sdk-3.3.24/src/examples/awsbedrock_examples/__init__.py

@@ -0,0 +1,18 @@
+from examples.awsbedrock_examples.converse import (
+    use_invoke_model_anthropic,
+    use_invoke_model_titan,
+    use_invoke_model_llama,
+)
+from langtrace_python_sdk import langtrace, with_langtrace_root_span
+
+
+class AWSBedrockRunner:
+    @with_langtrace_root_span("AWS_Bedrock")
+    def run(self):
+
+        # use_converse_stream()
+        # use_converse()
+        # use_invoke_model_anthropic(stream=True)
+        # use_invoke_model_cohere()
+        # use_invoke_model_llama(stream=False)
+        use_invoke_model_titan(stream=False)

{langtrace_python_sdk-3.3.22 → langtrace_python_sdk-3.3.24}/src/examples/awsbedrock_examples/converse.py

@@ -88,6 +88,12 @@ def use_invoke_model_titan(stream=False):
             response = brt.invoke_model_with_response_stream(
                 body=body, modelId=modelId, accept=accept, contentType=contentType
             )
+            # Extract and print the response text in real-time.
+            for event in response["body"]:
+                chunk = json.loads(event["chunk"]["bytes"])
+                if "outputText" in chunk:
+                    print(chunk["outputText"], end="")
+
         else:
             response = brt.invoke_model(
                 body=body, modelId=modelId, accept=accept, contentType=contentType
@@ -130,7 +136,8 @@ def use_invoke_model_anthropic(stream=False):
             for event in stream_response:
                 chunk = event.get("chunk")
                 if chunk:
-                    print(json.loads(chunk.get("bytes").decode()))
+                    # print(json.loads(chunk.get("bytes").decode()))
+                    pass
 
     else:
         response = brt.invoke_model(
@@ -141,7 +148,7 @@ def use_invoke_model_anthropic(stream=False):
         print(response_body.get("completion"))
 
 
-def use_invoke_model_llama():
+def use_invoke_model_llama(stream=False):
     model_id = "meta.llama3-8b-instruct-v1:0"
     prompt = "What is the capital of France?"
     max_gen_len = 128
@@ -157,11 +164,18 @@ def use_invoke_model_llama():
             "top_p": top_p,
         }
     )
-    response = brt.invoke_model(body=body, modelId=model_id)
-
-    response_body = json.loads(response.get("body").read())
 
-    return response_body
+    if stream:
+        response = brt.invoke_model_with_response_stream(body=body, modelId=model_id)
+        for event in response["body"]:
+            chunk = json.loads(event["chunk"]["bytes"])
+            if "generation" in chunk:
+                # print(chunk["generation"], end="")
+                pass
+    else:
+        response = brt.invoke_model(body=body, modelId=model_id)
+        response_body = json.loads(response.get("body").read())
+        return response_body
 
 
 # print(get_foundation_models())

{langtrace_python_sdk-3.3.22 → langtrace_python_sdk-3.3.24}/src/examples/gemini_example/main.py

@@ -20,16 +20,16 @@ async def async_demo():
 
 def basic():
     generate()
-    generate(stream=True, with_tools=True)
+    # generate(stream=True, with_tools=True)
 
     # image_to_text()
     # audio_to_text()
-    asyncio.run(async_demo())
+    # asyncio.run(async_demo())
 
 
 def generate(stream=False, with_tools=False):
     model = genai.GenerativeModel(
-        "gemini-1.5-pro", system_instruction="You are a cat. Your name is Neko."
+        "gemini-2.0-flash-exp", system_instruction="You are a cat. Your name is Neko."
     )
 
     response = model.generate_content(

langtrace_python_sdk-3.3.24/src/examples/google_genai_example/__init__.py

@@ -0,0 +1,7 @@
+from .main import generate_content, generate_content_streaming
+
+
+class GoogleGenaiRunner:
+    def run(self):
+        # generate_content()
+        generate_content_streaming()

langtrace_python_sdk-3.3.24/src/examples/google_genai_example/main.py

@@ -0,0 +1,27 @@
+from google import genai
+from dotenv import load_dotenv
+import os
+from langtrace_python_sdk import langtrace
+
+load_dotenv()
+langtrace.init(write_spans_to_console=False)
+
+
+def generate_content():
+    # Only run this block for Google AI API
+    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-exp", contents="What is your name?"
+    )
+
+    print(response.text)
+
+
+def generate_content_streaming():
+    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+    response = client.models.generate_content_stream(
+        model="gemini-2.0-flash-exp", contents="What is your name?"
+    )
+
+    for chunk in response:
+        pass

{langtrace_python_sdk-3.3.22 → langtrace_python_sdk-3.3.24}/src/langtrace_python_sdk/instrumentation/__init__.py

@@ -24,6 +24,7 @@ from .litellm import LiteLLMInstrumentation
 from .pymongo import PyMongoInstrumentation
 from .cerebras import CerebrasInstrumentation
 from .milvus import MilvusInstrumentation
+from .google_genai import GoogleGenaiInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -52,4 +53,5 @@ __all__ = [
     "AWSBedrockInstrumentation",
     "CerebrasInstrumentation",
     "MilvusInstrumentation",
+    "GoogleGenaiInstrumentation",
 ]

{langtrace_python_sdk-3.3.22 → langtrace_python_sdk-3.3.24}/src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py

@@ -16,9 +16,7 @@ limitations under the License.
 
 import json
 
-from langtrace_python_sdk.instrumentation.aws_bedrock.bedrock_streaming_wrapper import (
-    StreamingWrapper,
-)
+from wrapt import ObjectProxy
 from .stream_body_wrapper import BufferedStreamBody
 from functools import wraps
 from langtrace.trace_attributes import (
@@ -87,6 +85,11 @@ def patch_aws_bedrock(tracer, version):
 
         client = wrapped(*args, **kwargs)
         client.invoke_model = patch_invoke_model(client.invoke_model, tracer, version)
+        client.invoke_model_with_response_stream = (
+            patch_invoke_model_with_response_stream(
+                client.invoke_model_with_response_stream, tracer, version
+            )
+        )
 
         client.converse = patch_converse(client.converse, tracer, version)
         client.converse_stream = patch_converse_stream(
@@ -186,6 +189,56 @@ def patch_invoke_model(original_method, tracer, version):
     return traced_method
 
 
+def patch_invoke_model_with_response_stream(original_method, tracer, version):
+    @wraps(original_method)
+    def traced_method(*args, **kwargs):
+        modelId = kwargs.get("modelId")
+        (vendor, _) = modelId.split(".")
+        span_attributes = {
+            **get_langtrace_attributes(version, vendor, vendor_type="framework"),
+            **get_extra_attributes(),
+        }
+        span = tracer.start_span(
+            name=get_span_name("aws_bedrock.invoke_model_with_response_stream"),
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        )
+        set_span_attributes(span, span_attributes)
+        response = original_method(*args, **kwargs)
+        if span.is_recording():
+            handle_streaming_call(span, kwargs, response)
+        return response
+
+    return traced_method
+
+
+def handle_streaming_call(span, kwargs, response):
+
+    def stream_finished(response_body):
+        request_body = json.loads(kwargs.get("body"))
+
+        (vendor, model) = kwargs.get("modelId").split(".")
+
+        set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, model)
+        set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, model)
+
+        if vendor == "amazon":
+            set_amazon_attributes(span, request_body, response_body)
+
+        if vendor == "anthropic":
+            if "prompt" in request_body:
+                set_anthropic_completions_attributes(span, request_body, response_body)
+            elif "messages" in request_body:
+                set_anthropic_messages_attributes(span, request_body, response_body)
+
+        if vendor == "meta":
+            set_llama_meta_attributes(span, request_body, response_body)
+
+        span.end()
+
+    response["body"] = StreamingBedrockWrapper(response["body"], stream_finished)
+
+
 def handle_call(span, kwargs, response):
     modelId = kwargs.get("modelId")
     (vendor, model_name) = modelId.split(".")
@@ -195,7 +248,6 @@ def handle_call(span, kwargs, response):
     request_body = json.loads(kwargs.get("body"))
     response_body = json.loads(response.get("body").read())
 
-    set_span_attribute(span, SpanAttributes.LLM_SYSTEM, vendor)
     set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, modelId)
     set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, modelId)
 
@@ -222,12 +274,18 @@ def set_llama_meta_attributes(span, request_body, response_body):
     set_span_attribute(
         span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, request_body.get("max_gen_len")
     )
+    if "invocation_metrics" in response_body:
+        input_tokens = response_body.get("invocation_metrics").get("inputTokenCount")
+        output_tokens = response_body.get("invocation_metrics").get("outputTokenCount")
+    else:
+        input_tokens = response_body.get("prompt_token_count")
+        output_tokens = response_body.get("generation_token_count")
 
     set_usage_attributes(
         span,
         {
-            "input_tokens": response_body.get("prompt_token_count"),
-            "output_tokens": response_body.get("generation_token_count"),
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
         },
     )
 
@@ -245,7 +303,6 @@ def set_llama_meta_attributes(span, request_body, response_body):
         }
     ]
     set_span_attribute(span, SpanAttributes.LLM_PROMPTS, json.dumps(prompts))
-    print(completions)
     set_event_completion(span, completions)
 
 
@@ -257,13 +314,22 @@ def set_amazon_attributes(span, request_body, response_body):
             "content": request_body.get("inputText"),
         }
     ]
-    completions = [
-        {
-            "role": "assistant",
-            "content": result.get("outputText"),
-        }
-        for result in response_body.get("results")
-    ]
+    if "results" in response_body:
+        completions = [
+            {
+                "role": "assistant",
+                "content": result.get("outputText"),
+            }
+            for result in response_body.get("results")
+        ]
+
+    else:
+        completions = [
+            {
+                "role": "assistant",
+                "content": response_body.get("outputText"),
+            }
+        ]
     set_span_attribute(
         span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, config.get("maxTokenCount")
     )
@@ -272,13 +338,19 @@ def set_amazon_attributes(span, request_body, response_body):
     )
     set_span_attribute(span, SpanAttributes.LLM_REQUEST_TOP_P, config.get("topP"))
     set_span_attribute(span, SpanAttributes.LLM_PROMPTS, json.dumps(prompts))
+    input_tokens = response_body.get("inputTextTokenCount")
+    if "results" in response_body:
+        output_tokens = sum(
+            int(result.get("tokenCount")) for result in response_body.get("results")
+        )
+    else:
+        output_tokens = response_body.get("outputTextTokenCount")
+
     set_usage_attributes(
         span,
         {
-            "input_tokens": response_body.get("inputTextTokenCount"),
-            "output_tokens": sum(
-                int(result.get("tokenCount")) for result in response_body.get("results")
-            ),
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
         },
     )
     set_event_completion(span, completions)
@@ -320,7 +392,7 @@ def set_anthropic_messages_attributes(span, request_body, response_body):
     set_span_attribute(
         span,
         SpanAttributes.LLM_REQUEST_MAX_TOKENS,
-        request_body.get("max_tokens_to_sample"),
+        request_body.get("max_tokens_to_sample") or request_body.get("max_tokens"),
     )
     set_span_attribute(
         span,
@@ -394,3 +466,62 @@ def set_span_streaming_response(span, response):
         set_event_completion(
             span, [{"role": role or "assistant", "content": streaming_response}]
         )
+
+
+class StreamingBedrockWrapper(ObjectProxy):
+    def __init__(
+        self,
+        response,
+        stream_done_callback=None,
+    ):
+        super().__init__(response)
+
+        self._stream_done_callback = stream_done_callback
+        self._accumulating_body = {"generation": ""}
+
+    def __iter__(self):
+        for event in self.__wrapped__:
+            self._process_event(event)
+            yield event
+
+    def _process_event(self, event):
+        chunk = event.get("chunk")
+        if not chunk:
+            return
+
+        decoded_chunk = json.loads(chunk.get("bytes").decode())
+        type = decoded_chunk.get("type")
+
+        if type is None and "outputText" in decoded_chunk:
+            self._stream_done_callback(decoded_chunk)
+            return
+        if "generation" in decoded_chunk:
+            self._accumulating_body["generation"] += decoded_chunk.get("generation")
+
+        if type == "message_start":
+            self._accumulating_body = decoded_chunk.get("message")
+        elif type == "content_block_start":
+            self._accumulating_body["content"].append(
+                decoded_chunk.get("content_block")
+            )
+        elif type == "content_block_delta":
+            self._accumulating_body["content"][-1]["text"] += decoded_chunk.get(
+                "delta"
+            ).get("text")
+
+        elif self.has_finished(type, decoded_chunk):
+            self._accumulating_body["invocation_metrics"] = decoded_chunk.get(
+                "amazon-bedrock-invocationMetrics"
+            )
+            self._stream_done_callback(self._accumulating_body)
+
+    def has_finished(self, type, chunk):
+        if type and type == "message_stop":
+            return True
+
+        if "completionReason" in chunk and chunk.get("completionReason") == "FINISH":
+            return True
+
+        if "stop_reason" in chunk and chunk.get("stop_reason") is not None:
+            return True
+        return False

langtrace_python_sdk-3.3.24/src/langtrace_python_sdk/instrumentation/google_genai/__init__.py

@@ -0,0 +1,3 @@
+from .instrumentation import GoogleGenaiInstrumentation
+
+__all__ = ["GoogleGenaiInstrumentation"]

langtrace_python_sdk-3.3.24/src/langtrace_python_sdk/instrumentation/google_genai/instrumentation.py

@@ -0,0 +1,30 @@
+from typing import Collection
+from importlib_metadata import version as v
+from wrapt import wrap_function_wrapper as _W
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from .patch import patch_google_genai, patch_google_genai_streaming
+
+
+class GoogleGenaiInstrumentation(BaseInstrumentor):
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["google-genai >= 0.1.0", "google-generativeai < 1.0.0"]
+
+    def _instrument(self, **kwargs):
+        trace_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", trace_provider)
+        version = v("google-genai")
+
+        _W(
+            module="google.genai",
+            name="models.Models.generate_content",
+            wrapper=patch_google_genai(tracer, version),
+        )
+        _W(
+            module="google.genai",
+            name="models.Models.generate_content_stream",
+            wrapper=patch_google_genai_streaming(tracer, version),
+        )
+
+    def _uninstrument(self, **kwargs):
+        pass

langtrace_python_sdk-3.3.24/src/langtrace_python_sdk/instrumentation/google_genai/patch.py

@@ -0,0 +1,126 @@
+from langtrace_python_sdk.utils.llm import (
+    get_langtrace_attributes,
+    get_llm_request_attributes,
+    set_span_attributes,
+    set_usage_attributes,
+    set_span_attribute,
+    set_event_completion,
+)
+from langtrace_python_sdk.utils import handle_span_error
+
+from opentelemetry.trace import Tracer, SpanKind
+from opentelemetry.sdk.trace import Span
+from langtrace.trace_attributes import SpanAttributes
+
+from typing import Iterator
+
+
+def patch_google_genai(tracer: Tracer, version: str):
+    def traced_method(wrapped, instance, args, kwargs):
+        prompt = [
+            {
+                "role": "user",
+                "content": kwargs["contents"],
+            }
+        ]
+        span_attributes = {
+            **get_langtrace_attributes(
+                service_provider="google_genai", version=version
+            ),
+            **get_llm_request_attributes(kwargs=kwargs, prompts=prompt),
+        }
+        with tracer.start_as_current_span(
+            name="google.genai.generate_content",
+            kind=SpanKind.CLIENT,
+        ) as span:
+            try:
+                set_span_attributes(span, span_attributes)
+                response = wrapped(*args, **kwargs)
+                set_response_attributes(span, response)
+                return response
+            except Exception as error:
+                handle_span_error(span, error)
+                raise
+
+    return traced_method
+
+
+def patch_google_genai_streaming(tracer: Tracer, version: str):
+    def traced_method(wrapped, instance, args, kwargs):
+        prompt = [
+            {
+                "role": "user",
+                "content": kwargs["contents"],
+            }
+        ]
+        span_attributes = {
+            **get_langtrace_attributes(
+                service_provider="google_genai", version=version
+            ),
+            **get_llm_request_attributes(kwargs=kwargs, prompts=prompt),
+        }
+        with tracer.start_as_current_span(
+            name="google.genai.generate_content_stream",
+            kind=SpanKind.CLIENT,
+        ) as span:
+            set_span_attributes(span, span_attributes)
+            response = wrapped(*args, **kwargs)
+            set_streaming_response_attributes(span, response)
+            return response
+
+    return traced_method
+
+
+def set_streaming_response_attributes(span: Span, response):
+    accum_completion = ""
+    for chunk in response:
+        set_span_attribute(
+            span,
+            SpanAttributes.LLM_RESPONSE_MODEL,
+            chunk.model_version,
+        )
+        candidates = chunk.candidates
+        for candidate in candidates:
+            set_span_attribute(
+                span,
+                SpanAttributes.LLM_RESPONSE_FINISH_REASON,
+                candidate.finish_reason,
+            )
+
+            accum_completion += candidate.content.parts[0].text
+
+        if chunk.usage_metadata:
+            set_usage_attributes(
+                span,
+                {
+                    "input_tokens": chunk.usage_metadata.prompt_token_count,
+                    "output_tokens": chunk.usage_metadata.candidates_token_count,
+                },
+            )
+    set_event_completion(span, [{"role": "assistant", "content": accum_completion}])
+
+
+def set_response_attributes(span: Span, response):
+    completions = []
+    for candidate in response.candidates:
+        set_span_attribute(
+            span, SpanAttributes.LLM_RESPONSE_FINISH_REASON, candidate.finish_reason
+        )
+        parts = candidate.content.parts
+        role = candidate.content.role
+        completion = {
+            "role": role or "assistant",
+            "content": [part.text for part in parts],
+        }
+        completions.append(completion)
+
+    set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, response.model_version)
+    set_event_completion(span, completions)
+    if response.usage_metadata:
+        set_usage_attributes(
+            span,
+            {
+                "input_tokens": response.usage_metadata.prompt_token_count,
+                "output_tokens": response.usage_metadata.candidates_token_count,
+            },
+        )

{langtrace_python_sdk-3.3.22 → langtrace_python_sdk-3.3.24}/src/langtrace_python_sdk/langtrace.py

@@ -68,6 +68,7 @@ from langtrace_python_sdk.instrumentation import (
     PyMongoInstrumentation,
     CerebrasInstrumentation,
     MilvusInstrumentation,
+    GoogleGenaiInstrumentation,
 )
 from opentelemetry.util.re import parse_env_headers
 
@@ -301,6 +302,7 @@ def init(
         "vertexai": VertexAIInstrumentation(),
         "google-cloud-aiplatform": VertexAIInstrumentation(),
         "google-generativeai": GeminiInstrumentation(),
+        "google-genai": GoogleGenaiInstrumentation(),
         "mistralai": MistralInstrumentation(),
         "boto3": AWSBedrockInstrumentation(),
         "autogen": AutogenInstrumentation(),

langtrace_python_sdk-3.3.24/src/langtrace_python_sdk/version.py

@@ -0,0 +1 @@
+__version__ = "3.3.24"

{langtrace_python_sdk-3.3.22 → langtrace_python_sdk-3.3.24}/src/run_example.py

@@ -22,6 +22,7 @@ ENABLED_EXAMPLES = {
     "mistral": False,
     "awsbedrock": True,
     "cerebras": False,
+    "google_genai": False,
 }
 
 if ENABLED_EXAMPLES["anthropic"]:
@@ -137,3 +138,9 @@ if ENABLED_EXAMPLES["cerebras"]:
 
     print(Fore.BLUE + "Running Cerebras example" + Fore.RESET)
     CerebrasRunner().run()
+
+if ENABLED_EXAMPLES["google_genai"]:
+    from examples.google_genai_example import GoogleGenaiRunner
+
+    print(Fore.BLUE + "Running Google Genai example" + Fore.RESET)
+    GoogleGenaiRunner().run()

langtrace_python_sdk-3.3.22/src/examples/awsbedrock_examples/__init__.py

@@ -1,19 +0,0 @@
-from examples.awsbedrock_examples.converse import (
-    use_converse_stream,
-    use_converse,
-    use_invoke_model_anthropic,
-    use_invoke_model_cohere,
-    use_invoke_model_amazon,
-)
-from langtrace_python_sdk import langtrace, with_langtrace_root_span
-
-
-class AWSBedrockRunner:
-    @with_langtrace_root_span("AWS_Bedrock")
-    def run(self):
-
-        use_converse_stream()
-        use_converse()
-        use_invoke_model_anthropic()
-        use_invoke_model_cohere()
-        use_invoke_model_amazon()

langtrace_python_sdk-3.3.22/src/langtrace_python_sdk/version.py

@@ -1 +0,0 @@
-__version__ = "3.3.22"