from __future__ import annotations import logging from dataclasses import dataclass, field from typing import Any from google.adk.runners import InMemoryRunner from google.genai import types from agent.agent import root_agent from core.config import get_settings logger = logging.getLogger(__name__) @dataclass class SessionState: runner: InMemoryRunner @dataclass class TurnResult: text: str = "" tool_calls: list[dict[str, Any]] = field(default_factory=list) class ADKRuntime: def __init__(self) -> None: self.settings = get_settings() self._sessions: dict[tuple[str, str], SessionState] = {} async def run_turn(self, message: str, user_id: str, session_id: str) -> str: result = await self.run_turn_detailed( message=message, user_id=user_id, session_id=session_id, ) return result.text async def run_turn_detailed( self, message: str, user_id: str, session_id: str, ) -> TurnResult: runner = await self._get_or_create_runner(user_id=user_id, session_id=session_id) user_content = types.Content( role="user", parts=[types.Part.from_text(text=message)], ) latest_text = "" tool_calls: list[dict[str, Any]] = [] async for event in runner.run_async( user_id=user_id, session_id=session_id, new_message=user_content, ): extracted = self._extract_text(event) if extracted.strip(): latest_text = extracted new_tool_responses = self._extract_tool_responses(event) if new_tool_responses: tool_calls.extend(new_tool_responses) return TurnResult(text=latest_text.strip(), tool_calls=tool_calls) async def _get_or_create_runner(self, user_id: str, session_id: str) -> InMemoryRunner: key = (user_id, session_id) existing = self._sessions.get(key) if existing is not None: return existing.runner runner = InMemoryRunner(agent=root_agent, app_name=self.settings.app_name) await runner.session_service.create_session( app_name=self.settings.app_name, user_id=user_id, session_id=session_id, ) self._sessions[key] = SessionState(runner=runner) return runner @staticmethod def _extract_text(event: Any) -> str: content = getattr(event, "content", None) if content is None: return "" parts = getattr(content, "parts", None) if not parts: return "" chunks: list[str] = [] for part in parts: text = getattr(part, "text", None) if isinstance(text, str) and text: chunks.append(text) return "\n".join(chunks) @staticmethod def _debug_dump_event(event: Any) -> None: try: import json from pathlib import Path payload: dict[str, Any] = {"type": type(event).__name__} content = getattr(event, "content", None) if content is not None: payload["content_role"] = getattr(content, "role", None) parts_info = [] for part in getattr(content, "parts", None) or []: fr = getattr(part, "function_response", None) info: dict[str, Any] = { "has_text": bool(getattr(part, "text", None)), "has_function_call": bool(getattr(part, "function_call", None)), "has_function_response": bool(fr), } if fr is not None: info["fr_name"] = getattr(fr, "name", None) resp = getattr(fr, "response", None) info["fr_response_type"] = type(resp).__name__ info["fr_response_repr"] = repr(resp)[:1500] if isinstance(resp, dict): info["fr_response_keys"] = list(resp.keys()) parts_info.append(info) payload["parts"] = parts_info log_path = Path(__file__).resolve().parents[1] / "logs" / "adk-events.log" log_path.parent.mkdir(parents=True, exist_ok=True) with log_path.open("a", encoding="utf-8") as fh: fh.write(json.dumps(payload, default=str) + "\n") except Exception as exc: try: from pathlib import Path log_path = Path(__file__).resolve().parents[1] / "logs" / "adk-events.log" with log_path.open("a", encoding="utf-8") as fh: fh.write(f"DUMP-ERR: {exc}\n") except Exception: pass @staticmethod def _extract_tool_responses(event: Any) -> list[dict[str, Any]]: """Pull MCP tool responses out of an ADK event so we can render their UI payloads.""" content = getattr(event, "content", None) if content is None: return [] parts = getattr(content, "parts", None) or [] responses: list[dict[str, Any]] = [] for part in parts: function_response = getattr(part, "function_response", None) if function_response is None: continue name = getattr(function_response, "name", None) or "" response = getattr(function_response, "response", None) payload = ADKRuntime._coerce_tool_payload(response) if payload is None: continue responses.append({"name": name, "response": payload}) # ADK sometimes surfaces tool results via convenience attributes on the # event itself (e.g. event.actions.state_delta or event.tool_responses). # Dig those up too so we don't miss the UI payload. for attr in ("tool_responses", "function_responses"): extra = getattr(event, attr, None) if not extra: continue if isinstance(extra, list): for item in extra: name = getattr(item, "name", None) or "" response = getattr(item, "response", None) payload = ADKRuntime._coerce_tool_payload(response) if payload is not None: responses.append({"name": name, "response": payload}) return responses @staticmethod def _coerce_tool_payload(value: Any) -> dict[str, Any] | None: """Unwrap a tool response into a plain dict. ADK + MCP delivers tool results as one of several shapes: 1. ``{"result": }`` — direct ADK function tool 2. ``{"result": CallToolResult(...)}``— MCP toolset wraps the response into an ``mcp.types.CallToolResult`` whose ``content`` is a list of ``TextContent`` items carrying JSON-encoded text. 3. ``CallToolResult(...)`` — same as above, unwrapped. 4. plain dict already. """ import json # Step 1: pull off the optional "result" wrapper. if isinstance(value, dict) and "result" in value and len(value) == 1: value = value["result"] # Step 2: plain dict — done. if isinstance(value, dict): return value # Step 3: MCP CallToolResult / TextContent. Detect by duck-typing the # ``content`` attribute that holds a list of items with ``.text``. content_items = getattr(value, "content", None) if isinstance(content_items, list) and content_items: for item in content_items: text = getattr(item, "text", None) if isinstance(text, str) and text.strip(): try: parsed = json.loads(text) if isinstance(parsed, dict): return parsed except json.JSONDecodeError: continue # Step 4: protobuf Struct fallback. to_dict = getattr(value, "to_dict", None) if callable(to_dict): try: converted = to_dict() if isinstance(converted, dict): return converted.get("result", converted) except Exception: return None return None