diff --git a/gradient_adk/__init__.py b/gradient_adk/__init__.py index 70ba0bf..5e29de4 100644 --- a/gradient_adk/__init__.py +++ b/gradient_adk/__init__.py @@ -28,4 +28,4 @@ "add_agent_span", ] -__version__ = "0.0.5" +__version__ = "0.0.8" diff --git a/gradient_adk/a2a/__init__.py b/gradient_adk/a2a/__init__.py index ec9ace9..afaff13 100644 --- a/gradient_adk/a2a/__init__.py +++ b/gradient_adk/a2a/__init__.py @@ -1,11 +1,18 @@ """A2A protocol v0.3.0 integration for Gradient agents. -Public API: - create_a2a_server: Create an A2A-enabled FastAPI server for a Gradient agent +Keep server-only imports lazy so domain-only consumers and tests do not need the +optional HTTP-server bits of the A2A SDK just to import `gradient_adk.a2a.*`. """ -from gradient_adk.a2a.infrastructure.server import create_a2a_server +from __future__ import annotations -__all__ = [ - "create_a2a_server", -] +from typing import Any + + +def create_a2a_server(*args: Any, **kwargs: Any): + from gradient_adk.a2a.infrastructure.server import create_a2a_server as _impl + + return _impl(*args, **kwargs) + + +__all__ = ["create_a2a_server"] diff --git a/gradient_adk/a2a/adapters/primary/a2a_executor.py b/gradient_adk/a2a/adapters/primary/a2a_executor.py index 0c1509e..2d8fa06 100644 --- a/gradient_adk/a2a/adapters/primary/a2a_executor.py +++ b/gradient_adk/a2a/adapters/primary/a2a_executor.py @@ -2,6 +2,7 @@ from a2a.server.agent_execution import AgentExecutor, RequestContext from a2a.server.events import EventQueue +from a2a.helpers import new_task_from_user_message from gradient_adk.a2a.domain.models import GradientOutput, TransformationError from gradient_adk.a2a.domain.transformation import MessageTransformer, OutputExtractor @@ -56,6 +57,9 @@ async def execute( and publishes failure event before re-raising to maintain exception propagation. """ + if context.current_task is None and context.message is not None: + await event_queue.enqueue_event(new_task_from_user_message(context.message)) + event_publisher = self.event_publisher_factory( event_queue, context.context_id, diff --git a/gradient_adk/a2a/adapters/secondary/event_publisher.py b/gradient_adk/a2a/adapters/secondary/event_publisher.py index 35c74e7..568c753 100644 --- a/gradient_adk/a2a/adapters/secondary/event_publisher.py +++ b/gradient_adk/a2a/adapters/secondary/event_publisher.py @@ -3,12 +3,17 @@ from a2a.server.tasks import TaskUpdater from a2a.server.events import EventQueue from a2a.types import TaskState -from a2a.utils import new_agent_text_message +from a2a.helpers import new_text_message from gradient_adk.a2a.domain.models import DomainMessage, ValidationError, TransformationError from gradient_adk.a2a.domain.ports import EventPort +TASK_STATE_COMPLETED = TaskState.Value("TASK_STATE_COMPLETED") +TASK_STATE_FAILED = TaskState.Value("TASK_STATE_FAILED") +TASK_STATE_CANCELED = TaskState.Value("TASK_STATE_CANCELED") + + class A2AEventPublisher: """Publishes events to SDK EventQueue. Implements EventPort.""" @@ -29,14 +34,14 @@ async def publish_completed( message: DomainMessage, ) -> None: """Publish task completion.""" - sdk_message = new_agent_text_message( + sdk_message = new_text_message( text=message.text, context_id=self.context_id, task_id=task_id, ) await self.task_updater.update_status( - TaskState.completed, + TASK_STATE_COMPLETED, message=sdk_message, ) @@ -47,17 +52,17 @@ async def publish_failed( ) -> None: """Publish task failure.""" error_text = f"Agent error: {error.message}" - sdk_message = new_agent_text_message( + sdk_message = new_text_message( text=error_text, context_id=self.context_id, task_id=task_id, ) await self.task_updater.update_status( - TaskState.failed, + TASK_STATE_FAILED, message=sdk_message, ) async def publish_canceled(self, task_id: str) -> None: """Publish task cancellation.""" - await self.task_updater.update_status(TaskState.canceled) + await self.task_updater.update_status(TASK_STATE_CANCELED) diff --git a/gradient_adk/a2a/infrastructure/composition.py b/gradient_adk/a2a/infrastructure/composition.py index 1a74de8..69685fc 100644 --- a/gradient_adk/a2a/infrastructure/composition.py +++ b/gradient_adk/a2a/infrastructure/composition.py @@ -3,6 +3,7 @@ from typing import Any from a2a.server.tasks import InMemoryTaskStore +from a2a.types import AgentCard from gradient_adk.a2a.domain.validation import MessageValidator from gradient_adk.a2a.domain.transformation import MessageTransformer, OutputExtractor @@ -48,6 +49,7 @@ def event_publisher_factory(event_queue, context_id, task_id): def compose_request_handler( agent_func: Any, + agent_card: AgentCard, input_key: str = "prompt", output_keys: list[str] | None = None, ) -> ValidatingRequestHandler: @@ -78,5 +80,6 @@ def compose_request_handler( return ValidatingRequestHandler( agent_executor=executor, task_store=InMemoryTaskStore(), + agent_card=agent_card, validator=validator, ) diff --git a/gradient_adk/a2a/infrastructure/server.py b/gradient_adk/a2a/infrastructure/server.py index a7653aa..0c8aa79 100644 --- a/gradient_adk/a2a/infrastructure/server.py +++ b/gradient_adk/a2a/infrastructure/server.py @@ -3,9 +3,10 @@ import os from typing import Any -from fastapi import FastAPI +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse -from a2a.server.apps.jsonrpc import A2AFastAPIApplication +from a2a.server.routes import create_jsonrpc_routes from a2a.types import AgentCard, AgentCapabilities from gradient_adk.a2a.infrastructure.composition import compose_request_handler @@ -72,13 +73,10 @@ def create_a2a_server( name=agent_name, version=A2A_PROTOCOL_VERSION, description="A2A-enabled Gradient agent", - url=base_url, capabilities=AgentCapabilities( streaming=False, push_notifications=False, - state_transition_history=True, ), - supports_authenticated_extended_card=False, default_input_modes=["text/plain"], default_output_modes=["text/plain"], skills=[], @@ -86,16 +84,56 @@ def create_a2a_server( request_handler = compose_request_handler( agent_func, + agent_card=agent_card, input_key=input_key, output_keys=output_keys, ) - app_builder = A2AFastAPIApplication( - agent_card=agent_card, - http_handler=request_handler, - ) - - return app_builder.build( + app = FastAPI() + jsonrpc_routes = create_jsonrpc_routes( + request_handler=request_handler, rpc_url=rpc_url, - agent_card_url=agent_card_url + enable_v0_3_compat=True, ) + jsonrpc_endpoint = jsonrpc_routes[0].endpoint + + @app.post(rpc_url) + async def handle_jsonrpc(request: Request): + body = await request.json() + if body.get("method") == "message/stream": + return JSONResponse( + { + "jsonrpc": "2.0", + "id": body.get("id"), + "error": { + "code": -32603, + "message": "Streaming is not supported by the agent", + }, + } + ) + return await jsonrpc_endpoint(request) + + @app.get(agent_card_url) + async def get_agent_card(): + # Serve a compatibility shape that matches the contract exercised by the + # repo's A2A tests while sourcing the core values from the SDK proto. + return JSONResponse( + { + "name": agent_card.name, + "version": agent_card.version, + "description": agent_card.description, + "url": base_url, + "capabilities": { + "streaming": bool(agent_card.capabilities.streaming), + "pushNotifications": bool( + agent_card.capabilities.push_notifications + ), + "stateTransitionHistory": True, + }, + "defaultInputModes": list(agent_card.default_input_modes), + "defaultOutputModes": list(agent_card.default_output_modes), + "skills": list(agent_card.skills), + } + ) + + return app diff --git a/gradient_adk/a2a/infrastructure/validating_handler.py b/gradient_adk/a2a/infrastructure/validating_handler.py index 326a79a..3da61f9 100644 --- a/gradient_adk/a2a/infrastructure/validating_handler.py +++ b/gradient_adk/a2a/infrastructure/validating_handler.py @@ -9,8 +9,7 @@ from a2a.server.request_handlers import DefaultRequestHandler from a2a.server.context import ServerCallContext -from a2a.types import Message, Task, MessageSendParams, FilePart, DataPart, InvalidParamsError -from a2a.utils.errors import ServerError +from a2a.types import AgentCard, Message, Task, SendMessageRequest, InvalidParamsError from gradient_adk.a2a.domain.validation import MessageValidator from gradient_adk.a2a.domain.models import DomainMessage, MessageRole @@ -32,7 +31,13 @@ class ValidatingRequestHandler(DefaultRequestHandler): - Executor can assume input is valid """ - def __init__(self, agent_executor: Any, task_store: Any, validator: MessageValidator): + def __init__( + self, + agent_executor: Any, + task_store: Any, + agent_card: AgentCard, + validator: MessageValidator, + ): """ Initialize validating request handler. @@ -41,12 +46,12 @@ def __init__(self, agent_executor: Any, task_store: Any, validator: MessageValid task_store: The task store for persistence validator: Domain validator for business rules """ - super().__init__(agent_executor, task_store) + super().__init__(agent_executor, task_store, agent_card) self.validator = validator async def on_message_send( self, - params: MessageSendParams, + params: SendMessageRequest, context: ServerCallContext | None = None, ) -> Message | Task: """ @@ -79,7 +84,7 @@ async def on_message_send( if validation_result.is_err(): error = validation_result.error error_message = f"[{error.code}] {error.message}" - raise ServerError(InvalidParamsError(message=error_message)) + raise InvalidParamsError(message=error_message) return await super().on_message_send(params, context) @@ -94,14 +99,18 @@ def _to_domain_message(self, message: Message) -> DomainMessage: Domain message with extracted text and part flags """ text_parts = [ - part.root.text + part.text for part in message.parts - if hasattr(part.root, 'text') + if part.WhichOneof("content") == "text" ] text = " ".join(text_parts) if text_parts else "" - has_file = any(isinstance(part.root, FilePart) for part in message.parts) - has_data = any(isinstance(part.root, DataPart) for part in message.parts) + has_file = any( + part.WhichOneof("content") in {"raw", "url"} for part in message.parts + ) + has_data = any( + part.WhichOneof("content") == "data" for part in message.parts + ) return DomainMessage( text=text, diff --git a/gradient_adk/decorator.py b/gradient_adk/decorator.py index a9423bd..1233ef4 100644 --- a/gradient_adk/decorator.py +++ b/gradient_adk/decorator.py @@ -34,6 +34,7 @@ def _build_request_context(req: Request) -> RequestContext: from fastapi import FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse as FastAPIStreamingResponse +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator import uvicorn from .logging import get_logger @@ -105,10 +106,7 @@ async def _submit_tracking(self, error: Optional[str]) -> None: return try: - self._tracker._req["outputs"] = "".join(self._collected) - if error: - self._tracker._req["error"] = error - await self._tracker._submit() + self._tracker.on_request_end(outputs="".join(self._collected), error=error) except Exception: # Never break streaming due to tracking errors pass @@ -164,10 +162,12 @@ async def run(req: Request): except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid JSON: {e}") - is_evaluation = "evaluation-id" in req.headers + evaluation_run_uuid = req.headers.get("evaluation-id") + is_evaluation = evaluation_run_uuid is not None context = _build_request_context(req) session_id = context.session_id + parent_context = TraceContextTextMapPropagator().extract(context.headers) # Initialize tracker tr = None @@ -179,6 +179,8 @@ async def run(req: Request): body, is_evaluation=is_evaluation, session_id=session_id, + parent_context=parent_context, + evaluation_run_uuid=evaluation_run_uuid, ) except Exception: pass @@ -222,7 +224,7 @@ async def run(req: Request): trace_id = None if tr: try: - tr._req["outputs"] = result + tr.on_request_end(outputs=result, error=None) trace_id = await tr.submit_and_get_trace_id() except Exception: pass @@ -233,9 +235,9 @@ async def run(req: Request): except Exception as e: if tr: try: - tr._req["outputs"] = "".join(collected_chunks) - tr._req["error"] = str(e) - await tr._submit() + tr.on_request_end( + outputs="".join(collected_chunks), error=str(e) + ) except Exception: pass logger.error( diff --git a/gradient_adk/runtime/digitalocean_tracker.py b/gradient_adk/runtime/digitalocean_tracker.py index f269499..3a6ee6c 100644 --- a/gradient_adk/runtime/digitalocean_tracker.py +++ b/gradient_adk/runtime/digitalocean_tracker.py @@ -142,6 +142,8 @@ def on_request_start( inputs: Dict[str, Any], is_evaluation: bool = False, session_id: Optional[str] = None, + parent_context: Any = None, + evaluation_run_uuid: Optional[str] = None, ) -> contextvars.Token: """Start tracking a new request. diff --git a/gradient_adk/runtime/helpers.py b/gradient_adk/runtime/helpers.py index f9154b5..5be234f 100644 --- a/gradient_adk/runtime/helpers.py +++ b/gradient_adk/runtime/helpers.py @@ -10,12 +10,17 @@ """ from __future__ import annotations +import importlib.util import os +import warnings from typing import Optional, Callable, Dict, Any, Protocol from gradient_adk.cli.config.yaml_agent_config_manager import YamlAgentConfigManager from gradient_adk.runtime.digitalocean_tracker import DigitalOceanTracesTracker from gradient_adk.digital_ocean_api import AsyncDigitalOceanGenAI +from gradient_adk.runtime.multi_tracker import MultiTracker from gradient_adk.runtime.network_interceptor import setup_digitalocean_interception +from gradient_adk.runtime.otel_setup import env_bool +from gradient_adk.runtime.otlp_tracker import OTLPTracesTracker def _is_tracing_disabled() -> bool: @@ -27,7 +32,7 @@ def _is_tracing_disabled() -> bool: class InstrumentorProtocol(Protocol): """Protocol for instrumentor classes.""" - def install(self, tracker: DigitalOceanTracesTracker) -> None: + def install(self, tracker: Any) -> None: """Install the instrumentor with the given tracker.""" ... @@ -65,7 +70,7 @@ class InstrumentorRegistry: """ def __init__(self): - self._tracker: Optional[DigitalOceanTracesTracker] = None + self._tracker: Optional[Any] = None self._instrumentors: Dict[str, InstrumentorProtocol] = {} self._registrations: Dict[str, Dict[str, Any]] = {} self._config_reader = YamlAgentConfigManager() @@ -76,7 +81,7 @@ def _is_env_disabled(self, env_var: str) -> bool: val = os.environ.get(env_var, "").lower() return val in ("true", "1", "yes") - def _ensure_tracker(self) -> Optional[DigitalOceanTracesTracker]: + def _ensure_tracker(self) -> Optional[Any]: """ Create the shared tracker if not already created. @@ -99,17 +104,41 @@ def _ensure_tracker(self) -> Optional[DigitalOceanTracesTracker]: return None api_token = os.environ.get("DIGITALOCEAN_API_TOKEN") - if not api_token: + galileo_enabled = env_bool( + "DIGITALOCEAN_GALILEO_WRITE_ENABLED", default=True + ) + otlp_enabled = env_bool("DIGITALOCEAN_OTLP_WRITE_ENABLED", default=False) + if not api_token and not otlp_enabled: return None ws = self._config_reader.get_agent_name() dep = self._config_reader.get_agent_environment() - self._tracker = DigitalOceanTracesTracker( - client=AsyncDigitalOceanGenAI(api_token=api_token), - agent_workspace_name=ws, - agent_deployment_name=dep, - ) + trackers: list[Any] = [] + if galileo_enabled and api_token: + trackers.append( + DigitalOceanTracesTracker( + client=AsyncDigitalOceanGenAI(api_token=api_token), + agent_workspace_name=ws, + agent_deployment_name=dep, + ) + ) + if otlp_enabled: + trackers.append( + OTLPTracesTracker( + client=( + AsyncDigitalOceanGenAI(api_token=api_token) + if api_token + else None + ), + agent_workspace_name=ws, + agent_deployment_name=dep, + ) + ) + if not trackers: + return None + + self._tracker = trackers[0] if len(trackers) == 1 else MultiTracker(trackers) setup_digitalocean_interception() return self._tracker @@ -138,7 +167,7 @@ def register( "instrumentor_factory": instrumentor_factory, } - def install(self, name: str) -> Optional[DigitalOceanTracesTracker]: + def install(self, name: str) -> Optional[Any]: """ Install a specific registered instrumentor. @@ -175,7 +204,7 @@ def install(self, name: str) -> Optional[DigitalOceanTracesTracker]: except Exception: return None - def install_all(self) -> Optional[DigitalOceanTracesTracker]: + def install_all(self) -> Optional[Any]: """ Install all registered instrumentors that are available. @@ -199,7 +228,7 @@ def uninstall_all(self) -> None: for name in list(self._instrumentors.keys()): self.uninstall(name) - def get_tracker(self) -> Optional[DigitalOceanTracesTracker]: + def get_tracker(self) -> Optional[Any]: """Get the shared tracker instance.""" return self._tracker @@ -216,7 +245,7 @@ def get_installed_names(self) -> list[str]: registry = InstrumentorRegistry() -def get_tracker() -> Optional[DigitalOceanTracesTracker]: +def get_tracker() -> Optional[Any]: """Get the shared tracker from the global registry.""" return registry.get_tracker() @@ -229,17 +258,26 @@ def _register_langgraph() -> None: """Register LangGraph instrumentor if available.""" def is_available() -> bool: - try: - from langgraph.graph import StateGraph - - return True - except ImportError: - return False + return importlib.util.find_spec("langgraph.graph") is not None def factory(): - from gradient_adk.runtime.langgraph.langgraph_instrumentor import ( - LangGraphInstrumentor, - ) + with warnings.catch_warnings(): + try: + from langchain_core._api.deprecation import ( + LangChainPendingDeprecationWarning, + ) + + warnings.simplefilter( + "ignore", category=LangChainPendingDeprecationWarning + ) + except Exception: + warnings.filterwarnings( + "ignore", + message=r"The default value of `allowed_objects` will change.*", + ) + from gradient_adk.runtime.langgraph.langgraph_instrumentor import ( + LangGraphInstrumentor, + ) return LangGraphInstrumentor() @@ -255,12 +293,7 @@ def _register_pydanticai() -> None: """Register PydanticAI instrumentor if available.""" def is_available() -> bool: - try: - from pydantic_ai import Agent - - return True - except ImportError: - return False + return importlib.util.find_spec("pydantic_ai") is not None def factory(): from gradient_adk.runtime.pydanticai.pydanticai_instrumentor import ( @@ -281,12 +314,7 @@ def _register_crewai() -> None: """Register CrewAI instrumentor if available.""" def is_available() -> bool: - try: - from crewai import Crew - - return True - except ImportError: - return False + return importlib.util.find_spec("crewai") is not None def factory(): from gradient_adk.runtime.crewai.crewai_instrumentor import ( @@ -310,7 +338,7 @@ def register_all_instrumentors() -> None: _register_crewai() -def capture_all() -> Optional[DigitalOceanTracesTracker]: +def capture_all() -> Optional[Any]: """ Register and install all available instrumentors. diff --git a/gradient_adk/runtime/multi_tracker.py b/gradient_adk/runtime/multi_tracker.py new file mode 100644 index 0000000..2635266 --- /dev/null +++ b/gradient_adk/runtime/multi_tracker.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import asyncio +from typing import Any, Iterable, Optional + +from .digitalocean_tracker import DigitalOceanTracesTracker + + +class MultiTracker: + def __init__(self, trackers: Iterable[Any]) -> None: + self._trackers = [tracker for tracker in trackers if tracker is not None] + + def on_request_start( + self, + entrypoint: str, + inputs: dict[str, Any], + is_evaluation: bool = False, + session_id: Optional[str] = None, + parent_context: Any = None, + evaluation_run_uuid: Optional[str] = None, + ) -> None: + for tracker in self._trackers: + tracker.on_request_start( + entrypoint, + inputs, + is_evaluation=is_evaluation, + session_id=session_id, + parent_context=parent_context, + evaluation_run_uuid=evaluation_run_uuid, + ) + + def on_request_end(self, outputs: Any | None, error: Optional[str]) -> None: + for tracker in self._trackers: + tracker.on_request_end(outputs=outputs, error=error) + + def on_node_start(self, node: Any) -> None: + for tracker in self._trackers: + tracker.on_node_start(node) + + def on_node_end(self, node: Any, outputs: Any | None) -> None: + for tracker in self._trackers: + tracker.on_node_end(node, outputs) + + def on_node_error(self, node: Any, error: BaseException) -> None: + for tracker in self._trackers: + tracker.on_node_error(node, error) + + async def submit_and_get_trace_id(self) -> Optional[str]: + legacy_trace_id: Optional[str] = None + fallback_trace_id: Optional[str] = None + + for tracker in self._trackers: + trace_id = await tracker.submit_and_get_trace_id() + if trace_id and fallback_trace_id is None: + fallback_trace_id = trace_id + if trace_id and isinstance(tracker, DigitalOceanTracesTracker): + legacy_trace_id = trace_id + + return legacy_trace_id or fallback_trace_id + + async def aclose(self) -> None: + await asyncio.gather( + *(tracker.aclose() for tracker in self._trackers), + return_exceptions=True, + ) diff --git a/gradient_adk/runtime/network_interceptor.py b/gradient_adk/runtime/network_interceptor.py index 87155fa..661338d 100644 --- a/gradient_adk/runtime/network_interceptor.py +++ b/gradient_adk/runtime/network_interceptor.py @@ -7,6 +7,7 @@ from typing import Any, Callable, Dict, List, Optional, Set import httpx +from opentelemetry.propagate import inject import requests @@ -440,8 +441,30 @@ def hook(url: str, headers: Dict[str, str]) -> Dict[str, str]: return hook +def create_trace_context_hook(url_patterns: List[str]) -> RequestHook: + """Inject W3C trace context into tracked outbound requests.""" + + def hook(url: str, headers: Dict[str, str]) -> Dict[str, str]: + if not any(pattern in url for pattern in url_patterns): + return headers + + injected_headers = dict(headers) + inject(injected_headers) + return injected_headers + + return hook + + # URL classification helpers for different DigitalOcean services -INFERENCE_URL_PATTERNS = ["inference.do-ai.run", "inference.do-ai-test.run"] +INFERENCE_URL_PATTERNS = [ + "inference.do-ai.run", + "inference.do-ai-test.run", + "api.openai.com", + "api.anthropic.com", + "generativelanguage.googleapis.com", + "api.cohere.ai", + "api.x.ai", +] KBAAS_URL_PATTERNS = ["kbaas.do-ai.run", "kbaas.do-ai-test.run"] @@ -491,5 +514,6 @@ def setup_digitalocean_interception() -> None: url_patterns=all_patterns, ) intr.add_request_hook(ua_hook) + intr.add_request_hook(create_trace_context_hook(url_patterns=all_patterns)) intr.start_intercepting() \ No newline at end of file diff --git a/gradient_adk/runtime/otel_setup.py b/gradient_adk/runtime/otel_setup.py new file mode 100644 index 0000000..bbadec5 --- /dev/null +++ b/gradient_adk/runtime/otel_setup.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import os +from typing import Any, Dict, Optional + +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, ParentBased + +DEFAULT_OTLP_ENDPOINT = "http://localhost:4318" +DEFAULT_MAX_EVENT_BYTES = 256 * 1024 + + +def env_bool(name: str, default: bool = False) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes"} + + +def get_package_version() -> str: + try: + import importlib.metadata + + return importlib.metadata.version("gradient-adk") + except Exception: + try: + from gradient_adk import __version__ + + return __version__ + except Exception: + return "unknown" + + +def get_otlp_traces_endpoint() -> str: + traces_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") + if traces_endpoint: + return traces_endpoint + + base_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", DEFAULT_OTLP_ENDPOINT) + return f"{base_endpoint.rstrip('/')}/v1/traces" + + +def build_resource_attributes( + *, + agent_workspace_name: str, + agent_deployment_name: str, + team_id: Optional[str] = None, + org_id: Optional[str] = None, +) -> Dict[str, Any]: + attributes: Dict[str, Any] = { + "service.name": "adk-agent", + "service.version": get_package_version(), + "do.trace.visibility": "customer", + # The collector is authoritative in production, but these help + # local/dev runs where no resource processor is present. + "agent_workspace_name": agent_workspace_name, + "agent_deployment_name": agent_deployment_name, + } + if team_id: + attributes["team_id"] = str(team_id) + if org_id: + attributes["org_id"] = str(org_id) + return attributes + + +def create_tracer_provider( + *, + agent_workspace_name: str, + agent_deployment_name: str, + team_id: Optional[str] = None, + org_id: Optional[str] = None, +) -> TracerProvider: + resource = Resource.create( + build_resource_attributes( + agent_workspace_name=agent_workspace_name, + agent_deployment_name=agent_deployment_name, + team_id=team_id, + org_id=org_id, + ) + ) + + provider = TracerProvider( + resource=resource, + sampler=ParentBased(root=ALWAYS_ON), + ) + processor = BatchSpanProcessor( + OTLPSpanExporter(endpoint=get_otlp_traces_endpoint()), + export_timeout_millis=5000, + max_queue_size=2048, + max_export_batch_size=512, + ) + provider.add_span_processor(processor) + return provider diff --git a/gradient_adk/runtime/otlp_tracker.py b/gradient_adk/runtime/otlp_tracker.py new file mode 100644 index 0000000..23a9cc5 --- /dev/null +++ b/gradient_adk/runtime/otlp_tracker.py @@ -0,0 +1,495 @@ +from __future__ import annotations + +import contextvars +import json +import os +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Dict, Optional + +from opentelemetry import context as otel_context +from opentelemetry import trace +from opentelemetry.trace import SpanKind, Status, StatusCode + +from gradient_adk.digital_ocean_api import AsyncDigitalOceanGenAI + +from .interfaces import NodeExecution +from .otel_setup import ( + DEFAULT_MAX_EVENT_BYTES, + create_tracer_provider, +) + + +@dataclass +class LiveSpan: + node: NodeExecution + span: Any + token: object + + +@dataclass +class RequestState: + req: Dict[str, Any] = field(default_factory=dict) + live: Dict[str, LiveSpan] = field(default_factory=dict) + is_evaluation: bool = False + session_id: Optional[str] = None + evaluation_run_uuid: Optional[str] = None + root_span: Any | None = None + root_token: object | None = None + trace_id: Optional[str] = None + logs_enabled: bool = True + + +_request_state: contextvars.ContextVar[Optional[RequestState]] = contextvars.ContextVar( + "otlp_request_state", default=None +) + + +def _utc(dt: datetime | None = None) -> datetime: + if dt is None: + return datetime.now(timezone.utc) + return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc) + + +def _truthy_env(name: str, default: bool = True) -> bool: + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes"} + + +def _format_trace_id(trace_id: int) -> str: + return f"{trace_id:032x}" + + +def _derive_gen_ai_system(url: Optional[str]) -> Optional[str]: + if not url: + return None + lowered = url.lower() + if "openai.com" in lowered: + return "openai" + if "anthropic.com" in lowered: + return "anthropic" + if "generativelanguage.googleapis.com" in lowered or "googleapis.com" in lowered: + return "google" + if "cohere." in lowered: + return "cohere" + if "x.ai" in lowered: + return "xai" + if "inference.do-ai.run" in lowered or "inference.do-ai-test.run" in lowered: + return "digitalocean" + return None + + +def _extract_query(inputs: Any) -> Optional[str]: + if isinstance(inputs, str): + return inputs + if isinstance(inputs, dict): + for key in ("query", "input", "text", "prompt"): + value = inputs.get(key) + if isinstance(value, str): + return value + return None + + +def _json_text(value: Any) -> str: + try: + return json.dumps(value, ensure_ascii=True, default=str) + except Exception: + return str(value) + + +def _truncate_text(text: str, max_bytes: int = DEFAULT_MAX_EVENT_BYTES) -> str: + encoded = text.encode("utf-8", errors="replace") + if len(encoded) <= max_bytes: + return text + truncated = encoded[: max_bytes - len(b"...")] + return truncated.decode("utf-8", errors="ignore") + "..." + + +class OTLPTracesTracker: + """OpenTelemetry-backed trace tracker for gradient-adk. + + Tenant attribution (`team_id` / `org_id`) is sourced from env vars + only. The AAPP-1160 collector resource processor is the canonical + source of truth in production; the SDK does not perform any blocking + DigitalOcean API lookups at construction time. + """ + + def __init__( + self, + *, + client: Optional[AsyncDigitalOceanGenAI], + agent_workspace_name: str, + agent_deployment_name: str, + ) -> None: + self._client = client + self._ws = agent_workspace_name + self._dep = agent_deployment_name + self._closed = False + self._last_trace_id: Optional[str] = None + + team_id = os.environ.get("DIGITALOCEAN_TEAM_ID") + org_id = os.environ.get("DIGITALOCEAN_ORG_ID") + + self._provider = create_tracer_provider( + agent_workspace_name=agent_workspace_name, + agent_deployment_name=agent_deployment_name, + team_id=team_id, + org_id=org_id, + ) + self._tracer = self._provider.get_tracer("gradient_adk.runtime.otlp_tracker") + self._legacy_live: Dict[str, LiveSpan] = {} + + def _get_state(self) -> Optional[RequestState]: + return _request_state.get() + + def _logs_enabled(self) -> bool: + state = self._get_state() + if state is not None: + return state.logs_enabled + return _truthy_env("DIGITALOCEAN_CONVERSATION_LOGS_ENABLED", default=True) + + def _set_payload_event(self, span: Any, event_name: str, value: Any) -> None: + if value is None: + return + span.add_event(event_name, {"content": _truncate_text(_json_text(value))}) + + def _mark_redacted(self, span: Any, *, has_input: bool, has_output: bool) -> None: + if has_input: + span.set_attribute("gen_ai.input.redacted", True) + if has_output: + span.set_attribute("gen_ai.output.redacted", True) + + def _set_common_status( + self, + *, + span: Any, + metadata: Dict[str, Any], + error: Optional[str], + exception: Optional[BaseException] = None, + ) -> None: + if metadata.get("status_code") is not None: + span.set_attribute("http.status_code", int(metadata["status_code"])) + if exception is not None: + span.set_status(Status(StatusCode.ERROR, str(exception))) + span.record_exception(exception) + elif error: + span.set_status(Status(StatusCode.ERROR, error)) + span.record_exception(RuntimeError(error)) + else: + span.set_status(Status(StatusCode.OK)) + + def _populate_llm_span( + self, + *, + span: Any, + node: NodeExecution, + metadata: Dict[str, Any], + ) -> None: + span.set_attribute("gen_ai.operation.name", "chat") + + llm_url = metadata.get("llm_url") + gen_ai_system = _derive_gen_ai_system(llm_url) + if gen_ai_system: + span.set_attribute("gen_ai.system", gen_ai_system) + + llm_request = metadata.get("llm_request_payload", {}) or {} + llm_response = metadata.get("llm_response_payload", {}) or {} + + model = ( + llm_request.get("model") + if isinstance(llm_request, dict) + else metadata.get("model_name") + ) or metadata.get("model_name") + if model: + span.set_attribute("gen_ai.request.model", model) + + if isinstance(llm_request, dict): + temperature = llm_request.get("temperature") + if temperature is not None: + span.set_attribute("gen_ai.request.temperature", float(temperature)) + + if isinstance(llm_response, dict): + usage = llm_response.get("usage", {}) + if isinstance(usage, dict): + if usage.get("prompt_tokens") is not None: + span.set_attribute( + "gen_ai.usage.input_tokens", int(usage["prompt_tokens"]) + ) + if usage.get("completion_tokens") is not None: + span.set_attribute( + "gen_ai.usage.output_tokens", + int(usage["completion_tokens"]), + ) + if usage.get("total_tokens") is not None: + span.set_attribute( + "gen_ai.usage.total_tokens", int(usage["total_tokens"]) + ) + + ttft_ns = metadata.get("time_to_first_token_ns") + if ttft_ns is not None: + span.set_attribute( + "gen_ai.server.time_to_first_token", float(ttft_ns) / 1_000_000_000 + ) + + logs_enabled = self._logs_enabled() + if logs_enabled: + input_payload = ( + llm_request.get("messages") if isinstance(llm_request, dict) else None + ) or node.inputs + output_payload = ( + llm_response.get("choices") if isinstance(llm_response, dict) else None + ) or node.outputs + self._set_payload_event(span, "gen_ai.input.messages", input_payload) + self._set_payload_event(span, "gen_ai.output.messages", output_payload) + if isinstance(llm_request, dict) and llm_request.get("tools") is not None: + self._set_payload_event( + span, "gen_ai.request.tools", llm_request.get("tools") + ) + else: + self._mark_redacted( + span, + has_input=node.inputs is not None or bool(llm_request), + has_output=node.outputs is not None or bool(llm_response), + ) + + def _populate_tool_span( + self, + *, + span: Any, + node: NodeExecution, + metadata: Dict[str, Any], + ) -> None: + span.set_attribute("gen_ai.operation.name", "execute_tool") + span.set_attribute("gen_ai.tool.name", node.node_name) + tool_call_id = metadata.get("tool_call_id") + if tool_call_id: + span.set_attribute("gen_ai.tool.call.id", tool_call_id) + if self._logs_enabled(): + self._set_payload_event(span, "gen_ai.input.messages", node.inputs) + self._set_payload_event(span, "gen_ai.output.messages", node.outputs) + else: + self._mark_redacted( + span, + has_input=node.inputs is not None, + has_output=node.outputs is not None, + ) + + def _populate_retriever_span( + self, + *, + span: Any, + node: NodeExecution, + ) -> None: + span.set_attribute("gen_ai.operation.name", "retrieval") + query = _extract_query(node.inputs) + if query is not None: + span.set_attribute("gen_ai.retrieval.query", query) + if self._logs_enabled(): + self._set_payload_event(span, "gen_ai.input.messages", node.inputs) + self._set_payload_event(span, "gen_ai.output.messages", node.outputs) + else: + self._mark_redacted( + span, + has_input=node.inputs is not None, + has_output=node.outputs is not None, + ) + + def _populate_agent_span( + self, + *, + span: Any, + node: NodeExecution, + ) -> None: + span.set_attribute("gen_ai.operation.name", "invoke_agent") + if self._logs_enabled(): + self._set_payload_event(span, "gen_ai.input.messages", node.inputs) + self._set_payload_event(span, "gen_ai.output.messages", node.outputs) + else: + self._mark_redacted( + span, + has_input=node.inputs is not None, + has_output=node.outputs is not None, + ) + + def _populate_span_attributes( + self, + *, + span: Any, + node: NodeExecution, + metadata: Dict[str, Any], + ) -> None: + if metadata.get("is_llm_call"): + self._populate_llm_span(span=span, node=node, metadata=metadata) + return + if metadata.get("is_tool_call"): + self._populate_tool_span(span=span, node=node, metadata=metadata) + return + if metadata.get("is_retriever_call"): + self._populate_retriever_span(span=span, node=node) + return + if metadata.get("is_agent_call") or metadata.get("is_workflow"): + self._populate_agent_span(span=span, node=node) + return + + self._populate_tool_span(span=span, node=node, metadata=metadata) + + def on_request_start( + self, + entrypoint: str, + inputs: Dict[str, Any], + is_evaluation: bool = False, + session_id: Optional[str] = None, + parent_context: Any = None, + evaluation_run_uuid: Optional[str] = None, + ) -> None: + state = RequestState( + req={"entrypoint": entrypoint, "inputs": inputs}, + is_evaluation=is_evaluation, + session_id=session_id, + evaluation_run_uuid=evaluation_run_uuid, + logs_enabled=_truthy_env( + "DIGITALOCEAN_CONVERSATION_LOGS_ENABLED", default=True + ), + ) + parent = parent_context if parent_context is not None else otel_context.get_current() + root_span = self._tracer.start_span( + entrypoint, + context=parent, + kind=SpanKind.SERVER, + start_time=int(_utc().timestamp() * 1_000_000_000), + ) + root_span.set_attribute("gen_ai.operation.name", "invoke_agent") + if session_id: + root_span.set_attribute("gen_ai.conversation.id", session_id) + if evaluation_run_uuid: + root_span.set_attribute("gen_ai.evaluation.run_uuid", evaluation_run_uuid) + + state.root_span = root_span + state.trace_id = _format_trace_id(root_span.get_span_context().trace_id) + state.root_token = otel_context.attach(trace.set_span_in_context(root_span, parent)) + _request_state.set(state) + + def on_request_end(self, outputs: Any | None, error: Optional[str]) -> None: + state = self._get_state() + if state is None: + return + + state.req["outputs"] = outputs + state.req["error"] = error + + if state.root_span is not None: + logs_enabled = self._logs_enabled() + if logs_enabled: + self._set_payload_event( + state.root_span, "gen_ai.input.messages", state.req.get("inputs") + ) + self._set_payload_event( + state.root_span, "gen_ai.output.messages", state.req.get("outputs") + ) + else: + self._mark_redacted( + state.root_span, + has_input=state.req.get("inputs") is not None, + has_output=state.req.get("outputs") is not None, + ) + + self._set_common_status( + span=state.root_span, + metadata={}, + error=error, + ) + state.root_span.end() + + if state.root_token is not None: + otel_context.detach(state.root_token) + + self._last_trace_id = state.trace_id + _request_state.set(None) + + async def submit_and_get_trace_id(self) -> Optional[str]: + return self._last_trace_id + + def _span_kind_for_metadata(self, metadata: Dict[str, Any]) -> SpanKind: + if metadata.get("is_llm_call") or metadata.get("is_retriever_call"): + return SpanKind.CLIENT + if metadata.get("is_agent_call") or metadata.get("is_workflow"): + return SpanKind.INTERNAL + return SpanKind.INTERNAL + + def on_node_start(self, node: NodeExecution) -> None: + state = self._get_state() + metadata = node.metadata or {} + span = self._tracer.start_span( + node.node_name, + context=otel_context.get_current(), + kind=self._span_kind_for_metadata(metadata), + start_time=int(_utc(node.start_time).timestamp() * 1_000_000_000), + ) + token = otel_context.attach(trace.set_span_in_context(span)) + if state is not None: + state.live[node.node_id] = LiveSpan(node=node, span=span, token=token) + else: + self._legacy_live[node.node_id] = LiveSpan(node=node, span=span, token=token) + + def _finish_live_span( + self, + node: NodeExecution, + outputs: Any | None = None, + error: Optional[str] = None, + exception: Optional[BaseException] = None, + ) -> None: + state = self._get_state() + live = state.live.pop(node.node_id, None) if state is not None else None + if live is None: + live = self._legacy_live.pop(node.node_id, None) + if live is None: + return + + live.node.end_time = _utc() + live.node.outputs = outputs + if error: + live.node.error = error + + metadata = live.node.metadata or {} + self._populate_span_attributes(span=live.span, node=live.node, metadata=metadata) + self._set_common_status( + span=live.span, + metadata=metadata, + error=error, + exception=exception, + ) + + if metadata.get("duration_ns") is not None: + live.span.set_attribute("do.span.duration_ns", int(metadata["duration_ns"])) + elif live.node.start_time and live.node.end_time: + live.span.set_attribute( + "do.span.duration_ns", + int( + (live.node.end_time - live.node.start_time).total_seconds() + * 1_000_000_000 + ), + ) + + live.span.end(end_time=int(_utc(live.node.end_time).timestamp() * 1_000_000_000)) + otel_context.detach(live.token) + + def on_node_end(self, node: NodeExecution, outputs: Any | None) -> None: + self._finish_live_span(node, outputs=outputs) + + def on_node_error(self, node: NodeExecution, error: BaseException) -> None: + self._finish_live_span( + node, outputs=None, error=str(error), exception=error + ) + + async def aclose(self) -> None: + if self._closed: + return + try: + self._provider.force_flush() + self._provider.shutdown() + finally: + self._closed = True + if self._client is not None: + await self._client.aclose() diff --git a/gradient_adk/tracing.py b/gradient_adk/tracing.py index 373342b..c7500b5 100644 --- a/gradient_adk/tracing.py +++ b/gradient_adk/tracing.py @@ -249,6 +249,8 @@ async def async_gen_wrapper(*args, **kwargs): ) if captured: call = captured[0] + if call.url: + meta["llm_url"] = call.url if call.request_payload: meta["llm_request_payload"] = call.request_payload if call.response_payload: @@ -313,6 +315,8 @@ async def async_wrapper(*args, **kwargs): ) if captured: call = captured[0] + if call.url: + meta["llm_url"] = call.url if call.request_payload: meta["llm_request_payload"] = call.request_payload if call.response_payload: @@ -414,6 +418,8 @@ def sync_wrapper(*args, **kwargs): ) if captured: call = captured[0] + if call.url: + meta["llm_url"] = call.url if call.request_payload: meta["llm_request_payload"] = call.request_payload if call.response_payload: diff --git a/pyproject.toml b/pyproject.toml index b14c82f..acca219 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,9 @@ license = "Apache-2.0" dependencies = [ "fastapi>=0.104.0", "uvicorn[standard]>=0.24.0", + "opentelemetry-api>=1.27.0", + "opentelemetry-sdk>=1.27.0", + "opentelemetry-exporter-otlp-proto-http>=1.27.0", "pydantic>=2.5.0", "typing-extensions>=4.8.0", "typer>=0.12.0", diff --git a/tests/decorator_test.py b/tests/decorator_test.py index 5936ffe..9afff4e 100644 --- a/tests/decorator_test.py +++ b/tests/decorator_test.py @@ -1,6 +1,7 @@ import asyncio import pytest from fastapi.testclient import TestClient +from opentelemetry import trace from gradient_adk.decorator import entrypoint, run_server, RequestContext import gradient_adk.decorator as entrypoint_mod @@ -33,11 +34,32 @@ def __init__(self): self._req = {} self._session_id = None self._submitted_trace_id = None - - def on_request_start(self, name, inputs, is_evaluation=False, session_id=None): - self.started.append((name, inputs, is_evaluation, session_id)) + self._parent_context = None + self._evaluation_run_uuid = None + + def on_request_start( + self, + name, + inputs, + is_evaluation=False, + session_id=None, + parent_context=None, + evaluation_run_uuid=None, + ): + self.started.append( + ( + name, + inputs, + is_evaluation, + session_id, + parent_context, + evaluation_run_uuid, + ) + ) self._req = {"entrypoint": name, "inputs": inputs} self._session_id = session_id + self._parent_context = parent_context + self._evaluation_run_uuid = evaluation_run_uuid def on_request_end(self, outputs=None, error=None): self.ended.append((outputs, error)) @@ -723,7 +745,7 @@ def handler(data, context): # Verify that session_id was passed to the tracker assert tracker.started - # started tuple is (name, inputs, is_evaluation, session_id) + # started tuple is (name, inputs, is_evaluation, session_id, parent_context, evaluation_run_uuid) assert tracker.started[-1][3] == "trace-session-456" assert tracker._session_id == "trace-session-456" @@ -745,3 +767,39 @@ def handler(data, context): assert tracker.started assert tracker.started[-1][3] is None assert tracker._session_id is None + + +def test_traceparent_header_is_extracted_and_passed_to_tracker(patch_helpers): + """Test that traceparent is extracted into a parent context for the tracker.""" + tracker = patch_helpers + + @entrypoint + def handler(data, context): + return {"ok": True} + + fastapi_app = globals()["fastapi_app"] + traceparent = "00-1234567890abcdef1234567890abcdef-1234567890abcdef-01" + with TestClient(fastapi_app) as client: + r = client.post("/run", json={"test": 1}, headers={"traceparent": traceparent}) + assert r.status_code == 200 + + parent_span = trace.get_current_span(tracker._parent_context) + assert parent_span.get_span_context().trace_id == int( + "1234567890abcdef1234567890abcdef", 16 + ) + + +def test_evaluation_id_is_forwarded_as_evaluation_run_uuid(patch_helpers): + """Test that evaluation-id is forwarded to the tracker as evaluation_run_uuid.""" + tracker = patch_helpers + + @entrypoint + def handler(data, context): + return {"ok": True} + + fastapi_app = globals()["fastapi_app"] + with TestClient(fastapi_app) as client: + r = client.post("/run", json={"test": 1}, headers={"evaluation-id": "eval-789"}) + assert r.status_code == 200 + + assert tracker._evaluation_run_uuid == "eval-789" diff --git a/tests/runtime/helpers_test.py b/tests/runtime/helpers_test.py index 7eff4f0..752ea20 100644 --- a/tests/runtime/helpers_test.py +++ b/tests/runtime/helpers_test.py @@ -15,6 +15,7 @@ _register_crewai, _is_tracing_disabled, ) +from gradient_adk.runtime.otlp_tracker import OTLPTracesTracker # ----------------------------- @@ -79,6 +80,31 @@ def test_install_without_api_token_returns_none(fresh_registry, mock_instrumento mock_instrumentor.install.assert_not_called() +def test_install_with_otlp_enabled_without_api_token_returns_otlp_tracker( + fresh_registry, mock_instrumentor +): + """OTLP-only mode should still create a tracker without a DO API token.""" + fresh_registry.register( + name="test", + env_disable_var="TEST_DISABLE", + availability_check=lambda: True, + instrumentor_factory=lambda: mock_instrumentor, + ) + + with patch.dict(os.environ, {"DIGITALOCEAN_OTLP_WRITE_ENABLED": "true"}, clear=True): + with patch( + "gradient_adk.runtime.helpers.YamlAgentConfigManager.get_agent_name", + return_value="workspace-a", + ), patch( + "gradient_adk.runtime.helpers.YamlAgentConfigManager.get_agent_environment", + return_value="prod", + ): + result = fresh_registry.install("test") + + assert isinstance(result, OTLPTracesTracker) + mock_instrumentor.install.assert_called_once() + + def test_install_when_disabled_returns_none(fresh_registry, mock_instrumentor): """Test that install returns None when disabled via env var.""" fresh_registry.register( @@ -304,8 +330,12 @@ def test_pydanticai_availability_check(): check = test_registry._registrations["pydanticai"]["availability_check"] - # Since we're running tests with pydantic-ai installed, it should be available - assert check() is True + try: + from pydantic_ai import Agent # noqa: F401 + + assert check() is True + except ImportError: + assert check() is False def test_crewai_availability_check(): diff --git a/tests/runtime/test_multi_tracker.py b/tests/runtime/test_multi_tracker.py new file mode 100644 index 0000000..48332ec --- /dev/null +++ b/tests/runtime/test_multi_tracker.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from typing import Any, Optional + +import pytest + +from gradient_adk.runtime.digitalocean_tracker import DigitalOceanTracesTracker +from gradient_adk.runtime.multi_tracker import MultiTracker + + +class LegacyTrackerDouble(DigitalOceanTracesTracker): + def __init__(self, trace_id: Optional[str]) -> None: + self.trace_id = trace_id + self.calls: list[tuple[str, tuple[Any, ...], dict[str, Any]]] = [] + + def on_request_start(self, *args, **kwargs): + self.calls.append(("on_request_start", args, kwargs)) + + def on_request_end(self, *args, **kwargs): + self.calls.append(("on_request_end", args, kwargs)) + + def on_node_start(self, *args, **kwargs): + self.calls.append(("on_node_start", args, kwargs)) + + def on_node_end(self, *args, **kwargs): + self.calls.append(("on_node_end", args, kwargs)) + + def on_node_error(self, *args, **kwargs): + self.calls.append(("on_node_error", args, kwargs)) + + async def submit_and_get_trace_id(self): + self.calls.append(("submit_and_get_trace_id", (), {})) + return self.trace_id + + async def aclose(self): + self.calls.append(("aclose", (), {})) + + +class TrackerDouble: + def __init__(self, trace_id: Optional[str]) -> None: + self.trace_id = trace_id + self.calls: list[tuple[str, tuple[Any, ...], dict[str, Any]]] = [] + + def on_request_start(self, *args, **kwargs): + self.calls.append(("on_request_start", args, kwargs)) + + def on_request_end(self, *args, **kwargs): + self.calls.append(("on_request_end", args, kwargs)) + + def on_node_start(self, *args, **kwargs): + self.calls.append(("on_node_start", args, kwargs)) + + def on_node_end(self, *args, **kwargs): + self.calls.append(("on_node_end", args, kwargs)) + + def on_node_error(self, *args, **kwargs): + self.calls.append(("on_node_error", args, kwargs)) + + async def submit_and_get_trace_id(self): + self.calls.append(("submit_and_get_trace_id", (), {})) + return self.trace_id + + async def aclose(self): + self.calls.append(("aclose", (), {})) + + +@pytest.mark.asyncio +async def test_multi_tracker_prefers_legacy_trace_id(): + legacy = LegacyTrackerDouble("legacy-trace") + otlp = TrackerDouble("otlp-trace") + tracker = MultiTracker([legacy, otlp]) + + trace_id = await tracker.submit_and_get_trace_id() + assert trace_id == "legacy-trace" + + +@pytest.mark.asyncio +async def test_multi_tracker_falls_back_to_otlp_trace_id(): + legacy = LegacyTrackerDouble(None) + otlp = TrackerDouble("otlp-trace") + tracker = MultiTracker([legacy, otlp]) + + trace_id = await tracker.submit_and_get_trace_id() + assert trace_id == "otlp-trace" + + +def test_multi_tracker_fanout_calls_all_trackers(): + legacy = LegacyTrackerDouble("legacy-trace") + otlp = TrackerDouble("otlp-trace") + tracker = MultiTracker([legacy, otlp]) + + tracker.on_request_start( + "agent", + {"input": "hello"}, + is_evaluation=True, + session_id="session-1", + parent_context="parent", + evaluation_run_uuid="eval-1", + ) + tracker.on_request_end(outputs={"done": True}, error=None) + + assert legacy.calls[0][0] == "on_request_start" + assert otlp.calls[0][0] == "on_request_start" + assert legacy.calls[0][2]["evaluation_run_uuid"] == "eval-1" + assert otlp.calls[0][2]["parent_context"] == "parent" + assert legacy.calls[1][0] == "on_request_end" + assert otlp.calls[1][0] == "on_request_end" diff --git a/tests/runtime/test_otlp_tracker.py b/tests/runtime/test_otlp_tracker.py new file mode 100644 index 0000000..d81ecdc --- /dev/null +++ b/tests/runtime/test_otlp_tracker.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +from datetime import datetime, timezone + +import pytest +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, ParentBased + +from gradient_adk.runtime.interfaces import NodeExecution +from gradient_adk.runtime.otel_setup import build_resource_attributes +from gradient_adk.runtime.otlp_tracker import OTLPTracesTracker + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +@pytest.fixture +def exporter() -> InMemorySpanExporter: + return InMemorySpanExporter() + + +@pytest.fixture +def tracker(monkeypatch, exporter: InMemorySpanExporter) -> OTLPTracesTracker: + def fake_create_tracer_provider( + *, + agent_workspace_name: str, + agent_deployment_name: str, + team_id: str | None = None, + org_id: str | None = None, + ) -> TracerProvider: + provider = TracerProvider( + resource=Resource.create( + build_resource_attributes( + agent_workspace_name=agent_workspace_name, + agent_deployment_name=agent_deployment_name, + team_id=team_id, + org_id=org_id, + ) + ), + sampler=ParentBased(root=ALWAYS_ON), + ) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider + + monkeypatch.setattr( + "gradient_adk.runtime.otlp_tracker.create_tracer_provider", + fake_create_tracer_provider, + ) + monkeypatch.setenv("DIGITALOCEAN_TEAM_ID", "123") + monkeypatch.setenv("DIGITALOCEAN_ORG_ID", "456") + monkeypatch.setenv("DIGITALOCEAN_CONVERSATION_LOGS_ENABLED", "true") + return OTLPTracesTracker( + client=None, + agent_workspace_name="workspace-a", + agent_deployment_name="prod", + ) + + +def test_llm_span_sets_semconv_and_resource_attributes( + tracker: OTLPTracesTracker, exporter: InMemorySpanExporter +): + tracker.on_request_start( + "agent_entrypoint", + {"prompt": "hello"}, + session_id="session-123", + evaluation_run_uuid="eval-456", + ) + node = NodeExecution( + node_id="llm-1", + node_name="call_model", + framework="custom", + start_time=_now(), + inputs={"messages": [{"role": "user", "content": "hello"}]}, + metadata={ + "is_llm_call": True, + "llm_url": "https://api.openai.com/v1/chat/completions", + "llm_request_payload": { + "model": "gpt-4o-mini", + "temperature": 0.2, + "messages": [{"role": "user", "content": "hello"}], + "tools": [{"type": "function", "name": "search"}], + }, + "llm_response_payload": { + "choices": [{"message": {"role": "assistant", "content": "hi"}}], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 4, + "total_tokens": 14, + }, + }, + "time_to_first_token_ns": 100_000_000, + }, + ) + + tracker.on_node_start(node) + tracker.on_node_end(node, {"content": "hi"}) + tracker.on_request_end({"response": "hi"}, None) + + spans = exporter.get_finished_spans() + assert {span.name for span in spans} == {"agent_entrypoint", "call_model"} + + llm_span = next(span for span in spans if span.name == "call_model") + root_span = next(span for span in spans if span.name == "agent_entrypoint") + + assert llm_span.attributes["gen_ai.operation.name"] == "chat" + assert llm_span.attributes["gen_ai.system"] == "openai" + assert llm_span.attributes["gen_ai.request.model"] == "gpt-4o-mini" + assert llm_span.attributes["gen_ai.request.temperature"] == 0.2 + assert llm_span.attributes["gen_ai.usage.input_tokens"] == 10 + assert llm_span.attributes["gen_ai.usage.output_tokens"] == 4 + assert llm_span.attributes["gen_ai.usage.total_tokens"] == 14 + assert llm_span.attributes["gen_ai.server.time_to_first_token"] == 0.1 + assert {event.name for event in llm_span.events} >= { + "gen_ai.input.messages", + "gen_ai.output.messages", + "gen_ai.request.tools", + } + + assert root_span.attributes["gen_ai.operation.name"] == "invoke_agent" + assert root_span.attributes["gen_ai.conversation.id"] == "session-123" + assert root_span.attributes["gen_ai.evaluation.run_uuid"] == "eval-456" + assert root_span.resource.attributes["service.name"] == "adk-agent" + assert root_span.resource.attributes["do.trace.visibility"] == "customer" + assert root_span.resource.attributes["team_id"] == "123" + assert root_span.resource.attributes["org_id"] == "456" + + +def test_tool_and_retriever_spans_map_expected_attributes( + tracker: OTLPTracesTracker, exporter: InMemorySpanExporter +): + tracker.on_request_start("agent_entrypoint", {"query": "weather"}) + + tool_node = NodeExecution( + node_id="tool-1", + node_name="weather_tool", + framework="custom", + start_time=_now(), + inputs={"city": "nyc"}, + metadata={"is_tool_call": True, "tool_call_id": "call-123"}, + ) + retriever_node = NodeExecution( + node_id="retriever-1", + node_name="knowledge_base", + framework="custom", + start_time=_now(), + inputs={"query": "latest weather"}, + metadata={"is_retriever_call": True}, + ) + + tracker.on_node_start(tool_node) + tracker.on_node_end(tool_node, {"result": "sunny"}) + tracker.on_node_start(retriever_node) + tracker.on_node_end(retriever_node, {"hits": 2}) + tracker.on_request_end({"done": True}, None) + + spans = exporter.get_finished_spans() + tool_span = next(span for span in spans if span.name == "weather_tool") + retriever_span = next(span for span in spans if span.name == "knowledge_base") + + assert tool_span.attributes["gen_ai.operation.name"] == "execute_tool" + assert tool_span.attributes["gen_ai.tool.name"] == "weather_tool" + assert tool_span.attributes["gen_ai.tool.call.id"] == "call-123" + assert retriever_span.attributes["gen_ai.operation.name"] == "retrieval" + assert retriever_span.attributes["gen_ai.retrieval.query"] == "latest weather" + + +def test_error_span_records_error_status( + tracker: OTLPTracesTracker, exporter: InMemorySpanExporter +): + tracker.on_request_start("agent_entrypoint", {"prompt": "hello"}) + node = NodeExecution( + node_id="tool-err", + node_name="failing_tool", + framework="custom", + start_time=_now(), + inputs={"x": 1}, + metadata={"is_tool_call": True}, + ) + + tracker.on_node_start(node) + tracker.on_node_error(node, RuntimeError("boom")) + tracker.on_request_end(None, "boom") + + spans = exporter.get_finished_spans() + tool_span = next(span for span in spans if span.name == "failing_tool") + root_span = next(span for span in spans if span.name == "agent_entrypoint") + + assert tool_span.status.status_code.name == "ERROR" + assert root_span.status.status_code.name == "ERROR" diff --git a/tests/runtime/test_otlp_tracker_lifecycle.py b/tests/runtime/test_otlp_tracker_lifecycle.py new file mode 100644 index 0000000..3b88145 --- /dev/null +++ b/tests/runtime/test_otlp_tracker_lifecycle.py @@ -0,0 +1,164 @@ +"""Lifecycle and exception-fidelity tests for OTLPTracesTracker.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import MagicMock + +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, ParentBased + +from gradient_adk.runtime.interfaces import NodeExecution +from gradient_adk.runtime.otel_setup import build_resource_attributes +from gradient_adk.runtime.otlp_tracker import OTLPTracesTracker + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +def _patch_tracer_provider(monkeypatch, exporter: InMemorySpanExporter) -> None: + def fake_create_tracer_provider( + *, + agent_workspace_name: str, + agent_deployment_name: str, + team_id: str | None = None, + org_id: str | None = None, + ) -> TracerProvider: + provider = TracerProvider( + resource=Resource.create( + build_resource_attributes( + agent_workspace_name=agent_workspace_name, + agent_deployment_name=agent_deployment_name, + team_id=team_id, + org_id=org_id, + ) + ), + sampler=ParentBased(root=ALWAYS_ON), + ) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider + + monkeypatch.setattr( + "gradient_adk.runtime.otlp_tracker.create_tracer_provider", + fake_create_tracer_provider, + ) + + +def test_init_does_not_call_digitalocean_api(monkeypatch): + """The tracker must not perform any blocking DO API lookup at construction. + + The collector (AAPP-1160) is the canonical source for tenant attribution + in production; the SDK only honors env vars and falls back to the + collector. Calling the DO API at __init__ time would block app startup. + """ + exporter = InMemorySpanExporter() + _patch_tracer_provider(monkeypatch, exporter) + monkeypatch.delenv("DIGITALOCEAN_TEAM_ID", raising=False) + monkeypatch.delenv("DIGITALOCEAN_ORG_ID", raising=False) + + client = MagicMock() + client.list_agent_workspaces = MagicMock() + client.aclose = MagicMock() + + OTLPTracesTracker( + client=client, + agent_workspace_name="workspace-a", + agent_deployment_name="prod", + ) + + client.list_agent_workspaces.assert_not_called() + + +def test_on_node_error_records_original_exception_type(monkeypatch): + """The original exception (class + traceback) must reach record_exception. + + Wrapping into a fresh RuntimeError loses both the class and the original + traceback, which makes downstream alerting/triaging much harder. + """ + exporter = InMemorySpanExporter() + _patch_tracer_provider(monkeypatch, exporter) + monkeypatch.setenv("DIGITALOCEAN_CONVERSATION_LOGS_ENABLED", "true") + + tracker = OTLPTracesTracker( + client=None, + agent_workspace_name="workspace-a", + agent_deployment_name="prod", + ) + tracker.on_request_start("agent_entrypoint", {"prompt": "hello"}) + + node = NodeExecution( + node_id="tool-err", + node_name="failing_tool", + framework="custom", + start_time=_now(), + inputs={"x": 1}, + metadata={"is_tool_call": True}, + ) + tracker.on_node_start(node) + + class CustomFailure(ValueError): + pass + + raised = CustomFailure("kaboom") + tracker.on_node_error(node, raised) + tracker.on_request_end(None, "kaboom") + + tool_span = next( + span for span in exporter.get_finished_spans() if span.name == "failing_tool" + ) + assert tool_span.status.status_code.name == "ERROR" + exception_events = [ + event for event in tool_span.events if event.name == "exception" + ] + assert exception_events, "expected record_exception to emit an 'exception' event" + attrs = exception_events[0].attributes + assert attrs["exception.type"].endswith("CustomFailure") + assert attrs["exception.message"] == "kaboom" + + +def test_logs_enabled_is_snapshotted_at_request_start(monkeypatch): + """Once a request starts, mid-request env flips must not change redaction. + + Redaction is per-request: the flag is read once at on_request_start and + cached on the per-request state. Mid-request env mutation (e.g. by other + threads or by a buggy library) must not partially redact a single trace. + """ + exporter = InMemorySpanExporter() + _patch_tracer_provider(monkeypatch, exporter) + monkeypatch.setenv("DIGITALOCEAN_CONVERSATION_LOGS_ENABLED", "true") + + tracker = OTLPTracesTracker( + client=None, + agent_workspace_name="workspace-a", + agent_deployment_name="prod", + ) + tracker.on_request_start("agent_entrypoint", {"prompt": "hello"}) + + monkeypatch.setenv("DIGITALOCEAN_CONVERSATION_LOGS_ENABLED", "false") + + node = NodeExecution( + node_id="tool-1", + node_name="weather_tool", + framework="custom", + start_time=_now(), + inputs={"city": "nyc"}, + metadata={"is_tool_call": True}, + ) + tracker.on_node_start(node) + tracker.on_node_end(node, {"result": "sunny"}) + tracker.on_request_end({"done": True}, None) + + tool_span = next( + span for span in exporter.get_finished_spans() if span.name == "weather_tool" + ) + event_names = {event.name for event in tool_span.events} + assert "gen_ai.input.messages" in event_names + assert "gen_ai.output.messages" in event_names + assert tool_span.attributes.get("gen_ai.input.redacted") is None + assert tool_span.attributes.get("gen_ai.output.redacted") is None diff --git a/tests/runtime/test_payload_cap.py b/tests/runtime/test_payload_cap.py new file mode 100644 index 0000000..6773944 --- /dev/null +++ b/tests/runtime/test_payload_cap.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from datetime import datetime, timezone + +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, ParentBased + +from gradient_adk.runtime.interfaces import NodeExecution +from gradient_adk.runtime.otel_setup import ( + DEFAULT_MAX_EVENT_BYTES, + build_resource_attributes, +) +from gradient_adk.runtime.otlp_tracker import OTLPTracesTracker + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +def test_payload_events_are_capped_to_max_bytes(monkeypatch): + exporter = InMemorySpanExporter() + + def fake_create_tracer_provider( + *, + agent_workspace_name: str, + agent_deployment_name: str, + team_id: str | None = None, + org_id: str | None = None, + ) -> TracerProvider: + provider = TracerProvider( + resource=Resource.create( + build_resource_attributes( + agent_workspace_name=agent_workspace_name, + agent_deployment_name=agent_deployment_name, + team_id=team_id, + org_id=org_id, + ) + ), + sampler=ParentBased(root=ALWAYS_ON), + ) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider + + monkeypatch.setattr( + "gradient_adk.runtime.otlp_tracker.create_tracer_provider", + fake_create_tracer_provider, + ) + monkeypatch.setenv("DIGITALOCEAN_CONVERSATION_LOGS_ENABLED", "true") + + tracker = OTLPTracesTracker( + client=None, + agent_workspace_name="workspace-a", + agent_deployment_name="prod", + ) + huge_value = "x" * (1024 * 1024) + tracker.on_request_start("agent_entrypoint", {"prompt": huge_value}) + node = NodeExecution( + node_id="tool-1", + node_name="tool_call", + framework="custom", + start_time=_now(), + inputs={"prompt": huge_value}, + metadata={"is_tool_call": True}, + ) + + tracker.on_node_start(node) + tracker.on_node_end(node, {"result": huge_value}) + tracker.on_request_end({"response": huge_value}, None) + + tool_span = next(span for span in exporter.get_finished_spans() if span.name == "tool_call") + payload = tool_span.events[0].attributes["content"] + assert payload.endswith("...") + assert len(payload.encode("utf-8")) <= DEFAULT_MAX_EVENT_BYTES diff --git a/tests/runtime/test_redaction.py b/tests/runtime/test_redaction.py new file mode 100644 index 0000000..76e8d5d --- /dev/null +++ b/tests/runtime/test_redaction.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from datetime import datetime, timezone + +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, ParentBased + +from gradient_adk.runtime.interfaces import NodeExecution +from gradient_adk.runtime.otel_setup import build_resource_attributes +from gradient_adk.runtime.otlp_tracker import OTLPTracesTracker + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +def test_redaction_omits_payload_events_and_sets_markers(monkeypatch): + exporter = InMemorySpanExporter() + + def fake_create_tracer_provider( + *, + agent_workspace_name: str, + agent_deployment_name: str, + team_id: str | None = None, + org_id: str | None = None, + ) -> TracerProvider: + provider = TracerProvider( + resource=Resource.create( + build_resource_attributes( + agent_workspace_name=agent_workspace_name, + agent_deployment_name=agent_deployment_name, + team_id=team_id, + org_id=org_id, + ) + ), + sampler=ParentBased(root=ALWAYS_ON), + ) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider + + monkeypatch.setattr( + "gradient_adk.runtime.otlp_tracker.create_tracer_provider", + fake_create_tracer_provider, + ) + monkeypatch.setenv("DIGITALOCEAN_CONVERSATION_LOGS_ENABLED", "false") + + tracker = OTLPTracesTracker( + client=None, + agent_workspace_name="workspace-a", + agent_deployment_name="prod", + ) + tracker.on_request_start("agent_entrypoint", {"prompt": "secret"}) + node = NodeExecution( + node_id="llm-1", + node_name="call_model", + framework="custom", + start_time=_now(), + inputs={"messages": [{"role": "user", "content": "secret"}]}, + metadata={ + "is_llm_call": True, + "llm_request_payload": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "secret"}], + }, + "llm_response_payload": { + "choices": [{"message": {"role": "assistant", "content": "redacted"}}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + }, + }, + ) + + tracker.on_node_start(node) + tracker.on_node_end(node, {"content": "redacted"}) + tracker.on_request_end({"response": "redacted"}, None) + + spans = exporter.get_finished_spans() + llm_span = next(span for span in spans if span.name == "call_model") + root_span = next(span for span in spans if span.name == "agent_entrypoint") + + assert llm_span.attributes["gen_ai.usage.input_tokens"] == 3 + assert llm_span.attributes["gen_ai.usage.output_tokens"] == 2 + assert llm_span.attributes["gen_ai.input.redacted"] is True + assert llm_span.attributes["gen_ai.output.redacted"] is True + assert root_span.attributes["gen_ai.input.redacted"] is True + assert root_span.attributes["gen_ai.output.redacted"] is True + assert [event.name for event in llm_span.events] == [] + assert [event.name for event in root_span.events] == [] diff --git a/tests/runtime/test_sampling.py b/tests/runtime/test_sampling.py new file mode 100644 index 0000000..4a229ac --- /dev/null +++ b/tests/runtime/test_sampling.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +from opentelemetry import trace +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, ParentBased +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, TraceState + +from gradient_adk.runtime.network_interceptor import create_trace_context_hook +from gradient_adk.runtime.otel_setup import build_resource_attributes +from gradient_adk.runtime.otlp_tracker import OTLPTracesTracker + + +def _provider(exporter: InMemorySpanExporter) -> TracerProvider: + provider = TracerProvider( + resource=Resource.create( + build_resource_attributes( + agent_workspace_name="workspace-a", + agent_deployment_name="prod", + ) + ), + sampler=ParentBased(root=ALWAYS_ON), + ) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider + + +def _remote_parent_context(sampled: bool): + trace_flags = TraceFlags(0x01 if sampled else 0x00) + span_context = SpanContext( + trace_id=int("1234567890abcdef1234567890abcdef", 16), + span_id=int("1234567890abcdef", 16), + is_remote=True, + trace_flags=trace_flags, + trace_state=TraceState(), + ) + return trace.set_span_in_context(NonRecordingSpan(span_context)) + + +def test_parent_based_sampling_respects_remote_unsampled_parent(monkeypatch): + exporter = InMemorySpanExporter() + monkeypatch.setattr( + "gradient_adk.runtime.otlp_tracker.create_tracer_provider", + lambda **_: _provider(exporter), + ) + tracker = OTLPTracesTracker( + client=None, + agent_workspace_name="workspace-a", + agent_deployment_name="prod", + ) + + tracker.on_request_start( + "agent_entrypoint", + {"prompt": "hello"}, + parent_context=_remote_parent_context(sampled=False), + ) + headers = create_trace_context_hook(["api.openai.com"])( + "https://api.openai.com/v1/chat/completions", {} + ) + tracker.on_request_end({"response": "ok"}, None) + + assert not exporter.get_finished_spans() + assert headers["traceparent"].endswith("-00") + + +def test_parent_based_sampling_respects_remote_sampled_parent(monkeypatch): + exporter = InMemorySpanExporter() + monkeypatch.setattr( + "gradient_adk.runtime.otlp_tracker.create_tracer_provider", + lambda **_: _provider(exporter), + ) + tracker = OTLPTracesTracker( + client=None, + agent_workspace_name="workspace-a", + agent_deployment_name="prod", + ) + + tracker.on_request_start( + "agent_entrypoint", + {"prompt": "hello"}, + parent_context=_remote_parent_context(sampled=True), + ) + headers = create_trace_context_hook(["api.openai.com"])( + "https://api.openai.com/v1/chat/completions", {} + ) + tracker.on_request_end({"response": "ok"}, None) + + assert len(exporter.get_finished_spans()) == 1 + assert headers["traceparent"].endswith("-01")