diff --git a/README.md b/README.md index 664161c..fcdb8f9 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ Generic Python agent SDK with OpenTelemetry instrumentation and a plugin archite ```bash bash scripts/fetch-vendor.sh -pip install -e ".[dev,anthropic,openai,litellm]" +pip install -e ".[dev,anthropic,google-genai,openai,litellm]" ./scripts/run-unit-tests.sh ``` @@ -55,7 +55,7 @@ pip install harness-sdk Optional instrumentation extras: ```bash -pip install "harness-sdk[anthropic,openai,litellm]" +pip install "harness-sdk[anthropic,google-genai,openai,litellm]" ``` ### From TestPyPI (release candidates) @@ -75,7 +75,7 @@ With optional extras: pip install \ --index-url https://test.pypi.org/simple/ \ --extra-index-url https://pypi.org/simple/ \ - "harness-sdk[anthropic,openai,litellm]==1.0.0-rc.1" + "harness-sdk[anthropic,google-genai,openai,litellm]==1.0.0-rc.1" ``` Pin the version to the RC you want (see [published files](https://test.pypi.org/project/harness-sdk/#history) on TestPyPI). The version string matches the tag without the leading `v` (tag `v1.0.0-rc.1` → `1.0.0-rc.1`). In a virtualenv or `requirements.txt`, the same flags apply: diff --git a/docs/instrumentation/google-genai.md b/docs/instrumentation/google-genai.md new file mode 100644 index 0000000..21114cb --- /dev/null +++ b/docs/instrumentation/google-genai.md @@ -0,0 +1,111 @@ +# Google Gen AI (`google-genai`) instrumentation + +Instruments the unified **Google Gen AI SDK** (`from google import genai`) for +both backends: + +- **Gemini Developer API** — `genai.Client(api_key=...)` → `gen_ai.provider.name = gcp.gemini` +- **Vertex AI** — `genai.Client(vertexai=True, project=..., location=...)` → `gen_ai.provider.name = gcp.vertex_ai` + +> The legacy `vertexai.generative_models` / `google-cloud-aiplatform` generative +> modules and `google-generativeai` are **not** instrumented. Google deprecated +> them (generative modules removal 2026-06-24); `google-genai` is the supported client. + +Enable with `pip install "harness-sdk[google-genai]"`. Telemetry is produced via +`opentelemetry-util-genai` (`TelemetryHandler` → `LLMInvocation` / `EmbeddingInvocation`), +so spans, metrics, and attributes follow the OpenTelemetry GenAI semantic conventions, +identical in shape to the OpenAI/Anthropic wrappers. + +## Covered call surface + +| Method (class) | Span | +|---|---| +| `Models.generate_content` / `AsyncModels.generate_content` | `chat {model}` | +| `Models.generate_content_stream` / `AsyncModels.generate_content_stream` | `chat {model}` (streaming) | +| `Models.embed_content` / `AsyncModels.embed_content` | `embeddings {model}` | + +Not covered: the Live API (`client.aio.live`), batch/prediction, and tuning. Chat +sessions (`client.chats`) are covered indirectly because they delegate to +`generate_content`. + +## Configuration gates + +| Config key | Env | Effect | +|---|---|---| +| `gen_ai.enabled` | `HA_GEN_AI_ENABLED` | Master switch; off = passthrough, no span | +| `gen_ai.payload_capture_enabled` | `HA_GEN_AI_PAYLOAD_CAPTURE_ENABLED` | Enables prompt/response/tool **content** capture (sets OTEL experimental + `SPAN_ONLY`) | +| `gen_ai.payload_evaluation_enabled` | `HA_GEN_AI_PAYLOAD_EVALUATION_ENABLED` | Enables pre-call control-plugin evaluation (blocking) | + +--- + +## Captured fields, by confidence + +### ✅ Will **definitely** be captured + +Set on every span whenever `gen_ai.enabled = true`, independent of content capture +and independent of the response contents (they derive from the span lifecycle and +the request arguments the SDK requires): + +| Attribute | Source | +|---|---| +| `gen_ai.operation.name` (`chat` / `embeddings`) | wrapper (fixed per method) | +| `gen_ai.provider.name` (`gcp.vertex_ai` \| `gcp.gemini` \| `gcp.gen_ai`) | client backend detection | +| `gen_ai.request.model` | `model` argument | +| `gen_ai.framework` = `google-genai` | wrapper (fixed) | +| `gen_ai.request.streaming` = `true` | streaming methods only | +| span name (`chat {model}`), span kind `CLIENT`, status/error | `TelemetryHandler` | + +### 🟡 **Likely** captured + +Present whenever the request/response actually carries the value — i.e. the normal +case for a successful call, but not guaranteed (a field may be omitted by the API, +the caller may not set a request param, or an error may end the span early): + +| Attribute | Source | Why not guaranteed | +|---|---|---| +| `gen_ai.response.model` | `response.model_version` | Some responses omit `model_version` | +| `gen_ai.usage.input_tokens` | `usage_metadata.prompt_token_count` | Absent on failed calls | +| `gen_ai.usage.output_tokens` | `usage_metadata.candidates_token_count` | Absent on failed calls | +| `gen_ai.response.finish_reasons` | `candidate.finish_reason` | Absent if no candidates returned | +| `gen_ai.request.temperature` / `top_p` / `max_tokens` / `stop_sequences` / `seed` | request `config` | Only when the caller sets them | +| `gen_ai.input.messages` | `contents` | **Only when `payload_capture_enabled`** | +| `gen_ai.output.messages` | `response.candidates[*].content.parts` | **Only when `payload_capture_enabled`** | +| `gen_ai.system_instructions` | `config.system_instruction` | Only when set **and** content capture on | + +Content attributes (`*.messages`, `system_instructions`) are serialized only in +OTEL experimental mode with capture mode `SPAN_ONLY`/`SPAN_AND_EVENT`, which the SDK +enables from `payload_capture_enabled`. + +### 🟠 **Probable** (best-effort / conditional / version-sensitive) + +Captured only under specific conditions or subject to SDK/backend variability: + +| Attribute / behavior | Condition it depends on | +|---|---| +| **Tool calls** in `gen_ai.output.messages` (`ToolCallRequest`: name, arguments, id) | model actually returns `function_call` parts **and** content capture on | +| Tool results (`ToolCallResponse`) in `gen_ai.input.messages` | caller passes `function_response` parts back in `contents` | +| `gen_ai.response.id` | populated by response; the Gemini Developer API sometimes omits `response_id` | +| Streaming token usage / finish reasons | only if the **final** stream chunk carries `usage_metadata` (Gemini normally does; not contractual) | +| Streaming output text | reconstructed by concatenating `chunk.text`; non-text parts in streams are not reassembled | +| Embeddings `gen_ai.usage.input_tokens` | from `embeddings[0].statistics.token_count` — populated on Vertex, often absent on the Gemini Developer API | +| Embeddings dimension count | derived from `embeddings[0].values` length | +| Input message roles/parts for complex `contents` | best-effort mapping of str / `Part` / `Content`; exotic part types (blobs, files, URIs) are not fully expanded | + +### ❌ Not captured (current limitations) + +- **Automatic Function Calling (AFC):** a single `generate_content` that internally + loops through tool calls is recorded as **one** span (the final response); + intermediate turns / `automatic_function_calling_history` are not expanded into + child spans. +- Cache token counts (`gen_ai.usage.cache_*`) and `tool_use_prompt_token_count`. +- `gen_ai.tool.definitions` (declared tools) — not currently mapped from `config.tools`. +- Live API bidirectional streaming. + +--- + +## Blocking (control plugins) + +When `payload_evaluation_enabled = true` and a blocking control plugin is +registered, each call is evaluated **before** the request is sent. A blocking +decision raises `ControlEvaluationBlocked`, marks the span as errored, and the +underlying Google API is never called — identical to the OpenAI/Anthropic/LiteLLM +GenAI wrappers. diff --git a/pyproject.toml b/pyproject.toml index c218362..3e8858a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ anthropic = ["anthropic>=0.34.0"] openai = ["openai>=1.40.0"] litellm = ["litellm>=1.60.0"] +google-genai = ["google-genai>=1.0.0"] dev = [ "pytest==7.4.3", "pytest-cov==3.0.0", @@ -63,6 +64,7 @@ dev = [ "anthropic>=0.34.0", "openai>=1.40.0", "litellm>=1.60.0", + "google-genai>=1.0.0", "grpcio", ] diff --git a/src/harness_sdk/instrumentation/google_genai/__init__.py b/src/harness_sdk/instrumentation/google_genai/__init__.py new file mode 100644 index 0000000..4bf29a4 --- /dev/null +++ b/src/harness_sdk/instrumentation/google_genai/__init__.py @@ -0,0 +1,602 @@ +""" +Google Gen AI SDK (``google-genai``) instrumentation using +``opentelemetry-util-genai`` TelemetryHandler for span/metric telemetry, plus +Traceable pre-call policy evaluation. + +Covers the unified Google Gen AI SDK (``from google import genai``) against both +the Gemini Developer API and the Gemini API on Vertex AI +(``genai.Client(vertexai=True, ...)``). + +Coverage: + - Models.generate_content / AsyncModels.generate_content (non-streaming) + - Models.generate_content_stream / AsyncModels.generate_content_stream (streaming) + - Models.embed_content / AsyncModels.embed_content + +The classic ``vertexai.generative_models`` / ``google-cloud-aiplatform`` +generative modules are deprecated by Google (removal 2026-06-24) and are NOT +instrumented here; ``google-genai`` is the supported client. + +Optional: ``pip install harness-sdk[google-genai]`` +""" +# pylint: disable=duplicate-code + +from __future__ import annotations + +from typing import Any, Callable, Iterator, List, Optional + +import wrapt +from opentelemetry.instrumentation.utils import unwrap +from opentelemetry.util.genai.handler import TelemetryHandler +from opentelemetry.util.genai.types import ( + Error, + InputMessage, + LLMInvocation, + OutputMessage, + Text, + ToolCallRequest, + ToolCallResponse, +) +from opentelemetry.util.genai.utils import should_capture_content_on_spans_in_experimental_mode + +from harness_sdk.config.config import Config +from harness_sdk.custom_logger import get_custom_logger +from harness_sdk.plugins.control import get_control_registry +from harness_sdk.gen_ai.exceptions import ControlEvaluationBlocked +from harness_sdk.instrumentation import BaseInstrumentorWrapper + +logger = get_custom_logger(__name__) + +_MODELS_MODULE = "google.genai.models" + +_PROVIDER_VERTEX = "gcp.vertex_ai" +_PROVIDER_GEMINI = "gcp.gemini" +_PROVIDER_FALLBACK = "gcp.gen_ai" + + +def _get_handler() -> TelemetryHandler: + return TelemetryHandler() + + +def _get(source: Any, key: str, default: Any = None) -> Any: + """Read ``key`` from a dict or an object, returning ``default`` if absent.""" + if source is None: + return default + if isinstance(source, dict): + return source.get(key, default) + return getattr(source, key, default) + + +def _resolve_provider(instance: Any) -> str: + """Return the gen_ai.provider.name based on the client backend (Vertex vs Gemini). + + ``BaseApiClient.vertexai`` is truthy only for the Vertex AI backend; it is + falsy (``None``/``False``) for the Gemini Developer API. When the client + cannot be located we fall back to the generic GCP provider. + """ + api_client = _get(instance, "_api_client") + if api_client is None: + return _PROVIDER_FALLBACK + return _PROVIDER_VERTEX if _get(api_client, "vertexai") else _PROVIDER_GEMINI + + +# --------------------------------------------------------------------------- # +# Request / response mapping +# --------------------------------------------------------------------------- # +def _extract_model(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Optional[str]: + model = kwargs.get("model") + if model is None and args: + model = args[0] + return str(model) if model is not None else None + + +def _extract_contents(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + contents = kwargs.get("contents") + if contents is None and len(args) > 1: + contents = args[1] + return contents + + +def _part_to_message_part(part: Any) -> Optional[Any]: + """Map a google-genai Part to a util-genai MessagePart (best-effort).""" + text = _get(part, "text") + if isinstance(text, str) and text: + return Text(content=text) + + function_call = _get(part, "function_call") + if function_call is not None: + return ToolCallRequest( + arguments=_get(function_call, "args"), + name=_get(function_call, "name") or "", + id=_get(function_call, "id"), + ) + + function_response = _get(part, "function_response") + if function_response is not None: + return ToolCallResponse( + response=_get(function_response, "response"), + id=_get(function_response, "id"), + ) + return None + + +def _content_to_parts(content: Any) -> List[Any]: + if isinstance(content, str): + return [Text(content=content)] + parts_out: List[Any] = [] + raw_parts = _get(content, "parts") + if raw_parts: + for raw in raw_parts: + mapped = _part_to_message_part(raw) + if mapped is not None: + parts_out.append(mapped) + return parts_out + # A bare Part (has text/function_call directly) + mapped = _part_to_message_part(content) + if mapped is not None: + parts_out.append(mapped) + return parts_out + + +def _to_input_messages(contents: Any, capture_content: bool) -> List[InputMessage]: + if not capture_content or contents is None: + return [] + items = contents if isinstance(contents, list) else [contents] + messages: List[InputMessage] = [] + for item in items: + role = _get(item, "role") or "user" + parts = _content_to_parts(item) + if parts: + messages.append(InputMessage(role=str(role), parts=parts)) + return messages + + +def _to_system_instruction(config: Any, capture_content: bool) -> List[Any]: + if not capture_content or config is None: + return [] + system = _get(config, "system_instruction") + if system is None: + return [] + return _content_to_parts(system) + + +def _build_invocation( + handler: TelemetryHandler, + provider: str, + model: Optional[str], + contents: Any, + config: Any, + capture_content: bool, + *, + streaming: bool, +) -> LLMInvocation: + attributes: dict[str, Any] = {"gen_ai.framework": "google-genai"} + if streaming: + attributes["gen_ai.request.streaming"] = True + invocation = LLMInvocation( + request_model=model, + provider=provider, + input_messages=_to_input_messages(contents, capture_content), + system_instruction=_to_system_instruction(config, capture_content), + temperature=_get(config, "temperature"), + top_p=_get(config, "top_p"), + max_tokens=_get(config, "max_output_tokens"), + stop_sequences=_get(config, "stop_sequences"), + seed=_get(config, "seed"), + attributes=attributes, + ) + handler.start_llm(invocation) + return invocation + + +def _candidate_to_output_message(candidate: Any, capture_content: bool) -> Optional[OutputMessage]: + finish_reason = _get(candidate, "finish_reason") + finish_reason_str = str(finish_reason) if finish_reason is not None else "stop" + parts: List[Any] = [] + if capture_content: + content = _get(candidate, "content") + parts = _content_to_parts(content) if content is not None else [] + role = _get(_get(candidate, "content"), "role") or "assistant" + return OutputMessage(role=str(role), parts=parts, finish_reason=finish_reason_str) + + +def _apply_response(invocation: LLMInvocation, response: Any, capture_content: bool) -> None: + """Populate the invocation from a google-genai GenerateContentResponse.""" + if response is None: + return + + invocation.response_id = _get(response, "response_id") + invocation.response_model_name = _get(response, "model_version") + + usage = _get(response, "usage_metadata") + if usage is not None: + invocation.input_tokens = _get(usage, "prompt_token_count") + invocation.output_tokens = _get(usage, "candidates_token_count") + + candidates = _get(response, "candidates") or [] + output_messages: List[OutputMessage] = [] + finish_reasons: List[str] = [] + for candidate in candidates: + message = _candidate_to_output_message(candidate, capture_content) + if message is not None: + output_messages.append(message) + finish_reasons.append(message.finish_reason) + if output_messages: + invocation.output_messages = output_messages + if finish_reasons: + invocation.finish_reasons = finish_reasons + + +class _StreamAccumulator: + """Accumulates streamed chunks so the span can be finalized once complete.""" + + def __init__(self, capture_content: bool) -> None: + self._capture_content = capture_content + self._last_response: Any = None + self._text_parts: List[str] = [] + self._finish_reasons: List[str] = [] + self._response_id: Optional[str] = None + self._model_version: Optional[str] = None + self._input_tokens: Optional[int] = None + self._output_tokens: Optional[int] = None + + def add(self, chunk: Any) -> None: + self._last_response = chunk + response_id = _get(chunk, "response_id") + if response_id is not None: + self._response_id = response_id + model_version = _get(chunk, "model_version") + if model_version is not None: + self._model_version = model_version + usage = _get(chunk, "usage_metadata") + if usage is not None: + input_tokens = _get(usage, "prompt_token_count") + output_tokens = _get(usage, "candidates_token_count") + if input_tokens is not None: + self._input_tokens = input_tokens + if output_tokens is not None: + self._output_tokens = output_tokens + if self._capture_content: + text = _get(chunk, "text") + if isinstance(text, str) and text: + self._text_parts.append(text) + for candidate in _get(chunk, "candidates") or []: + finish_reason = _get(candidate, "finish_reason") + if finish_reason is not None: + self._finish_reasons.append(str(finish_reason)) + + def apply(self, invocation: LLMInvocation) -> None: + invocation.response_id = self._response_id + invocation.response_model_name = self._model_version + invocation.input_tokens = self._input_tokens + invocation.output_tokens = self._output_tokens + if self._finish_reasons: + invocation.finish_reasons = self._finish_reasons + if self._capture_content and self._text_parts: + invocation.output_messages = [ + OutputMessage( + role="assistant", + parts=[Text(content="".join(self._text_parts))], + finish_reason=self._finish_reasons[-1] if self._finish_reasons else "stop", + ) + ] + + +def _evaluate_invocation(invocation: LLMInvocation) -> None: + """Run Traceable policy evaluation against the live span; raise if blocked.""" + if not Config().config.gen_ai.payload_evaluation_enabled.value: + logger.debug("google-genai: evaluate_body disabled, skipping policy evaluation") + return + span = getattr(invocation, "span", None) + if span is None: + inference_invocation = getattr(invocation, "_inference_invocation", None) + if inference_invocation is not None: + span = getattr(inference_invocation, "span", None) + if span is None or not span.is_recording(): + logger.debug("google-genai: no active span, skipping evaluation") + return + try: + result = get_control_registry().evaluate_agent_span(span, body="") + if result.block: + logger.debug("Traceable policy blocked google-genai request") + raise ControlEvaluationBlocked(result) + except ControlEvaluationBlocked: + raise + except Exception as err: # pylint: disable=broad-except + logger.debug("google-genai span evaluation error: %s", err) + + +# --------------------------------------------------------------------------- # +# generate_content (non-streaming) +# --------------------------------------------------------------------------- # +def _make_generate_sync(handler: TelemetryHandler) -> Callable[..., Any]: + capture_content = should_capture_content_on_spans_in_experimental_mode() + + def _wrapper(wrapped, instance, args, kwargs): + if not Config().config.gen_ai.enabled.value: + return wrapped(*args, **kwargs) + invocation = _build_invocation( + handler, _resolve_provider(instance), + _extract_model(args, kwargs), _extract_contents(args, kwargs), + kwargs.get("config"), capture_content, streaming=False, + ) + try: + _evaluate_invocation(invocation) + except ControlEvaluationBlocked: + handler.fail_llm(invocation, Error(message="blocked", type=ControlEvaluationBlocked)) + raise + try: + result = wrapped(*args, **kwargs) + _apply_response(invocation, result, capture_content) + handler.stop_llm(invocation) + return result + except ControlEvaluationBlocked: + raise + except Exception as exc: # pylint: disable=broad-except + handler.fail_llm(invocation, Error(message=str(exc), type=type(exc))) + raise + + return _wrapper + + +def _make_generate_async(handler: TelemetryHandler) -> Callable[..., Any]: + capture_content = should_capture_content_on_spans_in_experimental_mode() + + async def _wrapper(wrapped, instance, args, kwargs): + if not Config().config.gen_ai.enabled.value: + return await wrapped(*args, **kwargs) + invocation = _build_invocation( + handler, _resolve_provider(instance), + _extract_model(args, kwargs), _extract_contents(args, kwargs), + kwargs.get("config"), capture_content, streaming=False, + ) + try: + _evaluate_invocation(invocation) + except ControlEvaluationBlocked: + handler.fail_llm(invocation, Error(message="blocked", type=ControlEvaluationBlocked)) + raise + try: + result = await wrapped(*args, **kwargs) + _apply_response(invocation, result, capture_content) + handler.stop_llm(invocation) + return result + except ControlEvaluationBlocked: + raise + except Exception as exc: # pylint: disable=broad-except + handler.fail_llm(invocation, Error(message=str(exc), type=type(exc))) + raise + + return _wrapper + + +# --------------------------------------------------------------------------- # +# generate_content_stream (streaming) +# --------------------------------------------------------------------------- # +def _finalize_stream(handler, invocation, accumulator, error: Optional[BaseException]) -> None: + if error is not None: + handler.fail_llm(invocation, Error(message=str(error), type=type(error))) + return + accumulator.apply(invocation) + handler.stop_llm(invocation) + + +def _iterate_sync(result: Iterator[Any], handler, invocation, accumulator) -> Iterator[Any]: + finalized = False + try: + for chunk in result: + accumulator.add(chunk) + yield chunk + _finalize_stream(handler, invocation, accumulator, None) + finalized = True + except Exception as exc: # pylint: disable=broad-except + _finalize_stream(handler, invocation, accumulator, exc) + finalized = True + raise + finally: + if not finalized: + _finalize_stream(handler, invocation, accumulator, None) + + +async def _iterate_async(result, handler, invocation, accumulator): + finalized = False + try: + async for chunk in result: + accumulator.add(chunk) + yield chunk + _finalize_stream(handler, invocation, accumulator, None) + finalized = True + except Exception as exc: # pylint: disable=broad-except + _finalize_stream(handler, invocation, accumulator, exc) + finalized = True + raise + finally: + if not finalized: + _finalize_stream(handler, invocation, accumulator, None) + + +def _make_generate_stream_sync(handler: TelemetryHandler) -> Callable[..., Any]: + capture_content = should_capture_content_on_spans_in_experimental_mode() + + def _wrapper(wrapped, instance, args, kwargs): + if not Config().config.gen_ai.enabled.value: + return wrapped(*args, **kwargs) + invocation = _build_invocation( + handler, _resolve_provider(instance), + _extract_model(args, kwargs), _extract_contents(args, kwargs), + kwargs.get("config"), capture_content, streaming=True, + ) + try: + _evaluate_invocation(invocation) + except ControlEvaluationBlocked: + handler.fail_llm(invocation, Error(message="blocked", type=ControlEvaluationBlocked)) + raise + try: + result = wrapped(*args, **kwargs) + except Exception as exc: # pylint: disable=broad-except + handler.fail_llm(invocation, Error(message=str(exc), type=type(exc))) + raise + return _iterate_sync(result, handler, invocation, _StreamAccumulator(capture_content)) + + return _wrapper + + +def _make_generate_stream_async(handler: TelemetryHandler) -> Callable[..., Any]: + capture_content = should_capture_content_on_spans_in_experimental_mode() + + async def _wrapper(wrapped, instance, args, kwargs): + if not Config().config.gen_ai.enabled.value: + return await wrapped(*args, **kwargs) + invocation = _build_invocation( + handler, _resolve_provider(instance), + _extract_model(args, kwargs), _extract_contents(args, kwargs), + kwargs.get("config"), capture_content, streaming=True, + ) + try: + _evaluate_invocation(invocation) + except ControlEvaluationBlocked: + handler.fail_llm(invocation, Error(message="blocked", type=ControlEvaluationBlocked)) + raise + try: + result = await wrapped(*args, **kwargs) + except Exception as exc: # pylint: disable=broad-except + handler.fail_llm(invocation, Error(message=str(exc), type=type(exc))) + raise + return _iterate_async(result, handler, invocation, _StreamAccumulator(capture_content)) + + return _wrapper + + +# --------------------------------------------------------------------------- # +# embed_content +# --------------------------------------------------------------------------- # +def _apply_embedding_response(invocation: Any, result: Any) -> None: + invocation.response_model_name = _get(result, "model_version") + embeddings = _get(result, "embeddings") or [] + if embeddings: + values = _get(embeddings[0], "values") + if isinstance(values, list): + invocation.dimension_count = len(values) + stats = _get(embeddings[0], "statistics") + token_count = _get(stats, "token_count") + if token_count is not None: + invocation.input_tokens = int(token_count) + + +def _make_embed_sync(handler: TelemetryHandler) -> Callable[..., Any]: + def _wrapper(wrapped, instance, args, kwargs): + if not Config().config.gen_ai.enabled.value: + return wrapped(*args, **kwargs) + invocation = handler.start_embedding( + _resolve_provider(instance), + request_model=_extract_model(args, kwargs), + ) + try: + _evaluate_invocation(invocation) + except ControlEvaluationBlocked: + invocation.fail(Error(message="blocked", type=ControlEvaluationBlocked)) + raise + try: + result = wrapped(*args, **kwargs) + _apply_embedding_response(invocation, result) + invocation.stop() + return result + except ControlEvaluationBlocked: + raise + except Exception as exc: # pylint: disable=broad-except + invocation.fail(Error(message=str(exc), type=type(exc))) + raise + + return _wrapper + + +def _make_embed_async(handler: TelemetryHandler) -> Callable[..., Any]: + async def _wrapper(wrapped, instance, args, kwargs): + if not Config().config.gen_ai.enabled.value: + return await wrapped(*args, **kwargs) + invocation = handler.start_embedding( + _resolve_provider(instance), + request_model=_extract_model(args, kwargs), + ) + try: + _evaluate_invocation(invocation) + except ControlEvaluationBlocked: + invocation.fail(Error(message="blocked", type=ControlEvaluationBlocked)) + raise + try: + result = await wrapped(*args, **kwargs) + _apply_embedding_response(invocation, result) + invocation.stop() + return result + except ControlEvaluationBlocked: + raise + except Exception as exc: # pylint: disable=broad-except + invocation.fail(Error(message=str(exc), type=type(exc))) + raise + + return _wrapper + + +# --------------------------------------------------------------------------- # +# Instrumentor +# --------------------------------------------------------------------------- # +_WRAPPED_METHODS = [ + ("Models", "generate_content", _make_generate_sync), + ("Models", "generate_content_stream", _make_generate_stream_sync), + ("Models", "embed_content", _make_embed_sync), + ("AsyncModels", "generate_content", _make_generate_async), + ("AsyncModels", "generate_content_stream", _make_generate_stream_async), + ("AsyncModels", "embed_content", _make_embed_async), +] + + +class GoogleGenAIInstrumentorWrapper(BaseInstrumentorWrapper): + """Instrument the Google Gen AI SDK (google-genai) with OTEL GenAI telemetry and policy evaluation.""" # pylint: disable=line-too-long + + _applied: bool = False + + def __init__(self) -> None: + BaseInstrumentorWrapper.__init__(self) + + def instrument(self, **_kwargs: Any) -> None: + if self._applied: + logger.debug("google-genai instrumentation already applied.") + return + if not Config().config.gen_ai.enabled.value: + logger.debug("Gen AI instrumentation disabled; skip google-genai wraps.") + return + try: + handler = _get_handler() + for cls_name, method, factory in _WRAPPED_METHODS: + wrapped = factory(handler) + wrapt.wrap_function_wrapper(_MODELS_MODULE, f"{cls_name}.{method}", wrapped) + GoogleGenAIInstrumentorWrapper._applied = True + logger.debug("Traceable google-genai instrumentation applied.") + except ImportError as err: + logger.error("google-genai SDK not available: %s", err) + logger.error("Install with: pip install 'harness-sdk[google-genai]'") + raise + + def uninstrument(self, **_kwargs: Any) -> None: + if not self._applied: + return + from importlib import import_module # pylint: disable=import-outside-toplevel + + errors: List[Exception] = [] + try: + mod = import_module(_MODELS_MODULE) + except Exception as err: # pylint: disable=broad-except + logger.error("google-genai SDK not available: %s", err) + raise + for cls_name, method, _ in _WRAPPED_METHODS: + try: + unwrap(getattr(mod, cls_name), method) + except Exception as err: # pylint: disable=broad-except + logger.error("Failed to uninstrument google-genai %s.%s: %s", cls_name, method, err) + errors.append(err) + + GoogleGenAIInstrumentorWrapper._applied = False + logger.debug("google-genai instrumentation removed.") + if errors: + raise errors[0] + + +__all__ = ["GoogleGenAIInstrumentorWrapper"] diff --git a/src/harness_sdk/instrumentation/instrumentation_definitions.py b/src/harness_sdk/instrumentation/instrumentation_definitions.py index be9f813..91cf359 100644 --- a/src/harness_sdk/instrumentation/instrumentation_definitions.py +++ b/src/harness_sdk/instrumentation/instrumentation_definitions.py @@ -18,6 +18,7 @@ ANTHROPIC_KEY = 'anthropic' OPENAI_KEY = 'openai' LITELLM_KEY = 'litellm' +GOOGLE_GENAI_KEY = 'google_genai' MCP_KEY = 'mcp' SUPPORTED_LIBRARIES = [ @@ -29,6 +30,7 @@ ANTHROPIC_KEY, OPENAI_KEY, LITELLM_KEY, + GOOGLE_GENAI_KEY, MCP_KEY, ] @@ -84,6 +86,7 @@ def _mark_as_instrumented(library_key, wrapper_instance): ANTHROPIC_KEY: ("harness_sdk.instrumentation.anthropic", "AnthropicInstrumentorWrapper"), OPENAI_KEY: ("harness_sdk.instrumentation.openai", "OpenAIInstrumentorWrapper"), LITELLM_KEY: ("harness_sdk.instrumentation.litellm", "LiteLLMInstrumentorWrapper"), + GOOGLE_GENAI_KEY: ("harness_sdk.instrumentation.google_genai", "GoogleGenAIInstrumentorWrapper"), MCP_KEY: ("harness_sdk.instrumentation.mcp", "McpInstrumentorWrapper"), } diff --git a/test/instrumentation/google_genai/__init__.py b/test/instrumentation/google_genai/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/test/instrumentation/google_genai/__init__.py @@ -0,0 +1 @@ + diff --git a/test/instrumentation/google_genai/google_genai_instrumentation_test.py b/test/instrumentation/google_genai/google_genai_instrumentation_test.py new file mode 100644 index 0000000..4975ccf --- /dev/null +++ b/test/instrumentation/google_genai/google_genai_instrumentation_test.py @@ -0,0 +1,198 @@ +"""End-to-end google-genai instrumentation tests (skipped if SDK not installed).""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from harness_sdk.plugins.control import get_control_registry +from harness_sdk.gen_ai.exceptions import ControlEvaluationBlocked +from harness_sdk.instrumentation import google_genai as gg + +pytest.importorskip("google.genai") + +from google.genai.models import Models # noqa: E402 + + +def _fake_usage(prompt=5, candidates=8, total=13): + return SimpleNamespace( + prompt_token_count=prompt, + candidates_token_count=candidates, + total_token_count=total, + ) + + +def _fake_response(text="hello from gemini", finish_reason="STOP"): + candidate = SimpleNamespace( + finish_reason=finish_reason, + content=SimpleNamespace( + role="model", + parts=[SimpleNamespace(text=text, function_call=None, function_response=None)], + ), + ) + return SimpleNamespace( + response_id="resp-123", + model_version="gemini-2.0-flash", + usage_metadata=_fake_usage(), + candidates=[candidate], + text=text, + ) + + +def _client(): + from google import genai # pylint: disable=import-outside-toplevel + + return genai.Client(api_key="test-key") + + +def _vertex_client(): + from google import genai # pylint: disable=import-outside-toplevel + + return genai.Client(vertexai=True, project="p", location="us-central1") + + +@pytest.fixture +def google_genai_instrumentor(): + wrapper = gg.GoogleGenAIInstrumentorWrapper() + yield wrapper + if getattr(wrapper, "_applied", False): + wrapper.uninstrument() + get_control_registry().clear() + + +def test_generate_content_span_has_gen_ai_attributes(agent, exporter, google_genai_instrumentor): # pylint: disable=unused-argument + def fake_generate(_self, *_args, **_kwargs): + return _fake_response() + + with patch.object(Models, "generate_content", new=fake_generate): + google_genai_instrumentor.instrument() + client = _client() + client.models.generate_content(model="gemini-2.0-flash", contents="hi") + + spans = exporter.get_finished_spans() + exporter.clear() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs.get("gen_ai.operation.name") == "chat" + assert attrs.get("gen_ai.provider.name") == "gcp.gemini" + assert attrs.get("gen_ai.request.model") == "gemini-2.0-flash" + assert attrs.get("gen_ai.response.id") == "resp-123" + assert attrs.get("gen_ai.response.model") == "gemini-2.0-flash" + assert attrs.get("gen_ai.usage.input_tokens") == 5 + assert attrs.get("gen_ai.usage.output_tokens") == 8 + + +def test_vertex_generate_content_span_has_vertex_provider( + agent, exporter, google_genai_instrumentor +): # pylint: disable=unused-argument + def fake_generate(_self, *_args, **_kwargs): + return _fake_response() + + with patch.object(Models, "generate_content", new=fake_generate): + google_genai_instrumentor.instrument() + client = _vertex_client() + client.models.generate_content(model="gemini-2.0-flash", contents="hi") + + spans = exporter.get_finished_spans() + exporter.clear() + assert len(spans) == 1 + assert spans[0].attributes.get("gen_ai.provider.name") == "gcp.vertex_ai" + + +def test_embed_content_span_has_embedding_attributes( + agent, exporter, google_genai_instrumentor +): # pylint: disable=unused-argument + response = SimpleNamespace( + model_version="text-embedding-004", + embeddings=[ + SimpleNamespace( + values=[0.1, 0.2, 0.3], + statistics=SimpleNamespace(token_count=7), + ) + ], + ) + + def fake_embed(_self, *_args, **_kwargs): + return response + + with patch.object(Models, "embed_content", new=fake_embed): + google_genai_instrumentor.instrument() + client = _client() + client.models.embed_content(model="text-embedding-004", contents="hi") + + spans = exporter.get_finished_spans() + exporter.clear() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs.get("gen_ai.operation.name") == "embeddings" + assert attrs.get("gen_ai.usage.input_tokens") == 7 + + +def test_uninstrument_roundtrip( + agent, exporter, google_genai_instrumentor +): # pylint: disable=unused-argument + def fake_generate(_self, *_args, **_kwargs): + return _fake_response() + + with patch.object(Models, "generate_content", new=fake_generate): + client = _client() + google_genai_instrumentor.instrument() + client.models.generate_content(model="gemini-2.0-flash", contents="first") + google_genai_instrumentor.uninstrument() + client.models.generate_content(model="gemini-2.0-flash", contents="untraced") + google_genai_instrumentor.instrument() + client.models.generate_content(model="gemini-2.0-flash", contents="second") + + spans = exporter.get_finished_spans() + exporter.clear() + assert len(spans) == 2 + + +def test_generate_content_evaluate_blocks_before_call(agent, exporter, google_genai_instrumentor): # pylint: disable=unused-argument + from test.control_test_helpers import AlwaysBlockControlPlugin # pylint: disable=import-outside-toplevel + + calls = {"n": 0} + + def counting_fake(_self, *_a, **_k): + calls["n"] += 1 + return _fake_response() + + get_control_registry().register(AlwaysBlockControlPlugin()) + + with patch.object(Models, "generate_content", new=counting_fake): + google_genai_instrumentor.instrument() + client = _client() + with pytest.raises(ControlEvaluationBlocked): + client.models.generate_content(model="gemini-2.0-flash", contents="hi") + + assert calls["n"] == 0 + spans = exporter.get_finished_spans() + exporter.clear() + assert len(spans) == 1 + + +def test_generate_content_stream_accumulates(agent, exporter, google_genai_instrumentor): # pylint: disable=unused-argument + def fake_stream(_self, *_args, **_kwargs): + yield SimpleNamespace( + response_id="resp-s", model_version="gemini-2.0-flash", + usage_metadata=None, candidates=[], text="Hel", + ) + yield SimpleNamespace( + response_id=None, model_version=None, + usage_metadata=_fake_usage(prompt=1, candidates=2), + candidates=[SimpleNamespace(finish_reason="STOP")], text="lo", + ) + + with patch.object(Models, "generate_content_stream", new=fake_stream): + google_genai_instrumentor.instrument() + client = _client() + chunks = list(client.models.generate_content_stream(model="gemini-2.0-flash", contents="hi")) + + assert len(chunks) == 2 + spans = exporter.get_finished_spans() + exporter.clear() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs.get("gen_ai.request.streaming") is True + assert attrs.get("gen_ai.usage.input_tokens") == 1 + assert attrs.get("gen_ai.usage.output_tokens") == 2 diff --git a/test/instrumentation/google_genai/google_genai_mapping_test.py b/test/instrumentation/google_genai/google_genai_mapping_test.py new file mode 100644 index 0000000..b1439b8 --- /dev/null +++ b/test/instrumentation/google_genai/google_genai_mapping_test.py @@ -0,0 +1,112 @@ +"""Pure request/response -> gen_ai attribute mapping tests for google-genai. + +These do not require the ``google-genai`` package to be installed; they only +exercise the mapping helpers against fake response objects. +""" + +from types import SimpleNamespace + +from harness_sdk.instrumentation import google_genai as gg + + +def _fake_usage(prompt=5, candidates=8, total=13): + return SimpleNamespace( + prompt_token_count=prompt, + candidates_token_count=candidates, + total_token_count=total, + ) + + +def _fake_text_part(text): + return SimpleNamespace(text=text, function_call=None, function_response=None) + + +def _fake_function_call_part(name, args, call_id="call-1"): + return SimpleNamespace( + text=None, + function_call=SimpleNamespace(name=name, args=args, id=call_id), + function_response=None, + ) + + +def _fake_response(text="hello from gemini", finish_reason="STOP", with_tool=False): + parts = [_fake_text_part(text)] + if with_tool: + parts.append(_fake_function_call_part("get_weather", {"city": "SF"})) + candidate = SimpleNamespace( + finish_reason=finish_reason, + content=SimpleNamespace(role="model", parts=parts), + ) + return SimpleNamespace( + response_id="resp-123", + model_version="gemini-2.0-flash", + usage_metadata=_fake_usage(), + candidates=[candidate], + text=text, + ) + + +def test_apply_response_sets_metadata_fields(): + invocation = gg.LLMInvocation(request_model="gemini-2.0-flash", provider="gcp.gemini") + gg._apply_response(invocation, _fake_response(), capture_content=False) + + assert invocation.response_id == "resp-123" + assert invocation.response_model_name == "gemini-2.0-flash" + assert invocation.input_tokens == 5 + assert invocation.output_tokens == 8 + assert invocation.finish_reasons == ["STOP"] + assert invocation.output_messages[0].parts == [] + + +def test_apply_response_captures_content_and_tool_calls(): + invocation = gg.LLMInvocation(request_model="gemini-2.0-flash", provider="gcp.gemini") + gg._apply_response(invocation, _fake_response(with_tool=True), capture_content=True) + + parts = invocation.output_messages[0].parts + kinds = {p.type for p in parts} + assert "text" in kinds + assert "tool_call" in kinds + tool_part = next(p for p in parts if p.type == "tool_call") + assert tool_part.name == "get_weather" + assert tool_part.arguments == {"city": "SF"} + + +def test_to_input_messages_from_string(): + messages = gg._to_input_messages("what is the weather?", capture_content=True) + assert len(messages) == 1 + assert messages[0].role == "user" + assert messages[0].parts[0].content == "what is the weather?" + + +def test_to_input_messages_empty_when_capture_off(): + assert gg._to_input_messages("hi", capture_content=False) == [] + + +def test_stream_accumulator_merges_chunks(): + acc = gg._StreamAccumulator(capture_content=True) + acc.add(SimpleNamespace( + response_id="r1", model_version="gemini-2.0-flash", + usage_metadata=None, candidates=[], text="Hel", + )) + acc.add(SimpleNamespace( + response_id=None, model_version=None, + usage_metadata=_fake_usage(prompt=2, candidates=4), + candidates=[SimpleNamespace(finish_reason="STOP")], text="lo", + )) + invocation = gg.LLMInvocation(request_model="gemini-2.0-flash", provider="gcp.gemini") + acc.apply(invocation) + + assert invocation.response_id == "r1" + assert invocation.input_tokens == 2 + assert invocation.output_tokens == 4 + assert invocation.finish_reasons == ["STOP"] + assert invocation.output_messages[0].parts[0].content == "Hello" + + +def test_resolve_provider_detects_backend(): + vertex_instance = SimpleNamespace(_api_client=SimpleNamespace(vertexai=True)) + gemini_instance = SimpleNamespace(_api_client=SimpleNamespace(vertexai=False)) + unknown_instance = SimpleNamespace() + assert gg._resolve_provider(vertex_instance) == "gcp.vertex_ai" + assert gg._resolve_provider(gemini_instance) == "gcp.gemini" + assert gg._resolve_provider(unknown_instance) == "gcp.gen_ai"