lmnr 0.6.16__py3-none-any.whl → 0.7.26__py3-none-any.whl

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

@@ -8,15 +8,24 @@ from typing import AsyncGenerator, Callable, Collection, Generator
 
 from google.genai import types
 
+from lmnr.opentelemetry_lib.tracing.context import (
+    get_current_context,
+    get_event_attributes_from_context,
+)
+from lmnr.sdk.utils import json_dumps
+
 from .config import (
     Config,
 )
+from .schema_utils import SchemaJSONEncoder, process_schema
 from .utils import (
     dont_throw,
     get_content,
+    merge_text_parts,
+    process_content_union,
+    process_stream_chunk,
     role_from_content_union,
     set_span_attribute,
-    process_content_union,
     to_dict,
     with_tracer_wrapper,
 )
@@ -24,8 +33,9 @@ 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.trace import get_tracer, SpanKind, Span, Status, StatusCode
 from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
+from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
@@ -78,7 +88,7 @@ WRAPPED_METHODS = [
 
 def should_send_prompts():
     return (
-        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
+        os.getenv("LAMINAR_TRACE_CONTENT") or "true"
     ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
 
 
@@ -128,6 +138,25 @@ def _set_request_attributes(span, args, kwargs):
         span, gen_ai_attributes.GEN_AI_REQUEST_SEED, config_dict.get("seed")
     )
 
+    if schema := config_dict.get("response_schema"):
+        try:
+            set_span_attribute(
+                span,
+                SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA,
+                json.dumps(process_schema(schema), cls=SchemaJSONEncoder),
+            )
+        except Exception:
+            pass
+    elif json_schema := config_dict.get("response_json_schema"):
+        try:
+            set_span_attribute(
+                span,
+                SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA,
+                json_dumps(json_schema),
+            )
+        except Exception:
+            pass
+
     tools: list[types.FunctionDeclaration] = []
     arg_tools = config_dict.get("tools", kwargs.get("tools"))
     if arg_tools:
@@ -152,7 +181,7 @@ def _set_request_attributes(span, args, kwargs):
         set_span_attribute(
             span,
             f"{SpanAttributes.LLM_REQUEST_FUNCTIONS}.{tool_num}.parameters",
-            json.dumps(tool_dict.get("parameters")),
+            json_dumps(tool_dict.get("parameters")),
         )
 
     if should_send_prompts():
@@ -177,15 +206,17 @@ def _set_request_attributes(span, args, kwargs):
             contents = [contents]
         for content in contents:
             processed_content = process_content_union(content)
-            content_str = get_content(processed_content)
+            content_payload = get_content(processed_content)
+            if isinstance(content_payload, dict):
+                content_payload = [content_payload]
 
             set_span_attribute(
                 span,
                 f"{gen_ai_attributes.GEN_AI_PROMPT}.{i}.content",
                 (
-                    content_str
-                    if isinstance(content_str, str)
-                    else json.dumps(content_str)
+                    content_payload
+                    if isinstance(content_payload, str)
+                    else json_dumps(content_payload)
                 ),
             )
             blocks = (
@@ -218,7 +249,7 @@ def _set_request_attributes(span, args, kwargs):
                 set_span_attribute(
                     span,
                     f"{gen_ai_attributes.GEN_AI_PROMPT}.{i}.tool_calls.{tool_call_index}.arguments",
-                    json.dumps(function_call.get("arguments")),
+                    json_dumps(function_call.get("arguments")),
                 )
                 tool_call_index += 1
 
@@ -244,6 +275,16 @@ def _set_response_attributes(span, response: types.GenerateContentResponse):
 
     if response.usage_metadata:
         usage_dict = to_dict(response.usage_metadata)
+        candidates_token_count = usage_dict.get("candidates_token_count")
+        # unlike OpenAI, and unlike input cached tokens, thinking tokens are
+        # not counted as part of candidates token count, so we need to add them
+        # separately for consistency with other instrumentations
+        thoughts_token_count = usage_dict.get("thoughts_token_count")
+        output_token_count = (
+            (candidates_token_count or 0) + (thoughts_token_count or 0)
+            if candidates_token_count is not None or thoughts_token_count is not None
+            else None
+        )
         set_span_attribute(
             span,
             gen_ai_attributes.GEN_AI_USAGE_INPUT_TOKENS,
@@ -252,7 +293,7 @@ def _set_response_attributes(span, response: types.GenerateContentResponse):
         set_span_attribute(
             span,
             gen_ai_attributes.GEN_AI_USAGE_OUTPUT_TOKENS,
-            usage_dict.get("candidates_token_count"),
+            output_token_count,
         )
         set_span_attribute(
             span,
@@ -264,28 +305,39 @@ def _set_response_attributes(span, response: types.GenerateContentResponse):
             SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS,
             usage_dict.get("cached_content_token_count"),
         )
+        set_span_attribute(
+            span,
+            SpanAttributes.LLM_USAGE_REASONING_TOKENS,
+            thoughts_token_count,
+        )
 
     if should_send_prompts():
         set_span_attribute(
             span, f"{gen_ai_attributes.GEN_AI_COMPLETION}.0.role", "model"
         )
         candidates_list = candidates if isinstance(candidates, list) else [candidates]
-        for i, candidate in enumerate(candidates_list):
+        i = 0
+        for candidate in candidates_list:
+            has_content = False
             processed_content = process_content_union(candidate.content)
-            content_str = get_content(processed_content)
+            content_payload = get_content(processed_content)
+            if isinstance(content_payload, dict):
+                content_payload = [content_payload]
 
             set_span_attribute(
                 span, f"{gen_ai_attributes.GEN_AI_COMPLETION}.{i}.role", "model"
             )
-            set_span_attribute(
-                span,
-                f"{gen_ai_attributes.GEN_AI_COMPLETION}.{i}.content",
-                (
-                    content_str
-                    if isinstance(content_str, str)
-                    else json.dumps(content_str)
-                ),
-            )
+            if content_payload:
+                has_content = True
+                set_span_attribute(
+                    span,
+                    f"{gen_ai_attributes.GEN_AI_COMPLETION}.{i}.content",
+                    (
+                        content_payload
+                        if isinstance(content_payload, str)
+                        else json_dumps(content_payload)
+                    ),
+                )
             blocks = (
                 processed_content
                 if isinstance(processed_content, list)
@@ -298,6 +350,7 @@ def _set_response_attributes(span, response: types.GenerateContentResponse):
                 if not block_dict.get("function_call"):
                     continue
                 function_call = to_dict(block_dict.get("function_call", {}))
+                has_content = True
                 set_span_attribute(
                     span,
                     f"{gen_ai_attributes.GEN_AI_COMPLETION}.{i}.tool_calls.{tool_call_index}.name",
@@ -315,9 +368,11 @@ def _set_response_attributes(span, response: types.GenerateContentResponse):
                 set_span_attribute(
                     span,
                     f"{gen_ai_attributes.GEN_AI_COMPLETION}.{i}.tool_calls.{tool_call_index}.arguments",
-                    json.dumps(function_call.get("arguments")),
+                    json_dumps(function_call.get("arguments")),
                 )
                 tool_call_index += 1
+            if has_content:
+                i += 1
 
 
 @dont_throw
@@ -329,53 +384,49 @@ def _build_from_streaming_response(
     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
-            )
+        try:
+            span.add_event("llm.content.completion.chunk")
+        except Exception:
+            pass
+        # Important: do all processing in a separate sync function, that is
+        # wrapped in @dont_throw. If we did it here, the @dont_throw on top of
+        # this function would not be able to catch the errors, as they are
+        # raised later, after the generator is returned, and when it is being
+        # consumed.
+        chunk_result = process_stream_chunk(
+            chunk,
+            role,
+            model_version,
+            aggregated_usage_metadata,
+            final_parts,
+        )
+        # even though process_stream_chunk can't return None, the result can be
+        # None, if the processing throws an error (see @dont_throw)
+        if chunk_result:
+            role = chunk_result["role"]
+            model_version = chunk_result["model_version"]
         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()
+    try:
+        compound_response = types.GenerateContentResponse(
+            candidates=[
+                {
+                    "content": {
+                        "parts": merge_text_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)
+    finally:
+        if span.is_recording():
+            span.end()
 
 
 @dont_throw
@@ -387,52 +438,49 @@ async def _abuild_from_streaming_response(
     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
-            )
+        try:
+            span.add_event("llm.content.completion.chunk")
+        except Exception:
+            pass
+        # Important: do all processing in a separate sync function, that is
+        # wrapped in @dont_throw. If we did it here, the @dont_throw on top of
+        # this function would not be able to catch the errors, as they are
+        # raised later, after the generator is returned, and when it is being
+        # consumed.
+        chunk_result = process_stream_chunk(
+            chunk,
+            role,
+            model_version,
+            aggregated_usage_metadata,
+            final_parts,
+        )
+        # even though process_stream_chunk can't return None, the result can be
+        # None, if the processing throws an error (see @dont_throw)
+        if chunk_result:
+            role = chunk_result["role"]
+            model_version = chunk_result["model_version"]
         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()
+    try:
+        compound_response = types.GenerateContentResponse(
+            candidates=[
+                {
+                    "content": {
+                        "parts": merge_text_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)
+    finally:
+        if span.is_recording():
+            span.end()
 
 
 @with_tracer_wrapper
@@ -449,21 +497,27 @@ def _wrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
             SpanAttributes.LLM_SYSTEM: "gemini",
             SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
         },
+        context=get_current_context(),
     )
 
     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:
+    try:
         response = wrapped(*args, **kwargs)
-
-    if span.is_recording():
-        _set_response_attributes(span, response)
-
-    span.end()
-    return response
+        if to_wrap.get("is_streaming"):
+            return _build_from_streaming_response(span, response)
+        if span.is_recording():
+            _set_response_attributes(span, response)
+        span.end()
+        return response
+    except Exception as e:
+        attributes = get_event_attributes_from_context()
+        span.set_attribute(ERROR_TYPE, e.__class__.__name__)
+        span.record_exception(e, attributes=attributes)
+        span.set_status(Status(StatusCode.ERROR, str(e)))
+        span.end()
+        raise
 
 
 @with_tracer_wrapper
@@ -480,21 +534,29 @@ async def _awrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
             SpanAttributes.LLM_SYSTEM: "gemini",
             SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
         },
+        context=get_current_context(),
     )
 
     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:
+    try:
         response = await wrapped(*args, **kwargs)
-
-    if span.is_recording():
-        _set_response_attributes(span, response)
-
-    span.end()
-    return response
+        if to_wrap.get("is_streaming"):
+            return _abuild_from_streaming_response(span, response)
+        else:
+            if span.is_recording():
+                _set_response_attributes(span, response)
+
+            span.end()
+            return response
+    except Exception as e:
+        attributes = get_event_attributes_from_context()
+        span.set_attribute(ERROR_TYPE, e.__class__.__name__)
+        span.record_exception(e, attributes=attributes)
+        span.set_status(Status(StatusCode.ERROR, str(e)))
+        span.end()
+        raise
 
 
 class GoogleGenAiSdkInstrumentor(BaseInstrumentor):

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

@@ -0,0 +1,26 @@
+from typing import Any
+from google.genai._api_client import BaseApiClient
+from google.genai._transformers import t_schema
+from google.genai.types import JSONSchemaType
+
+import json
+
+DUMMY_CLIENT = BaseApiClient(api_key="dummy")
+
+
+def process_schema(schema: Any) -> dict[str, Any]:
+    # The only thing we need from the client is the t_schema function
+    try:
+        json_schema = t_schema(DUMMY_CLIENT, schema).json_schema.model_dump(
+            exclude_unset=True, exclude_none=True
+        )
+    except Exception:
+        json_schema = {}
+    return json_schema
+
+
+class SchemaJSONEncoder(json.JSONEncoder):
+    def default(self, o: Any) -> Any:
+        if isinstance(o, JSONSchemaType):
+            return o.value
+        return super().default(o)