lmnr 0.5.1a0__py3-none-any.whl → 0.5.3__py3-none-any.whl

lmnr/opentelemetry_lib/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/opentelemetry_lib/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

lmnr/opentelemetry_lib/opentelemetry/instrumentation/google_genai/utils.py

@@ -0,0 +1,216 @@
+import logging
+import traceback
+import json
+
+from .config import (
+    Config,
+)
+from google.genai import types
+from google.genai._common import BaseModel
+import pydantic
+from opentelemetry.trace import Span
+from typing import Any, Optional, Union
+
+
+def set_span_attribute(span: Span, name: str, value: str):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+
+def dont_throw(func):
+    """
+    A decorator that wraps the passed in function and logs exceptions instead of throwing them.
+
+    @param func: The function to wrap
+    @return: The wrapper function
+    """
+    # Obtain a logger specific to the function's module
+    logger = logging.getLogger(func.__module__)
+
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            logger.debug(
+                "Laminar failed to trace in %s, error: %s",
+                func.__name__,
+                traceback.format_exc(),
+            )
+            if Config.exception_logger:
+                Config.exception_logger(e)
+
+    return wrapper
+
+
+def to_dict(obj: Union[BaseModel, pydantic.BaseModel, dict]) -> dict[str, Any]:
+    try:
+        if isinstance(obj, BaseModel):
+            return obj.model_dump()
+        elif isinstance(obj, pydantic.BaseModel):
+            return obj.model_dump()
+        elif isinstance(obj, dict):
+            return obj
+        else:
+            return dict(obj)
+    except Exception:
+        return dict(obj)
+
+
+def process_content_union(
+    content: Union[types.ContentUnion, types.ContentUnionDict],
+    trace_id: Optional[str] = None,
+    span_id: Optional[str] = None,
+    message_index: int = 0,
+) -> Optional[str]:
+    parts = _process_content_union(content, trace_id, span_id, message_index)
+    if parts is None:
+        return None
+    if isinstance(parts, str):
+        return parts
+    elif isinstance(parts, list):
+        if len(parts) == 1 and isinstance(parts[0], str):
+            return parts[0]
+        return json.dumps(
+            [
+                {"type": "text", "text": part} if isinstance(part, str) else part
+                for part in parts
+            ]
+        )
+    else:
+        return None
+
+
+def _process_content_union(
+    content: Union[types.ContentUnion, types.ContentUnionDict],
+    trace_id: Optional[str] = None,
+    span_id: Optional[str] = None,
+    message_index: int = 0,
+) -> Union[str, list[str], None]:
+    if isinstance(content, types.Content):
+        parts = to_dict(content).get("parts", [])
+        return [_process_part(part) for part in parts]
+    elif isinstance(content, list):
+        return [_process_part_union(item) for item in content]
+    elif isinstance(content, (types.Part, types.File, str)):
+        return _process_part_union(content)
+    elif isinstance(content, dict):
+        if "parts" in content:
+            return [
+                _process_part_union(
+                    item, trace_id, span_id, message_index, content_index
+                )
+                for content_index, item in enumerate(content.get("parts", []))
+            ]
+        else:
+            # Assume it's PartDict
+            return _process_part_union(content, trace_id, span_id, message_index)
+    else:
+        return None
+
+
+def _process_part_union(
+    content: Union[types.PartDict, types.File, types.Part, str],
+    trace_id: Optional[str] = None,
+    span_id: Optional[str] = None,
+    message_index: int = 0,
+    content_index: int = 0,
+) -> Optional[str]:
+    if isinstance(content, str):
+        return content
+    elif isinstance(content, types.File):
+        content_dict = to_dict(content)
+        name = (
+            content_dict.get("name")
+            or content_dict.get("display_name")
+            or content_dict.get("uri")
+        )
+        return f"files/{name}"
+    elif isinstance(content, (types.Part, dict)):
+        return _process_part(content, trace_id, span_id, message_index, content_index)
+    else:
+        return None
+
+
+def _process_part(
+    content: types.Part,
+    trace_id: Optional[str] = None,
+    span_id: Optional[str] = None,
+    message_index: int = 0,
+    content_index: int = 0,
+) -> Optional[str]:
+    part_dict = to_dict(content)
+    if part_dict.get("text") is not None:
+        return part_dict.get("text")
+    elif part_dict.get("inline_data"):
+        blob = to_dict(part_dict.get("inline_data"))
+        if blob.get("mime_type").startswith("image/"):
+            return _process_image_item(
+                blob, trace_id, span_id, message_index, content_index
+            )
+        else:
+            # currently, only images are supported
+            return blob.get("mime_type") or "unknown_media"
+    else:
+        return None
+
+
+def role_from_content_union(
+    content: Union[types.ContentUnion, types.ContentUnionDict],
+) -> Optional[str]:
+    if isinstance(content, types.Content):
+        return to_dict(content).get("role")
+    elif isinstance(content, list) and len(content) > 0:
+        return role_from_content_union(content[0])
+    else:
+        return None
+
+
+def with_tracer_wrapper(func):
+    """Helper for providing tracer for wrapper functions."""
+
+    def _with_tracer(tracer, to_wrap):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer
+
+
+def _run_async(method):
+    import asyncio
+    import threading
+
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    if loop and loop.is_running():
+        thread = threading.Thread(target=lambda: asyncio.run(method))
+        thread.start()
+        thread.join()
+    else:
+        asyncio.run(method)
+
+
+def _process_image_item(
+    blob: dict[str, Any],
+    trace_id: str,
+    span_id: str,
+    message_index: int,
+    content_index: int,
+):
+    # Convert to openai format, so backends can handle it
+    return (
+        {
+            "type": "image_url",
+            "image_url": {
+                "url": f"data:image/{blob.get('mime_type').split('/')[1]};base64,{blob.get('data')}",
+            },
+        }
+        if Config.convert_image_to_openai_format
+        else blob
+    )

lmnr/opentelemetry_lib/tracing/__init__.py

@@ -0,0 +1 @@
+from lmnr.opentelemetry_lib.tracing.context_manager import get_tracer

lmnr/opentelemetry_lib/tracing/context_manager.py

@@ -0,0 +1,13 @@
+from contextlib import contextmanager
+
+from lmnr.opentelemetry_lib.tracing.tracing import TracerWrapper
+
+
+@contextmanager
+def get_tracer(flush_on_exit: bool = False):
+    wrapper = TracerWrapper()
+    try:
+        yield wrapper.get_tracer()
+    finally:
+        if flush_on_exit:
+            wrapper.flush()