codekingpro/portable-devtools

1"""LangSmith integration for Google ADK (Agent Development Kit)."""2 3from __future__ import annotations4 5import logging6from typing import Optional7 8from ._config import set_tracing_config9 10logger = logging.getLogger(__name__)11 12__all__ = ["configure_google_adk", "create_traced_session_context"]13 14_patched = False15 16 17def configure_google_adk(18    name: Optional[str] = None,19    project_name: Optional[str] = None,20    metadata: Optional[dict] = None,21    tags: Optional[list[str]] = None,22) -> bool:23    """Enable LangSmith tracing for Google ADK.24 25    Can be called before or after importing Runner (import-order agnostic).26 27    Args:28        name: Name of the root trace. Defaults to "google_adk.session".29        project_name: LangSmith project to trace to.30        metadata: Metadata to associate with all traces.31        tags: Tags to associate with all traces.32 33    Returns:34        True if configuration was successful, False otherwise.35    """36    global _patched37 38    if _patched:39        set_tracing_config(40            name=name, project_name=project_name, metadata=metadata, tags=tags41        )42        return True43 44    try:45        import google.adk  # noqa: F40146        from wrapt import wrap_function_wrapper  # type: ignore[import-untyped]47    except ImportError as e:48        logger.warning(f"Missing dependency: {e}")49        return False50 51    set_tracing_config(52        name=name, project_name=project_name, metadata=metadata, tags=tags53    )54 55    from ._client import (56        wrap_agent_run_async,57        wrap_flow_call_llm_async,58        wrap_runner_run,59        wrap_runner_run_async,60        wrap_tool_run_async,61    )62 63    _wraps = [64        (65            "google.adk.runners",66            "Runner.run",67            wrap_runner_run,68        ),69        (70            "google.adk.runners",71            "Runner.run_async",72            wrap_runner_run_async,73        ),74        (75            "google.adk.agents.base_agent",76            "BaseAgent.run_async",77            wrap_agent_run_async,78        ),79        (80            "google.adk.flows.llm_flows.base_llm_flow",81            "BaseLlmFlow._call_llm_async",82            wrap_flow_call_llm_async,83        ),84        (85            "google.adk.tools.base_tool",86            "BaseTool.run_async",87            wrap_tool_run_async,88        ),89        (90            "google.adk.tools.function_tool",91            "FunctionTool.run_async",92            wrap_tool_run_async,93        ),94        (95            "google.adk.tools.mcp_tool.mcp_tool",96            "McpTool.run_async",97            wrap_tool_run_async,98        ),99    ]100 101    for module, name, wrapper in _wraps:102        try:103            wrap_function_wrapper(module, name, wrapper)104        except Exception as e:105            logger.warning(f"Failed to wrap {name}: {e}")106 107    _patched = True108    return True109 110 111def create_traced_session_context(112    name: Optional[str] = None,113    project_name: Optional[str] = None,114    metadata: Optional[dict] = None,115    tags: Optional[list[str]] = None,116    inputs: Optional[dict] = None,117):118    """Create a trace context for manual session tracing."""119    from ._client import create_traced_session_context as _create_context120 121    return _create_context(122        name=name,123        project_name=project_name,124        metadata=metadata,125        tags=tags,126        inputs=inputs,127    )128