langtrace-python-sdk 2.1.29__py3-none-any.whl → 2.2.2__py3-none-any.whl

langtrace_python_sdk/instrumentation/dspy/patch.py

@@ -39,25 +39,25 @@ def patch_bootstrapfewshot_optimizer(operation_name, version, tracer):
                     ),
                 }
                 span_attributes["dspy.optimizer.module.prog"] = json.dumps(prog)
-        if "metric" in instance and instance.metric:
-            span_attributes["dspy.optimizer.metric"] = instance.metric.__name__
+        if hasattr(instance, 'metric'):
+            span_attributes["dspy.optimizer.metric"] = getattr(instance, 'metric').__name__
         if kwargs.get("trainset") and len(kwargs.get("trainset")) > 0:
             span_attributes["dspy.optimizer.trainset"] = str(kwargs.get("trainset"))
         config = {}
-        if "metric_threshold" in instance and instance.metric_threshold:
-            config["metric_threshold"] = instance.metric_threshold
-        if "teacher_settings" in instance and instance.teacher_settings:
-            config["teacher_settings"] = instance.teacher_settings
-        if "max_bootstrapped_demos" in instance and instance.max_bootstrapped_demos:
-            config["max_bootstrapped_demos"] = instance.max_bootstrapped_demos
-        if "max_labeled_demos" in instance and instance.max_labeled_demos:
-            config["max_labeled_demos"] = instance.max_labeled_demos
-        if "max_rounds" in instance and instance.max_rounds:
-            config["max_rounds"] = instance.max_rounds
-        if "max_errors" in instance and instance.max_errors:
-            config["max_errors"] = instance.max_errors
-        if "error_count" in instance and instance.error_count:
-            config["error_count"] = instance.error_count
+        if hasattr(instance, 'metric_threshold'):
+            config["metric_threshold"] = getattr(instance, 'metric_threshold')
+        if hasattr(instance, 'teacher_settings'):
+            config["teacher_settings"] = getattr(instance, 'teacher_settings')
+        if hasattr(instance, 'max_bootstrapped_demos'):
+            config["max_bootstrapped_demos"] = getattr(instance, 'max_bootstrapped_demos')
+        if hasattr(instance, 'max_labeled_demos'):
+            config["max_labeled_demos"] = getattr(instance, 'max_labeled_demos')
+        if hasattr(instance, 'max_rounds'):
+            config["max_rounds"] = getattr(instance, 'max_rounds')
+        if hasattr(instance, 'max_steps'):
+            config["max_errors"] = getattr(instance, 'max_errors')
+        if hasattr(instance, 'error_count'):
+            config["error_count"] = getattr(instance, 'error_count')
         if config and len(config) > 0:
             span_attributes["dspy.optimizer.config"] = json.dumps(config)
 
@@ -147,30 +147,30 @@ def patch_evaluate(operation_name, version, tracer):
             **(extra_attributes if extra_attributes is not None else {}),
         }
 
-        if "devset" in instance and instance.devset is not None:
-            span_attributes["dspy.evaluate.devset"] = str(instance.devset)
-        if "display" in instance and instance.display is not None:
-            span_attributes["dspy.evaluate.display"] = str(instance.display)
-        if "num_threads" in instance and instance.num_threads is not None:
-            span_attributes["dspy.evaluate.num_threads"] = str(instance.num_threads)
-        if "return_outputs" in instance and instance.return_outputs is not None:
+        if hasattr(instance, "devset"):
+            span_attributes["dspy.evaluate.devset"] = str(getattr(instance, "devset"))
+        if hasattr(instance, "trainset"):
+            span_attributes["dspy.evaluate.display"] = str(getattr(instance, "trainset"))
+        if hasattr(instance, "num_threads"):
+            span_attributes["dspy.evaluate.num_threads"] = str(getattr(instance, "num_threads"))
+        if hasattr(instance, "return_outputs"):
             span_attributes["dspy.evaluate.return_outputs"] = str(
-                instance.return_outputs
+                getattr(instance, "return_outputs")
             )
-        if "display_table" in instance and instance.display_table is not None:
-            span_attributes["dspy.evaluate.display_table"] = str(instance.display_table)
-        if "display_progress" in instance and instance.display_progress is not None:
+        if hasattr(instance, "display_table"):
+            span_attributes["dspy.evaluate.display_table"] = str(getattr(instance, "display_table"))
+        if hasattr(instance, "display_progress"):
             span_attributes["dspy.evaluate.display_progress"] = str(
-                instance.display_progress
+                getattr(instance, "display_progress")
             )
-        if "metric" in instance and instance.metric is not None:
-            span_attributes["dspy.evaluate.metric"] = instance.metric.__name__
-        if "error_count" in instance and instance.error_count is not None:
-            span_attributes["dspy.evaluate.error_count"] = str(instance.error_count)
-        if "error_lock" in instance and instance.error_lock is not None:
-            span_attributes["dspy.evaluate.error_lock"] = str(instance.error_lock)
-        if "max_errors" in instance and instance.max_errors is not None:
-            span_attributes["dspy.evaluate.max_errors"] = str(instance.max_errors)
+        if hasattr(instance, "metric"):
+            span_attributes["dspy.evaluate.metric"] = getattr(instance, "metric").__name__
+        if hasattr(instance, "error_count"):
+            span_attributes["dspy.evaluate.error_count"] = str(getattr(instance, "error_count"))
+        if hasattr(instance, "error_lock"):
+            span_attributes["dspy.evaluate.error_lock"] = str(getattr(instance, "error_lock"))
+        if hasattr(instance, "max_errors"):
+            span_attributes["dspy.evaluate.max_errors"] = str(getattr(instance, "max_errors"))
         if args and len(args) > 0:
             span_attributes["dspy.evaluate.args"] = str(args)
 

langtrace_python_sdk/instrumentation/gemini/__init__.py

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

langtrace_python_sdk/instrumentation/gemini/instrumentation.py

@@ -0,0 +1,36 @@
+from typing import Collection
+from importlib_metadata import version as v
+from langtrace_python_sdk.constants.instrumentation.gemini import APIS
+from wrapt import wrap_function_wrapper as _W
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from .patch import patch_gemini, apatch_gemini
+
+
+class GeminiInstrumentation(BaseInstrumentor):
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["google-generativeai >= 0.5.0"]
+
+    def _instrument(self, **kwargs):
+        trace_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", trace_provider)
+        version = v("google-cloud-aiplatform")
+
+        for _, api_config in APIS.items():
+            module = api_config.get("module")
+            operation = api_config.get("operation")
+            method = api_config.get("method")
+            name = f"{method}.{operation}"
+
+            _W(
+                module=module,
+                name=name,
+                wrapper=(
+                    apatch_gemini(name, version, tracer)
+                    if operation == "generate_content_async"
+                    else patch_gemini(name, version, tracer)
+                ),
+            )
+
+    def _uninstrument(self, **kwargs):
+        pass

langtrace_python_sdk/instrumentation/gemini/patch.py

@@ -0,0 +1,186 @@
+from langtrace.trace_attributes import LLMSpanAttributes, SpanAttributes
+from opentelemetry import trace
+from opentelemetry.trace import Span, SpanKind, Tracer
+from opentelemetry.trace.propagation import set_span_in_context
+from opentelemetry.trace.status import Status, StatusCode
+
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+from langtrace_python_sdk.utils.llm import (
+    get_extra_attributes,
+    get_langtrace_attributes,
+    get_llm_request_attributes,
+    get_llm_url,
+    is_streaming,
+    set_event_completion,
+    set_event_completion_chunk,
+    set_span_attributes,
+    set_usage_attributes,
+)
+
+
+def patch_gemini(name, version, tracer: Tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS["GEMINI"]
+        prompts = serialize_prompts(args, kwargs, instance)
+        span_attributes = {
+            **get_langtrace_attributes(version, service_provider),
+            **get_llm_request_attributes(
+                kwargs,
+                prompts=prompts,
+                model=get_llm_model(instance),
+            ),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_PATH: "",
+            **get_extra_attributes(),
+        }
+        attributes = LLMSpanAttributes(**span_attributes)
+        span = tracer.start_span(
+            name=name,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        )
+
+        try:
+            set_span_attributes(span, attributes)
+            result = wrapped(*args, **kwargs)
+            if is_streaming(kwargs):
+                return build_streaming_response(span, result)
+
+            else:
+                set_response_attributes(span, result)
+            span.end()
+            return result
+        except Exception as error:
+            span.record_exception(error)
+            span.set_status(Status(StatusCode.ERROR, str(error)))
+            span.end()
+            raise
+
+    return traced_method
+
+
+def apatch_gemini(name, version, tracer: Tracer):
+    async def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS["GEMINI"]
+        prompts = serialize_prompts(args, kwargs, instance)
+        span_attributes = {
+            **get_langtrace_attributes(version, service_provider),
+            **get_llm_request_attributes(
+                kwargs,
+                prompts=prompts,
+                model=get_llm_model(instance),
+            ),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_PATH: "",
+            **get_extra_attributes(),
+        }
+        attributes = LLMSpanAttributes(**span_attributes)
+        span = tracer.start_span(
+            name=name,
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        )
+
+        try:
+            set_span_attributes(span, attributes)
+            result = await wrapped(*args, **kwargs)
+            if is_streaming(kwargs):
+                return abuild_streaming_response(span, result)
+            else:
+                set_response_attributes(span, result)
+            span.end()
+            return result
+        except Exception as error:
+            span.record_exception(error)
+            span.set_status(Status(StatusCode.ERROR, str(error)))
+            span.end()
+            raise
+
+    return traced_method
+
+
+def get_llm_model(instance):
+    llm_model = "unknown"
+    if hasattr(instance, "_model_id"):
+        llm_model = instance._model_id
+    if hasattr(instance, "_model_name"):
+        llm_model = instance._model_name.replace("models/", "")
+    return llm_model
+
+
+def serialize_prompts(args, kwargs, instance):
+    prompts = []
+    if hasattr(instance, "_system_instruction") and instance._system_instruction is not None:
+        system_prompt = {
+            "role": "system",
+            "content": instance._system_instruction.__dict__["_pb"].parts[0].text,
+        }
+        prompts.append(system_prompt)
+
+    if args is not None and len(args) > 0:
+        content = ""
+        for arg in args:
+            if isinstance(arg, str):
+                content = f"{content}{arg}\n"
+            elif isinstance(arg, list):
+                for subarg in arg:
+                    content = f"{content}{subarg}\n"
+        prompts.append({"role": "user", "content": content})
+    return prompts
+
+
+def set_response_attributes(
+    span: Span,
+    result,
+):
+    span.set_status(Status(StatusCode.OK))
+    if hasattr(result, "text"):
+        set_event_completion(span, [{"role": "assistant", "content": result.text}])
+
+    if hasattr(result, "usage_metadata"):
+        usage = result.usage_metadata
+        input_tokens = usage.prompt_token_count
+        output_tokens = usage.candidates_token_count
+        set_usage_attributes(
+            span, {"input_tokens": input_tokens, "output_tokens": output_tokens}
+        )
+
+
+def build_streaming_response(span, response):
+    complete_response = ""
+    for item in response:
+        item_to_yield = item
+        complete_response += str(item.text)
+        yield item_to_yield
+        set_event_completion_chunk(span, item.text)
+        if hasattr(item, "usage_metadata"):
+            usage = item.usage_metadata
+            input_tokens = usage.prompt_token_count
+            output_tokens = usage.candidates_token_count
+            set_usage_attributes(
+                span, {"input_tokens": input_tokens, "output_tokens": output_tokens}
+            )
+
+    set_response_attributes(span, response)
+    span.set_status(Status(StatusCode.OK))
+    span.end()
+
+
+async def abuild_streaming_response(span, response):
+    complete_response = ""
+    async for item in response:
+        item_to_yield = item
+        complete_response += str(item.text)
+        yield item_to_yield
+        set_event_completion_chunk(span, item.text)
+        if hasattr(item, "usage_metadata"):
+            usage = item.usage_metadata
+            input_tokens = usage.prompt_token_count
+            output_tokens = usage.candidates_token_count
+            set_usage_attributes(
+                span, {"input_tokens": input_tokens, "output_tokens": output_tokens}
+            )
+
+    set_response_attributes(span, response)
+    span.set_status(Status(StatusCode.OK))
+    span.end()