Spaces:
Sleeping
Sleeping
| 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__) | |
| class SessionState: | |
| runner: InMemoryRunner | |
| 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 | |
| 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) | |
| 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 | |
| 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 | |
| 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": <plain dict>}`` — 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 | |