from __future__ import annotations import json from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Dict, List, Optional, Type, Any from openspace.tool_runtime.orchestration import run_tools from openspace.utils.logging import Logger if TYPE_CHECKING: from openspace.llm import LLMClient from openspace.grounding.core.grounding_client import GroundingClient from openspace.recording import RecordingManager logger = Logger.get_logger(__name__) class BaseAgent(ABC): def __init__( self, name: str, backend_scope: Optional[List[str]] = None, llm_client: Optional[LLMClient] = None, grounding_client: Optional[GroundingClient] = None, recording_manager: Optional[RecordingManager] = None, ) -> None: """ Initialize the BaseAgent. Args: name: Unique name for the agent backend_scope: List of backend types this agent can access (e.g., ["gui", "shell", "mcp", "web", "meta"]) llm_client: LLM client for agent reasoning (optional, can be set later) grounding_client: Reference to GroundingClient for tool execution recording_manager: RecordingManager for recording execution """ self._name = name self._grounding_client: Optional[GroundingClient] = grounding_client self._backend_scope = backend_scope or [] self._llm_client = llm_client self._recording_manager: Optional[RecordingManager] = recording_manager self._step = 0 self._status = AgentStatus.ACTIVE self._register_self() logger.info(f"Initialized {self.__class__.__name__}: {name}") @property def name(self) -> str: return self._name @property def grounding_client(self) -> Optional[GroundingClient]: """Get the grounding client.""" return self._grounding_client @property def backend_scope(self) -> List[str]: return self._backend_scope @property def llm_client(self) -> Optional[LLMClient]: return self._llm_client @llm_client.setter def llm_client(self, client: LLMClient) -> None: self._llm_client = client @property def recording_manager(self) -> Optional[RecordingManager]: """Get the recording manager.""" return self._recording_manager @property def step(self) -> int: return self._step @property def status(self) -> str: return self._status @abstractmethod async def process(self, context: Dict[str, Any]) -> Dict[str, Any]: pass @abstractmethod def construct_messages(self, context: Dict[str, Any]) -> List[Dict[str, Any]]: """ Construct messages for LLM reasoning. Context must contain 'instruction' key. """ pass async def get_llm_response( self, messages: List[Dict[str, Any]], tools: Optional[List] = None, **kwargs ) -> Dict[str, Any]: if not self._llm_client: raise ValueError(f"LLM client not initialized for agent {self.name}") try: current_messages = list(messages) # Use the fallback-aware variant so that consecutive-overloaded # errors don't hard-fail auxiliary agents — matches the # ``FallbackTriggeredError`` handling in ``GroundingAgent.process``. model_response = await self._llm_client.call_model_with_fallback( messages=messages, tools=tools, **kwargs ) effective_model = ( model_response.effective_model or kwargs.get("model") or self._llm_client.model ) assistant_message = model_response.assistant_message current_messages.append(assistant_message) current_messages.extend( self._llm_client.get_model_response_followup_messages(model_response) ) tool_results: list[dict[str, Any]] = [] if model_response.tool_calls and tools: tool_context = self._llm_client.build_auxiliary_tool_use_context( tools=tools, messages=current_messages, model=effective_model, agent_id=self.name, agent_type=self.__class__.__name__, hook_registry=getattr(self, "_hook_registry", None), event_sink=getattr(self, "_runtime_event_sink", None), recording_manager=self._recording_manager, quality_manager=getattr(self._grounding_client, "quality_manager", None), task_description=str(current_messages[0].get("content", "") or ""), ) tools_result = await run_tools( model_response.tool_calls, model_response.tool_map, tool_context, assistant_message=assistant_message, ) current_messages.extend(tools_result.messages) tool_results = self._llm_client.collect_tool_results( model_response.tool_calls, model_response.tool_map, tools_result.messages, ) return { "message": assistant_message, "tool_results": tool_results, "messages": current_messages, "has_tool_calls": bool(model_response.tool_calls), } except Exception as e: logger.error(f"{self.name}: LLM call failed: {e}", exc_info=True) raise def response_to_dict(self, response: str) -> Dict[str, Any]: try: if response.strip().startswith("```json") or response.strip().startswith("```"): lines = response.strip().split('\n') if lines and lines[0].startswith('```'): lines = lines[1:] end_idx = len(lines) for i, line in enumerate(lines): if line.strip() == '```': end_idx = i break response = '\n'.join(lines[:end_idx]) return json.loads(response) except json.JSONDecodeError as e: # If parsing fails, try to find and extract just the JSON object/array if "Extra data" in str(e): try: decoder = json.JSONDecoder() obj, idx = decoder.raw_decode(response) logger.warning( f"{self.name}: Successfully extracted JSON but found extra text after position {idx}. " f"Extra text: {response[idx:idx+100]}..." ) return obj except Exception as e2: logger.error(f"{self.name}: Failed to extract JSON even with raw_decode: {e2}") logger.error(f"{self.name}: Failed to parse response: {e}") logger.error(f"{self.name}: Response content: {response[:500]}") return {"error": "Failed to parse response", "raw": response} def increment_step(self) -> None: self._step += 1 @classmethod def _register_self(cls) -> None: """Register the agent class in the registry upon instantiation.""" # Get the actual instance class, not BaseAgent if cls.__name__ != "BaseAgent" and cls.__name__ not in AgentRegistry._registry: AgentRegistry.register(cls.__name__, cls) def __repr__(self) -> str: return f"<{self.__class__.__name__}(name={self.name}, step={self.step}, status={self.status})>" class AgentStatus: """Constants for agent status.""" ACTIVE = "active" IDLE = "idle" WAITING = "waiting" class AgentRegistry: """ Registry for managing agent classes. Allows dynamic registration and retrieval of agent types. """ _registry: Dict[str, Type[BaseAgent]] = {} @classmethod def register(cls, name: str, agent_cls: Type[BaseAgent]) -> None: if name in cls._registry: logger.warning(f"Agent class '{name}' already registered, overwriting") cls._registry[name] = agent_cls logger.debug(f"Registered agent class: {name}") @classmethod def get_cls(cls, name: str) -> Type[BaseAgent]: if name not in cls._registry: raise ValueError(f"No agent class registered under '{name}'") return cls._registry[name] @classmethod def list_registered(cls) -> List[str]: return list(cls._registry.keys()) @classmethod def clear(cls) -> None: cls._registry.clear() logger.debug("Agent registry cleared")