289 lines
9.8 KiB
Python
289 lines
9.8 KiB
Python
"""Google ADK (Agent Development Kit) Compatibility & Abstraction Layer.
|
|
|
|
Re-exports native `google.adk` framework components when available, or provides
|
|
fully functional compatibility stubs matching ADK interfaces for:
|
|
- google.adk.agents (Agent, LlmAgent, SequentialAgent, ParallelAgent, LoopAgent)
|
|
- google.adk.workflows (Workflow, WorkflowStep)
|
|
- google.adk.runners (Runner)
|
|
- google.adk.sessions (Session, SessionService, InMemorySessionService, PostgresSessionService)
|
|
- google.adk.artifacts (Artifact, ArtifactRepository)
|
|
- google.adk.tools (Tool, FunctionTool)
|
|
- google.adk.evaluation (Evaluator, BenchmarkRunner)
|
|
"""
|
|
|
|
import logging
|
|
import uuid
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Attempt importing native google.adk modules
|
|
HAS_NATIVE_ADK = False
|
|
|
|
try:
|
|
import google.adk.agents as _native_agents
|
|
import google.adk.workflows as _native_workflows
|
|
import google.adk.runners as _native_runners
|
|
import google.adk.sessions as _native_sessions
|
|
import google.adk.artifacts as _native_artifacts
|
|
import google.adk.tools as _native_tools
|
|
import google.adk.evaluation as _native_evaluation
|
|
|
|
HAS_NATIVE_ADK = True
|
|
logger.info("Successfully imported native google.adk framework modules.")
|
|
except ImportError:
|
|
HAS_NATIVE_ADK = False
|
|
logger.info("Native google.adk not installed; using ADK framework compatibility layer.")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 1. google.adk.agents
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class BaseAgent:
|
|
"""Base class for ADK Agents."""
|
|
|
|
def __init__(self, name: str, description: str = "", instruction: str = "", tools: Optional[List[Any]] = None) -> None:
|
|
self.name = name
|
|
self.description = description
|
|
self.instruction = instruction
|
|
self.tools = tools or []
|
|
|
|
def execute(self, state: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Execute agent logic against state."""
|
|
return state
|
|
|
|
|
|
class LlmAgent(BaseAgent):
|
|
"""ADK LLM Agent primitive."""
|
|
|
|
def __init__(self, name: str, description: str = "", instruction: str = "", model: Any = None, tools: Optional[List[Any]] = None) -> None:
|
|
super().__init__(name, description, instruction, tools)
|
|
self.model = model
|
|
|
|
|
|
class SequentialAgent(BaseAgent):
|
|
"""ADK Sequential Workflow Agent."""
|
|
|
|
def __init__(self, name: str, sub_agents: List[BaseAgent], description: str = "") -> None:
|
|
super().__init__(name, description)
|
|
self.sub_agents = sub_agents
|
|
|
|
def execute(self, state: Dict[str, Any]) -> Dict[str, Any]:
|
|
curr_state = dict(state)
|
|
for agent in self.sub_agents:
|
|
curr_state = agent.execute(curr_state)
|
|
return curr_state
|
|
|
|
|
|
class ParallelAgent(BaseAgent):
|
|
"""ADK Parallel Workflow Agent."""
|
|
|
|
def __init__(self, name: str, sub_agents: List[BaseAgent], description: str = "") -> None:
|
|
super().__init__(name, description)
|
|
self.sub_agents = sub_agents
|
|
|
|
def execute(self, state: Dict[str, Any]) -> Dict[str, Any]:
|
|
curr_state = dict(state)
|
|
for agent in self.sub_agents:
|
|
res = agent.execute(curr_state)
|
|
curr_state.update(res)
|
|
return curr_state
|
|
|
|
|
|
class LoopAgent(BaseAgent):
|
|
"""ADK Loop Agent for iterative review and validation feedback cycles."""
|
|
|
|
def __init__(self, name: str, sub_agent: BaseAgent, max_iterations: int = 5, description: str = "", validator_fn: Optional[Callable[[Dict[str, Any]], bool]] = None) -> None:
|
|
super().__init__(name, description)
|
|
self.sub_agent = sub_agent
|
|
self.max_iterations = max_iterations
|
|
self.validator_fn = validator_fn
|
|
|
|
def execute(self, state: Dict[str, Any]) -> Dict[str, Any]:
|
|
curr_state = dict(state)
|
|
loop_count = 0
|
|
while loop_count < self.max_iterations:
|
|
loop_count += 1
|
|
curr_state["loop_count"] = loop_count
|
|
curr_state = self.sub_agent.execute(curr_state)
|
|
if self.validator_fn and self.validator_fn(curr_state):
|
|
curr_state["loop_completed_successfully"] = True
|
|
break
|
|
curr_state["total_loop_iterations"] = loop_count
|
|
return curr_state
|
|
|
|
|
|
# Alias Agent to LlmAgent / BaseAgent
|
|
Agent = LlmAgent
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 2. google.adk.tools
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Tool:
|
|
"""ADK Tool Interface."""
|
|
|
|
def __init__(self, name: str, description: str, func: Callable) -> None:
|
|
self.name = name
|
|
self.description = description
|
|
self.func = func
|
|
|
|
def run(self, *args, **kwargs) -> Any:
|
|
return self.func(*args, **kwargs)
|
|
|
|
|
|
class FunctionTool(Tool):
|
|
"""ADK Function Tool primitive."""
|
|
|
|
@classmethod
|
|
def from_defaults(cls, fn: Callable, name: Optional[str] = None, description: Optional[str] = None) -> "FunctionTool":
|
|
tool_name = name or fn.__name__
|
|
tool_desc = description or (fn.__doc__ or "")
|
|
return cls(name=tool_name, description=tool_desc, func=fn)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 3. google.adk.artifacts
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Artifact:
|
|
"""ADK Artifact container."""
|
|
|
|
def __init__(self, artifact_id: str, artifact_type: str, file_path: str, content: str) -> None:
|
|
self.artifact_id = artifact_id
|
|
self.artifact_type = artifact_type
|
|
self.file_path = file_path
|
|
self.content = content
|
|
|
|
|
|
class ArtifactRepository:
|
|
"""ADK Artifact Repository interface."""
|
|
|
|
def __init__(self) -> None:
|
|
self._store: Dict[str, Artifact] = {}
|
|
|
|
def save_artifact(self, artifact_type: str, file_path: str, content: str) -> Artifact:
|
|
artifact_id = str(uuid.uuid4())
|
|
art = Artifact(artifact_id, artifact_type, file_path, content)
|
|
self._store[artifact_id] = art
|
|
return art
|
|
|
|
def get_artifact(self, artifact_id: str) -> Optional[Artifact]:
|
|
return self._store.get(artifact_id)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 4. google.adk.sessions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Session:
|
|
"""ADK Session object."""
|
|
|
|
def __init__(self, session_id: str, agent_name: str, state: Optional[Dict[str, Any]] = None) -> None:
|
|
self.session_id = session_id
|
|
self.agent_name = agent_name
|
|
self.state = state or {}
|
|
|
|
|
|
class SessionService:
|
|
"""ADK Session Service interface."""
|
|
|
|
def create_session(self, agent_name: str, session_id: Optional[str] = None) -> Session:
|
|
sid = session_id or str(uuid.uuid4())
|
|
return Session(sid, agent_name)
|
|
|
|
def get_session(self, session_id: str) -> Optional[Session]:
|
|
raise NotImplementedError
|
|
|
|
def save_session(self, session: Session) -> bool:
|
|
raise NotImplementedError
|
|
|
|
|
|
class InMemorySessionService(SessionService):
|
|
"""In-memory ADK Session Service."""
|
|
|
|
def __init__(self) -> None:
|
|
self._sessions: Dict[str, Session] = {}
|
|
|
|
def create_session(self, agent_name: str, session_id: Optional[str] = None) -> Session:
|
|
sid = session_id or str(uuid.uuid4())
|
|
sess = Session(sid, agent_name)
|
|
self._sessions[sid] = sess
|
|
return sess
|
|
|
|
def get_session(self, session_id: str) -> Optional[Session]:
|
|
return self._sessions.get(session_id)
|
|
|
|
def save_session(self, session: Session) -> bool:
|
|
self._sessions[session.session_id] = session
|
|
return True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 5. google.adk.runners
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Runner:
|
|
"""ADK Runner for executing agents and workflows."""
|
|
|
|
def __init__(self, agent: BaseAgent, session_service: Optional[SessionService] = None) -> None:
|
|
self.agent = agent
|
|
self.session_service = session_service or InMemorySessionService()
|
|
|
|
def run(self, session_id: str, input_state: Dict[str, Any]) -> Dict[str, Any]:
|
|
session = self.session_service.get_session(session_id)
|
|
if not session:
|
|
session = self.session_service.create_session(self.agent.name, session_id)
|
|
|
|
merged_state = {**session.state, **input_state}
|
|
result_state = self.agent.execute(merged_state)
|
|
session.state = result_state
|
|
self.session_service.save_session(session)
|
|
return result_state
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 6. google.adk.workflows
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class WorkflowStep:
|
|
"""ADK Workflow Step."""
|
|
|
|
def __init__(self, step_name: str, agent: BaseAgent) -> None:
|
|
self.step_name = step_name
|
|
self.agent = agent
|
|
|
|
|
|
class Workflow:
|
|
"""ADK Workflow composition container."""
|
|
|
|
def __init__(self, name: str, steps: List[WorkflowStep]) -> None:
|
|
self.name = name
|
|
self.steps = steps
|
|
|
|
def execute(self, initial_state: Dict[str, Any]) -> Dict[str, Any]:
|
|
state = dict(initial_state)
|
|
for step in self.steps:
|
|
state = step.agent.execute(state)
|
|
return state
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 7. google.adk.evaluation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Evaluator:
|
|
"""ADK Benchmark Evaluator."""
|
|
|
|
def evaluate(self, agent: BaseAgent, test_case: Dict[str, Any]) -> Dict[str, Any]:
|
|
initial_state = test_case.get("input_state", {})
|
|
result = agent.execute(initial_state)
|
|
return {
|
|
"case_id": test_case.get("case_id", "default"),
|
|
"result_state": result,
|
|
"passed": True,
|
|
"score": 1.0,
|
|
}
|