komodor-agentops 0.1.0__py3-none-any.whl

komodor_agentops/__init__.py

@@ -0,0 +1,92 @@
+from agentops_rpc import (
+    GRADER_CAPABILITY,
+    REVIEW_CAPABILITY,
+    USE_CASE_INPUT_KEY,
+    A2AMessage,
+    A2AMessagePart,
+    A2ATask,
+    GradeContext,
+    GradeResult,
+    RunJob,
+    UseCase,
+)
+from komodor_agentops.artifacts import ArtifactClient, get_artifact_client
+from komodor_agentops.client import AgentOpsClient as AgentOps
+from komodor_agentops.client import get_client
+from komodor_agentops.core.context import TraceContext
+from komodor_agentops.core.events import AgentOpsEvent
+from komodor_agentops.core.observe import observe
+from komodor_agentops.core.secrets import apply_secret_to_env, get_run_secrets, get_secret
+from komodor_agentops.core.skills import AgentMarkdown
+from komodor_agentops.core.transport import AgentOpsClient as TransportClient
+from komodor_agentops.core.types import AgentSpec
+from komodor_agentops.core.usage import drain_usage, record_usage, reset_usage, usage_from_anthropic_response
+from komodor_agentops.github import (
+    GitHubIntegrationError,
+    github_api_get,
+    github_api_patch,
+    github_api_post,
+    github_graphql,
+    github_installation_token,
+)
+from komodor_agentops.grade_context import fetch_grade_context
+from komodor_agentops.worker import (
+    AgentOpsWorker,
+    collect_env_secrets,
+    fetch_controlplane_secrets,
+    load_env_files,
+    run_worker_once,
+    run_worker_once_sync,
+)
+from komodor_agentops.worker.a2a.streaming import stream_status_text
+from komodor_agentops.worker.result import DESCRIPTIVE_OUTPUT_KEY, result_task
+from komodor_agentops.worker.server import create_worker_app, event_sink_from_client, register_worker_endpoint
+
+__all__ = [
+    "ArtifactClient",
+    "get_artifact_client",
+    "AgentOps",
+    "AgentMarkdown",
+    "AgentOpsEvent",
+    "AgentOpsWorker",
+    "A2AMessage",
+    "A2AMessagePart",
+    "A2ATask",
+    "AgentSpec",
+    "DESCRIPTIVE_OUTPUT_KEY",
+    "GRADER_CAPABILITY",
+    "REVIEW_CAPABILITY",
+    "USE_CASE_INPUT_KEY",
+    "GitHubIntegrationError",
+    "GradeContext",
+    "GradeResult",
+    "fetch_grade_context",
+    "RunJob",
+    "UseCase",
+    "apply_secret_to_env",
+    "collect_env_secrets",
+    "fetch_controlplane_secrets",
+    "TraceContext",
+    "TransportClient",
+    "create_worker_app",
+    "event_sink_from_client",
+    "get_client",
+    "get_run_secrets",
+    "github_api_get",
+    "github_api_patch",
+    "github_api_post",
+    "github_graphql",
+    "github_installation_token",
+    "get_secret",
+    "load_env_files",
+    "observe",
+    "record_usage",
+    "register_worker_endpoint",
+    "result_task",
+    "reset_usage",
+    "run_worker_once",
+    "run_worker_once_sync",
+    "drain_usage",
+    "stream_status_text",
+    "usage_from_anthropic_response",
+]

komodor_agentops/adapters/__init__.py

@@ -0,0 +1 @@
+"""Adapter modules for third-party agent frameworks."""

komodor_agentops/adapters/adk/__init__.py

@@ -0,0 +1,4 @@
+from komodor_agentops.adapters.adk.a2a_adapter import make_agentops_adk_interceptor, run_adk_pipeline
+from komodor_agentops.adapters.adk.adapter import instrument_adk_run
+
+__all__ = ["instrument_adk_run", "make_agentops_adk_interceptor", "run_adk_pipeline"]

komodor_agentops/adapters/adk/a2a_adapter.py

@@ -0,0 +1,313 @@
+from __future__ import annotations
+
+import uuid
+from typing import Any
+
+from google.adk.runners import Runner
+from google.adk.sessions import InMemorySessionService
+from google.genai import types
+
+from komodor_agentops import client as client_mod
+from komodor_agentops.adapters.adk.adapter import (
+    instrument_adk_run,
+    instrument_agent_model_calls,
+    usage_from_adk_event,
+    usage_from_adk_model,
+)
+from komodor_agentops.core.events import AgentOpsEvent
+
+
+def _extract_prompt(message: Any) -> str:
+    if message is None:
+        return ""
+    parts = []
+    for part in getattr(message, "parts", None) or []:
+        actual = getattr(part, "root", part)
+        text = getattr(actual, "text", None)
+        if text:
+            parts.append(str(text))
+    return " ".join(parts).strip()
+
+
+def make_agentops_adk_interceptor(root_agent_name: str) -> Any:
+    """Return an ADK ExecuteInterceptor that emits all AgentOps hook events.
+
+    Uses per-session state keyed by context_id so concurrent requests don't
+    interfere. No-op when no AgentOpsClient is active.
+    """
+    # Must stay inline: tests patch sys.modules["google.adk.a2a.executor.config"]
+    # before calling this function to inject a fake ExecuteInterceptor.
+    # A top-level import would resolve before the patch is applied.
+    from google.adk.a2a.executor.config import ExecuteInterceptor  # noqa: PLC0415
+
+    _states: dict[str, dict] = {}
+
+    async def before_agent(context: Any) -> Any:
+        ctx_id = context.context_id
+        prompt = _extract_prompt(context.message)
+        _states[ctx_id] = {
+            "prompt": prompt,
+            "current_subagent": None,
+            "subagent_span_ids": {},
+            "pending_tool_spans": {},
+            # Per-LLM-turn index (see adapter.py): bumped per model event issuing tool
+            # calls, stamped on each tool message so the UI can group the transcript by
+            # round; kept in per-session state since after_event is called once per event.
+            "turn": 0,
+            "pending_tool_turns": {},
+            "default_usage": usage_from_adk_model(getattr(context, "agent", None)),
+        }
+        client = client_mod._active_client
+        if client is not None:
+            await client.enqueue(
+                AgentOpsEvent(
+                    run_id=ctx_id,
+                    event_type="span.started",
+                    span_id=ctx_id,
+                    kind="agent",
+                    name=root_agent_name,
+                    hook_event_type="session.start",
+                )
+            )
+            if prompt:
+                await client.enqueue(
+                    AgentOpsEvent(
+                        run_id=ctx_id,
+                        event_type="message.created",
+                        role="user",
+                        content=prompt,
+                        hook_event_type="prompt.submit",
+                    )
+                )
+        return context
+
+    async def after_event(exec_ctx: Any, a2a_event: Any, adk_event: Any) -> Any:
+        session_id = exec_ctx.session_id
+        state = _states.get(session_id)
+        if state is None:
+            return a2a_event
+        client = client_mod._active_client
+        if client is None:
+            return a2a_event
+
+        run_id = session_id
+        author: str = getattr(adk_event, "author", "") or ""
+        usage = usage_from_adk_event(adk_event)
+        if usage:
+            usage = {**state.get("default_usage", {}), **usage}
+            await client.enqueue(
+                AgentOpsEvent(
+                    run_id=run_id,
+                    event_type="token.created",
+                    payload=usage,
+                )
+            )
+
+        if author and author != root_agent_name and author != state["current_subagent"]:
+            if state["current_subagent"] is not None:
+                await client.enqueue(
+                    AgentOpsEvent(
+                        run_id=run_id,
+                        event_type="span.ended",
+                        span_id=state["subagent_span_ids"].get(state["current_subagent"], state["current_subagent"]),
+                        parent_span_id=run_id,
+                        kind="agent",
+                        name=state["current_subagent"],
+                        hook_event_type="subagent.end",
+                    )
+                )
+            state["current_subagent"] = author
+            span_id = str(uuid.uuid4())
+            state["subagent_span_ids"][author] = span_id
+            await client.enqueue(
+                AgentOpsEvent(
+                    run_id=run_id,
+                    event_type="span.started",
+                    span_id=span_id,
+                    parent_span_id=run_id,
+                    kind="agent",
+                    name=author,
+                    hook_event_type="subagent.start",
+                )
+            )
+
+        content = getattr(adk_event, "content", None)
+        parts = getattr(content, "parts", None) or []
+        parent_span = state["subagent_span_ids"].get(state["current_subagent"] or "", run_id)
+        # A model event that issues tool calls is one LLM turn; all its calls share it.
+        if any(getattr(part, "function_call", None) is not None for part in parts):
+            state["turn"] += 1
+
+        for part in parts:
+            fc = getattr(part, "function_call", None)
+            fr = getattr(part, "function_response", None)
+            text = getattr(part, "text", None)
+
+            if fc is not None:
+                call_key = getattr(fc, "id", None) or fc.name
+                tool_span_id = str(uuid.uuid4())
+                state["pending_tool_spans"][call_key] = tool_span_id
+                state["pending_tool_turns"][call_key] = state["turn"]
+                args = dict(fc.args or {})
+                await client.enqueue(
+                    AgentOpsEvent(
+                        run_id=run_id,
+                        event_type="span.started",
+                        span_id=tool_span_id,
+                        parent_span_id=parent_span,
+                        kind="tool",
+                        name=fc.name,
+                        hook_event_type="tool.start",
+                        payload={"args": args},
+                    )
+                )
+                await client.message(
+                    fc.name,
+                    role="tool",
+                    run_id=run_id,
+                    message_type="tool_call",
+                    payload={"input": args, "tool_use_id": call_key, "step": state["turn"]},
+                )
+            elif fr is not None:
+                call_key = getattr(fr, "id", None) or fr.name
+                tool_span_id = state["pending_tool_spans"].pop(call_key, str(uuid.uuid4()))
+                step = state["pending_tool_turns"].pop(call_key, state["turn"])
+                response_val = fr.response or {}
+                result_str = str(response_val.get("result", ""))
+                is_error = result_str.startswith("Error:") or "error" in response_val
+                await client.enqueue(
+                    AgentOpsEvent(
+                        run_id=run_id,
+                        event_type="span.ended",
+                        span_id=tool_span_id,
+                        kind="tool",
+                        name=fr.name,
+                        status="error" if is_error else "ok",
+                        hook_event_type="tool.error" if is_error else "tool.end",
+                        payload={"response": result_str[:4096]},
+                    )
+                )
+                await client.message(
+                    f"{fr.name} returned",
+                    role="tool",
+                    run_id=run_id,
+                    message_type="data",
+                    payload={"output": result_str[:4096], "tool_use_id": call_key, "step": step},
+                )
+                if is_error:
+                    await client.log(
+                        f"Tool {fr.name} error: {result_str[:300]}",
+                        level="ERROR",
+                        run_id=run_id,
+                    )
+            elif text:
+                try:
+                    is_final = adk_event.is_final_response()
+                except AttributeError:
+                    is_final = False
+                if is_final:
+                    await client.enqueue(
+                        AgentOpsEvent(
+                            run_id=run_id,
+                            event_type="message.created",
+                            role="assistant",
+                            content=text[:4096],
+                            hook_event_type="notification",
+                            payload={"agent": author},
+                        )
+                    )
+
+        return a2a_event
+
+    async def after_agent(exec_ctx: Any, final_event: Any) -> Any:
+        session_id = exec_ctx.session_id
+        state = _states.pop(session_id, {})
+        client = client_mod._active_client
+        if client is None:
+            return final_event
+
+        run_id = session_id
+        current_subagent = state.get("current_subagent")
+        subagent_span_ids = state.get("subagent_span_ids", {})
+
+        if current_subagent is not None:
+            await client.enqueue(
+                AgentOpsEvent(
+                    run_id=run_id,
+                    event_type="span.ended",
+                    span_id=subagent_span_ids.get(current_subagent, current_subagent),
+                    parent_span_id=run_id,
+                    kind="agent",
+                    name=current_subagent,
+                    hook_event_type="subagent.end",
+                )
+            )
+        await client.enqueue(
+            AgentOpsEvent(
+                run_id=run_id,
+                event_type="span.ended",
+                span_id=run_id,
+                kind="agent",
+                name=root_agent_name,
+                hook_event_type="session.stop",
+            )
+        )
+        await client.flush()
+        return final_event
+
+    return ExecuteInterceptor(
+        before_agent=before_agent,
+        after_event=after_event,
+        after_agent=after_agent,
+    )
+
+
+async def run_adk_pipeline(
+    pipeline: Any,
+    *,
+    run_id: str,
+    initial_prompt: str,
+    root_agent_name: str,
+    initial_state: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Run an ADK agent pipeline, emitting all AgentOps hook events.
+
+    Returns the final session state dict (output_key values from each sub-agent).
+    If no AgentOpsClient is active, events are silently dropped.
+    """
+    app_name = f"adk-{run_id}"
+    user_id = "worker"
+
+    # Emit an `llm` span around each model call (the run stream only sees completed turns).
+    instrument_agent_model_calls(pipeline, run_id=run_id, parent_span_id=run_id)
+
+    session_service = InMemorySessionService()
+    runner = Runner(agent=pipeline, app_name=app_name, session_service=session_service)
+    session = await session_service.create_session(
+        app_name=app_name,
+        user_id=user_id,
+        state=initial_state or {},
+    )
+    message = types.Content(
+        role="user",
+        parts=[types.Part(text=initial_prompt)],
+    )
+
+    async for _ in instrument_adk_run(
+        runner.run_async(user_id=user_id, session_id=session.id, new_message=message),
+        run_id=run_id,
+        root_agent_name=root_agent_name,
+        initial_prompt=initial_prompt,
+        default_usage=usage_from_adk_model(pipeline),
+    ):
+        pass
+
+    refreshed = await session_service.get_session(
+        app_name=app_name,
+        user_id=user_id,
+        session_id=session.id,
+    )
+    return refreshed.state if refreshed else {}
+
+
+__all__ = ["make_agentops_adk_interceptor", "run_adk_pipeline"]