lmnr 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

lmnr/__init__.py

@@ -6,11 +6,7 @@ from .sdk.laminar import Laminar
 from .sdk.types import (
     AgentOutput,
     FinalOutputChunkContent,
-    ChatMessage,
     HumanEvaluator,
-    NodeInput,
-    PipelineRunError,
-    PipelineRunResponse,
     RunAgentResponseChunk,
     StepChunkContent,
     TracingLevel,
@@ -25,7 +21,6 @@ __all__ = [
     "AgentOutput",
     "AsyncLaminarClient",
     "Attributes",
-    "ChatMessage",
     "EvaluationDataset",
     "FinalOutputChunkContent",
     "HumanEvaluator",
@@ -34,9 +29,6 @@ __all__ = [
     "LaminarClient",
     "LaminarDataset",
     "LaminarSpanContext",
-    "NodeInput",
-    "PipelineRunError",
-    "PipelineRunResponse",
     "RunAgentResponseChunk",
     "StepChunkContent",
     "TracingLevel",

lmnr/openllmetry_sdk/__init__.py

@@ -69,3 +69,7 @@ class TracerManager:
     @staticmethod
     def flush() -> bool:
         return TracerManager.__tracer_wrapper.flush()
+
+    @staticmethod
+    def shutdown():
+        TracerManager.__tracer_wrapper.shutdown()

lmnr/openllmetry_sdk/decorators/base.py

@@ -1,7 +1,6 @@
 import json
 from functools import wraps
 import logging
-import os
 import pydantic
 import types
 from typing import Any, Literal, Optional, Union
@@ -40,6 +39,7 @@ def json_dumps(data: dict) -> str:
 def entity_method(
     name: Optional[str] = None,
     ignore_input: bool = False,
+    ignore_inputs: Optional[list[str]] = None,
     ignore_output: bool = False,
     span_type: Union[Literal["DEFAULT"], Literal["LLM"], Literal["TOOL"]] = "DEFAULT",
 ):
@@ -58,9 +58,15 @@ def entity_method(
                 ctx_token = context_api.attach(ctx)
 
                 try:
-                    if _should_send_prompts() and not ignore_input:
+                    if not ignore_input:
                         inp = json_dumps(
-                            get_input_from_func_args(fn, is_method(fn), args, kwargs)
+                            get_input_from_func_args(
+                                fn,
+                                is_method=is_method(fn),
+                                func_args=args,
+                                func_kwargs=kwargs,
+                                ignore_inputs=ignore_inputs,
+                            )
                         )
                         if len(inp) > MAX_MANUAL_SPAN_PAYLOAD_SIZE:
                             span.set_attribute(
@@ -83,7 +89,7 @@ def entity_method(
                     return _handle_generator(span, res)
 
                 try:
-                    if _should_send_prompts() and not ignore_output:
+                    if not ignore_output:
                         output = json_dumps(res)
                         if len(output) > MAX_MANUAL_SPAN_PAYLOAD_SIZE:
                             span.set_attribute(
@@ -108,6 +114,7 @@ def entity_method(
 def aentity_method(
     name: Optional[str] = None,
     ignore_input: bool = False,
+    ignore_inputs: Optional[list[str]] = None,
     ignore_output: bool = False,
     span_type: Union[Literal["DEFAULT"], Literal["LLM"], Literal["TOOL"]] = "DEFAULT",
 ):
@@ -126,9 +133,15 @@ def aentity_method(
                 ctx_token = context_api.attach(ctx)
 
                 try:
-                    if _should_send_prompts() and not ignore_input:
+                    if not ignore_input:
                         inp = json_dumps(
-                            get_input_from_func_args(fn, is_method(fn), args, kwargs)
+                            get_input_from_func_args(
+                                fn,
+                                is_method=is_method(fn),
+                                func_args=args,
+                                func_kwargs=kwargs,
+                                ignore_inputs=ignore_inputs,
+                            )
                         )
                         if len(inp) > MAX_MANUAL_SPAN_PAYLOAD_SIZE:
                             span.set_attribute(
@@ -151,7 +164,7 @@ def aentity_method(
                     return await _ahandle_generator(span, ctx_token, res)
 
                 try:
-                    if _should_send_prompts() and not ignore_output:
+                    if not ignore_output:
                         output = json_dumps(res)
                         if len(output) > MAX_MANUAL_SPAN_PAYLOAD_SIZE:
                             span.set_attribute(
@@ -192,12 +205,6 @@ async def _ahandle_generator(span, ctx_token, res):
     context_api.detach(ctx_token)
 
 
-def _should_send_prompts():
-    return (
-        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
-    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
-
-
 def _process_exception(span: Span, e: Exception):
     # Note that this `escaped` is sent as a StringValue("True"), not a boolean.
     span.record_exception(e, escaped=True)

lmnr/openllmetry_sdk/instruments.py

@@ -11,6 +11,7 @@ class Instruments(Enum):
     CHROMA = "chroma"
     COHERE = "cohere"
     GOOGLE_GENERATIVEAI = "google_generativeai"
+    GOOGLE_GENAI = "google_genai"
     GROQ = "groq"
     HAYSTACK = "haystack"
     LANCEDB = "lancedb"

lmnr/openllmetry_sdk/opentelemetry/instrumentation/google_genai/__init__.py

@@ -0,0 +1,454 @@
+"""OpenTelemetry Google Generative AI API instrumentation"""
+
+from collections import defaultdict
+import logging
+import os
+from typing import AsyncGenerator, Callable, Collection, Generator, Optional
+
+from google.genai import types
+
+from .config import (
+    Config,
+)
+from .utils import (
+    dont_throw,
+    role_from_content_union,
+    set_span_attribute,
+    process_content_union,
+    to_dict,
+    with_tracer_wrapper,
+)
+from opentelemetry.trace import Tracer
+from wrapt import wrap_function_wrapper
+
+from opentelemetry import context as context_api
+from opentelemetry.trace import get_tracer, SpanKind, Span
+from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
+
+from opentelemetry.semconv_ai import (
+    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
+    SpanAttributes,
+    LLMRequestTypeValues,
+)
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("google-genai >= 1.0.0",)
+
+WRAPPED_METHODS = [
+    {
+        "package": "google.genai.models",
+        "object": "Models",
+        "method": "generate_content",
+        "span_name": "gemini.generate_content",
+        "is_streaming": False,
+        "is_async": False,
+    },
+    {
+        "package": "google.genai.models",
+        "object": "AsyncModels",
+        "method": "generate_content",
+        "span_name": "gemini.generate_content",
+        "is_streaming": False,
+        "is_async": True,
+    },
+    {
+        "package": "google.genai.models",
+        "object": "Models",
+        "method": "generate_content_stream",
+        "span_name": "gemini.generate_content_stream",
+        "is_streaming": True,
+        "is_async": False,
+    },
+    {
+        "package": "google.genai.models",
+        "object": "AsyncModels",
+        "method": "generate_content_stream",
+        "span_name": "gemini.generate_content_stream",
+        "is_streaming": True,
+        "is_async": True,
+    },
+]
+
+
+def should_send_prompts():
+    return (
+        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
+    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
+
+
+@dont_throw
+def _set_request_attributes(span, args, kwargs):
+    config_dict = to_dict(kwargs.get("config", {}))
+    set_span_attribute(
+        span, gen_ai_attributes.GEN_AI_REQUEST_MODEL, kwargs.get("model")
+    )
+    set_span_attribute(
+        span,
+        gen_ai_attributes.GEN_AI_REQUEST_TEMPERATURE,
+        config_dict.get("temperature"),
+    )
+    set_span_attribute(
+        span, gen_ai_attributes.GEN_AI_REQUEST_TOP_P, config_dict.get("top_p")
+    )
+    set_span_attribute(
+        span, gen_ai_attributes.GEN_AI_REQUEST_TOP_K, config_dict.get("top_k")
+    )
+    set_span_attribute(
+        span,
+        gen_ai_attributes.GEN_AI_REQUEST_CHOICE_COUNT,
+        config_dict.get("candidate_count"),
+    )
+    set_span_attribute(
+        span,
+        gen_ai_attributes.GEN_AI_REQUEST_MAX_TOKENS,
+        config_dict.get("max_output_tokens"),
+    )
+    set_span_attribute(
+        span,
+        gen_ai_attributes.GEN_AI_REQUEST_STOP_SEQUENCES,
+        config_dict.get("stop_sequences"),
+    )
+    set_span_attribute(
+        span,
+        gen_ai_attributes.GEN_AI_REQUEST_FREQUENCY_PENALTY,
+        config_dict.get("frequency_penalty"),
+    )
+    set_span_attribute(
+        span,
+        gen_ai_attributes.GEN_AI_REQUEST_PRESENCE_PENALTY,
+        config_dict.get("presence_penalty"),
+    )
+    set_span_attribute(
+        span, gen_ai_attributes.GEN_AI_REQUEST_SEED, config_dict.get("seed")
+    )
+
+    tools: list[types.FunctionDeclaration] = []
+    if kwargs.get("tools"):
+        for tool in kwargs.get("tools"):
+            if isinstance(tool, types.Tool):
+                tools += tool.function_declarations or []
+            elif isinstance(tool, Callable):
+                tools.append(types.FunctionDeclaration.from_callable(tool))
+    for tool_num, tool in enumerate(tools):
+        set_span_attribute(
+            span,
+            f"{SpanAttributes.LLM_REQUEST_FUNCTIONS}.{tool_num}.name",
+            to_dict(tool).get("name"),
+        )
+        set_span_attribute(
+            span,
+            f"{SpanAttributes.LLM_REQUEST_FUNCTIONS}.{tool_num}.description",
+            to_dict(tool).get("description"),
+        )
+        set_span_attribute(
+            span,
+            f"{SpanAttributes.LLM_REQUEST_FUNCTIONS}.{tool_num}.parameters",
+            to_dict(tool).get("parameters"),
+        )
+
+    if should_send_prompts():
+        i = 0
+        system_instruction: Optional[types.ContentUnion] = config_dict.get(
+            "system_instruction"
+        )
+        if system_instruction:
+            set_span_attribute(
+                span,
+                f"{gen_ai_attributes.GEN_AI_PROMPT}.{i}.content",
+                process_content_union(system_instruction),
+            )
+            set_span_attribute(
+                span, f"{gen_ai_attributes.GEN_AI_PROMPT}.{i}.role", "system"
+            )
+            i += 1
+        contents = kwargs.get("contents", [])
+        if not isinstance(contents, list):
+            contents = [contents]
+        for content in contents:
+            set_span_attribute(
+                span,
+                f"{gen_ai_attributes.GEN_AI_PROMPT}.{i}.content",
+                process_content_union(content),
+            )
+            set_span_attribute(
+                span,
+                f"{gen_ai_attributes.GEN_AI_PROMPT}.{i}.role",
+                role_from_content_union(content) or "user",
+            )
+            i += 1
+
+
+@dont_throw
+def _set_response_attributes(span, response: types.GenerateContentResponse):
+    candidates = response.candidates or []
+    set_span_attribute(
+        span, gen_ai_attributes.GEN_AI_RESPONSE_ID, to_dict(response).get("response_id")
+    )
+    set_span_attribute(
+        span,
+        gen_ai_attributes.GEN_AI_RESPONSE_MODEL,
+        to_dict(response).get("model_version"),
+    )
+
+    if response.usage_metadata:
+        usage_dict = to_dict(response.usage_metadata)
+        set_span_attribute(
+            span,
+            gen_ai_attributes.GEN_AI_USAGE_INPUT_TOKENS,
+            usage_dict.get("prompt_token_count"),
+        )
+        set_span_attribute(
+            span,
+            gen_ai_attributes.GEN_AI_USAGE_OUTPUT_TOKENS,
+            usage_dict.get("candidates_token_count"),
+        )
+        set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
+            usage_dict.get("total_token_count"),
+        )
+        set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS,
+            usage_dict.get("cached_content_token_count"),
+        )
+
+    if should_send_prompts():
+        if len(candidates) > 1:
+            for i, candidate in enumerate(candidates):
+                set_span_attribute(
+                    span,
+                    f"{gen_ai_attributes.GEN_AI_COMPLETION}.{i}.content",
+                    process_content_union(candidate.content),
+                )
+                set_span_attribute(
+                    span, f"{gen_ai_attributes.GEN_AI_COMPLETION}.{i}.role", "assistant"
+                )
+        else:
+            set_span_attribute(
+                span, f"{gen_ai_attributes.GEN_AI_COMPLETION}.0.content", response.text
+            )
+            set_span_attribute(
+                span, f"{gen_ai_attributes.GEN_AI_COMPLETION}.0.role", "assistant"
+            )
+
+
+@dont_throw
+def _build_from_streaming_response(
+    span: Span, response: Generator[types.GenerateContentResponse, None, None]
+) -> Generator[types.GenerateContentResponse, None, None]:
+    final_parts = []
+    role = "model"
+    aggregated_usage_metadata = defaultdict(int)
+    model_version = None
+    for chunk in response:
+        if chunk.model_version:
+            model_version = chunk.model_version
+
+        if chunk.candidates:
+            # Currently gemini throws an error if you pass more than one candidate
+            # with streaming
+            if chunk.candidates and len(chunk.candidates) > 0:
+                final_parts += chunk.candidates[0].content.parts or []
+                role = chunk.candidates[0].content.role or role
+        if chunk.usage_metadata:
+            usage_dict = to_dict(chunk.usage_metadata)
+            # prompt token count is sent in every chunk
+            # (and is less by 1 in the last chunk, so we set it once);
+            # total token count in every chunk is greater by prompt token count than it should be,
+            # thus this awkward logic here
+            if aggregated_usage_metadata.get("prompt_token_count") is None:
+                aggregated_usage_metadata["prompt_token_count"] = (
+                    usage_dict.get("prompt_token_count") or 0
+                )
+                aggregated_usage_metadata["total_token_count"] = (
+                    usage_dict.get("total_token_count") or 0
+                )
+            aggregated_usage_metadata["candidates_token_count"] += (
+                usage_dict.get("candidates_token_count") or 0
+            )
+            aggregated_usage_metadata["total_token_count"] += (
+                usage_dict.get("candidates_token_count") or 0
+            )
+        yield chunk
+
+    compound_response = types.GenerateContentResponse(
+        candidates=[
+            {
+                "content": {
+                    "parts": final_parts,
+                    "role": role,
+                },
+            }
+        ],
+        usage_metadata=types.GenerateContentResponseUsageMetadataDict(
+            **aggregated_usage_metadata
+        ),
+        model_version=model_version,
+    )
+    if span.is_recording():
+        _set_response_attributes(span, compound_response)
+    span.end()
+
+
+@dont_throw
+async def _abuild_from_streaming_response(
+    span: Span, response: AsyncGenerator[types.GenerateContentResponse, None]
+) -> AsyncGenerator[types.GenerateContentResponse, None]:
+    final_parts = []
+    role = "model"
+    aggregated_usage_metadata = defaultdict(int)
+    model_version = None
+    async for chunk in response:
+        if chunk.candidates:
+            # Currently gemini throws an error if you pass more than one candidate
+            # with streaming
+            if chunk.candidates and len(chunk.candidates) > 0:
+                final_parts += chunk.candidates[0].content.parts or []
+                role = chunk.candidates[0].content.role or role
+                if chunk.model_version:
+                    model_version = chunk.model_version
+        if chunk.usage_metadata:
+            usage_dict = to_dict(chunk.usage_metadata)
+            # prompt token count is sent in every chunk
+            # (and is less by 1 in the last chunk, so we set it once);
+            # total token count in every chunk is greater by prompt token count than it should be,
+            # thus this awkward logic here
+            if aggregated_usage_metadata.get("prompt_token_count") is None:
+                aggregated_usage_metadata["prompt_token_count"] = usage_dict.get(
+                    "prompt_token_count"
+                )
+                aggregated_usage_metadata["total_token_count"] = usage_dict.get(
+                    "total_token_count"
+                )
+            aggregated_usage_metadata["candidates_token_count"] += (
+                usage_dict.get("candidates_token_count") or 0
+            )
+            aggregated_usage_metadata["total_token_count"] += (
+                usage_dict.get("candidates_token_count") or 0
+            )
+        yield chunk
+
+    compound_response = types.GenerateContentResponse(
+        candidates=[
+            {
+                "content": {
+                    "parts": final_parts,
+                    "role": role,
+                },
+            }
+        ],
+        usage_metadata=types.GenerateContentResponseUsageMetadataDict(
+            **aggregated_usage_metadata
+        ),
+        model_version=model_version,
+    )
+    if span.is_recording():
+        _set_response_attributes(span, compound_response)
+    span.end()
+
+
+@with_tracer_wrapper
+def _wrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) or context_api.get_value(
+        SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
+    ):
+        return wrapped(*args, **kwargs)
+
+    span = tracer.start_span(
+        to_wrap.get("span_name"),
+        kind=SpanKind.CLIENT,
+        attributes={
+            SpanAttributes.LLM_SYSTEM: "gemini",
+            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
+        },
+    )
+
+    if span.is_recording():
+        _set_request_attributes(span, args, kwargs)
+
+    if to_wrap.get("is_streaming"):
+        return _build_from_streaming_response(span, wrapped(*args, **kwargs))
+    else:
+        response = wrapped(*args, **kwargs)
+
+    if span.is_recording():
+        _set_response_attributes(span, response)
+
+    span.end()
+    return response
+
+
+@with_tracer_wrapper
+async def _awrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY) or context_api.get_value(
+        SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
+    ):
+        return await wrapped(*args, **kwargs)
+
+    span = tracer.start_span(
+        to_wrap.get("span_name"),
+        kind=SpanKind.CLIENT,
+        attributes={
+            SpanAttributes.LLM_SYSTEM: "gemini",
+            SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
+        },
+    )
+
+    if span.is_recording():
+        _set_request_attributes(span, args, kwargs)
+
+    if to_wrap.get("is_streaming"):
+        return _abuild_from_streaming_response(span, await wrapped(*args, **kwargs))
+    else:
+        response = await wrapped(*args, **kwargs)
+
+    if span.is_recording():
+        _set_response_attributes(span, response)
+
+    span.end()
+    return response
+
+
+class GoogleGenAiSdkInstrumentor(BaseInstrumentor):
+    """An instrumentor for Google GenAI's client library."""
+
+    def __init__(
+        self,
+        exception_logger=None,
+        upload_base64_image=None,
+        convert_image_to_openai_format=True,
+    ):
+        super().__init__()
+        Config.exception_logger = exception_logger
+        Config.upload_base64_image = upload_base64_image
+        Config.convert_image_to_openai_format = convert_image_to_openai_format
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "0.0.1a0", tracer_provider)
+
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_function_wrapper(
+                wrapped_method.get("package"),
+                f"{wrapped_method.get('object')}.{wrapped_method.get('method')}",
+                (
+                    _awrap(tracer, wrapped_method)
+                    if wrapped_method.get("is_async")
+                    else _wrap(tracer, wrapped_method)
+                ),
+            )
+
+    def _uninstrument(self, **kwargs):
+        for wrapped_method in WRAPPED_METHODS:
+            unwrap(
+                f"{wrapped_method.get('package')}.{wrapped_method.get('object')}",
+                wrapped_method.get("method"),
+            )

lmnr/openllmetry_sdk/opentelemetry/instrumentation/google_genai/config.py

@@ -0,0 +1,9 @@
+from typing import Callable, Coroutine, Optional
+
+
+class Config:
+    exception_logger = None
+    upload_base64_image: Optional[
+        Callable[[str, str, str, str], Coroutine[None, None, str]]
+    ] = None
+    convert_image_to_openai_format: bool = True