"""Simple, ready-to-use Goose client""" import logging from contextlib import contextmanager from typing import Any, Dict, Iterator, List, Optional, Union import httpx from ._generated._generated.models import ( Session, StartAgentRequest, ) from .core import BaseClient from .streaming import StreamingMixin logger = logging.getLogger(__name__) class GooseClient(BaseClient, StreamingMixin): """Simple Goose client for 80% of use cases""" def __init__(self, api_key: str, base_url: str = "http://localhost:3000"): """Initialize the Goose client. Args: api_key: The API key for authentication base_url: The base URL for the Goose server """ super().__init__(api_key, base_url) self._active_session: Optional[str] = None # --- Session Management --- def create_session(self, working_dir: str = "/tmp") -> Session: """Create a new session. Args: working_dir: Working directory for the session Returns: The created Session object """ request = StartAgentRequest(working_dir=working_dir) session = self._agent_api.start_agent(start_agent_request=request) self._active_session = session.id logger.info(f"Created session {session.id} with working_dir={working_dir}") return session @contextmanager def session(self, working_dir: str = "/tmp"): """Context manager for session lifecycle. Args: working_dir: Working directory for the session Yields: The created Session object """ session = self.create_session(working_dir) old_active = self._active_session self._active_session = session.id logger.debug(f"Entering session context {session.id}") try: yield session finally: self._active_session = old_active try: self._sessions_api.delete_session(session.id) logger.debug(f"Deleted session {session.id}") except Exception as e: logger.warning(f"Failed to delete session {session.id}: {e}") def list_sessions(self) -> List[Session]: """List all sessions. Returns: List of Session objects """ response = self._sessions_api.list_sessions() sessions = response.sessions if hasattr(response, "sessions") else [] logger.debug(f"Listed {len(sessions)} sessions") return sessions def get_session(self, session_id: str) -> Session: """Get a specific session. Args: session_id: The session ID to retrieve Returns: The Session object """ session = self._sessions_api.get_session(session_id) logger.debug(f"Retrieved session {session_id}") return session def delete_session(self, session_id: str) -> bool: """Delete a session. Args: session_id: The session ID to delete Returns: True if successful, False otherwise """ try: self._sessions_api.delete_session(session_id) logger.info(f"Deleted session {session_id}") if self._active_session == session_id: self._active_session = None return True except Exception as e: logger.error(f"Failed to delete session {session_id}: {e}") return False # --- Messaging --- def send_message( self, text: str, session_id: Optional[str] = None, stream: bool = False ) -> Union[str, Iterator[Dict[str, Any]]]: """Send a message to a session. Args: text: The message text to send session_id: Optional session ID (uses active session if not provided) stream: Whether to stream the response Returns: Either the complete response text or an iterator of stream events """ # Use provided session or active session or create temporary if not session_id: if not self._active_session: temp_session = self.create_session() session_id = temp_session.id logger.debug(f"Created temporary session {session_id}") else: session_id = self._active_session # Create message with proper structure message_dict = {"role": "user", "content": [{"type": "text", "text": text}]} logger.debug(f"Sending message to session {session_id}: {text[:50]}...") # Stream or collect response if stream: return self.stream_reply(session_id, [message_dict]) else: # Collect full response response_text = "" for event in self.stream_reply(session_id, [message_dict]): if event.get("type") == "Message": msg = event.get("message", {}) for content in msg.get("content", []): if content.get("type") == "text": response_text += content.get("text", "") elif event.get("type") == "Error": error_msg = event.get("error", "Unknown error") logger.error(f"Error in response: {error_msg}") raise Exception(f"Error: {error_msg}") logger.debug(f"Received response: {response_text[:100]}...") return response_text def chat(self, text: str, session_id: Optional[str] = None) -> str: """Simple chat interface - always returns text. Args: text: The message text to send session_id: Optional session ID (uses active session if not provided) Returns: The response text """ return self.send_message(text, session_id, stream=False) def stream_chat(self, text: str, session_id: Optional[str] = None) -> Iterator[str]: """Stream chat responses as text chunks. Args: text: The message text to send session_id: Optional session ID (uses active session if not provided) Yields: Text chunks from the response """ for event in self.send_message(text, session_id, stream=True): if event.get("type") == "Message": msg = event.get("message", {}) for content in msg.get("content", []): if content.get("type") == "text": text_chunk = content.get("text", "") if text_chunk: yield text_chunk elif event.get("type") == "Error": error_msg = event.get("error", "Unknown error") logger.error(f"Stream error: {error_msg}") break # --- Tools & Extensions --- def list_tools(self, session_id: Optional[str] = None) -> List[Dict[str, Any]]: """List available tools. Args: session_id: Optional session ID (uses active session if not provided) Returns: List of tool information dictionaries """ if not session_id: session_id = self._active_session or self.create_session().id try: tools = self._agent_api.get_tools(session_id=session_id) logger.debug(f"Listed {len(tools)} tools for session {session_id}") return tools except Exception as e: logger.error(f"Failed to list tools: {e}") return [] def list_extensions(self) -> List[Dict[str, Any]]: """List available extensions. Returns: List of extension information dictionaries """ try: response = self._config_api.get_extensions() extensions = response.extensions if hasattr(response, "extensions") else [] logger.debug(f"Listed {len(extensions)} extensions") return extensions except Exception as e: logger.error(f"Failed to list extensions: {e}") return [] # --- Convenience Methods --- def ask(self, question: str) -> str: """One-shot question answering. Args: question: The question to ask Returns: The answer text """ with self.session() as session: return self.chat(question, session.id) @property def is_healthy(self) -> bool: """Check server health. Returns: True if server is healthy, False otherwise """ return self.health_check() @property def active_session_id(self) -> Optional[str]: """Get the active session ID. Returns: The active session ID or None """ return self._active_session def set_active_session(self, session_id: str) -> None: """Set the active session. Args: session_id: The session ID to make active """ self._active_session = session_id logger.debug(f"Set active session to {session_id}") def clear_active_session(self) -> None: """Clear the active session.""" self._active_session = None logger.debug("Cleared active session") # --- Tool Confirmation --- def confirm_permission( self, confirmation_id: str, action: str, session_id: Optional[str] = None, principal_type: str = "Tool", ) -> bool: """Confirm or deny a tool permission request. Args: confirmation_id: The ID of the tool confirmation request action: The action to take ("allow_once", "always_allow", or "deny") session_id: Optional session ID (uses active session if not provided) principal_type: Type of principal (default: "Tool") Returns: True if confirmation was successful, False otherwise """ if not session_id: session_id = self._active_session if not session_id: logger.error("No session ID provided and no active session") return False try: # Send confirmation directly via HTTP since we need the /confirm endpoint with httpx.Client() as client: response = client.post( f"{self.base_url}/confirm", headers={"X-Secret-Key": self.api_key}, json={ "id": confirmation_id, "principalType": principal_type, "action": action, "sessionId": session_id, }, ) response.raise_for_status() logger.info(f"Confirmed permission {confirmation_id} with action {action}") return True except Exception as e: logger.error(f"Failed to confirm permission: {e}") return False def stream_chat_with_confirmations( self, text: str, session_id: Optional[str] = None, auto_confirm: Optional[str] = None ) -> Iterator[Union[str, Dict[str, Any]]]: """Stream chat with tool confirmation support. This method yields both text chunks and tool confirmation requests. When a tool confirmation is received, it yields the confirmation dict so the caller can handle it (e.g., prompt user or auto-confirm). Args: text: The message text to send session_id: Optional session ID (uses active session if not provided) auto_confirm: Optional auto-confirmation action ("allow_once", "always_allow", "deny", or None) If provided, automatically confirms tools with this action Yields: Either text chunks (str) or tool confirmation dicts with structure: { 'type': 'toolConfirmationRequest', 'id': str, 'toolName': str, 'arguments': dict, 'prompt': Optional[str] } """ for event in self.send_message(text, session_id, stream=True): if event.get("type") == "Message": msg = event.get("message", {}) # Check for tool confirmations tool_confirmations = self.get_tool_confirmations(event) for confirmation in tool_confirmations: if auto_confirm: # Auto-confirm with the specified action self.confirm_permission( confirmation["id"], auto_confirm, session_id or self._active_session ) logger.debug( f"Auto-confirmed tool {confirmation['toolName']} with action {auto_confirm}" ) else: # Yield the confirmation for the caller to handle yield confirmation # Yield text content as before for content in msg.get("content", []): if content.get("type") == "text": text_chunk = content.get("text", "") if text_chunk: yield text_chunk elif event.get("type") == "Error": error_msg = event.get("error", "Unknown error") logger.error(f"Stream error: {error_msg}") break