lm-deluge 0.0.67__py3-none-any.whl → 0.0.90__py3-none-any.whl

lm_deluge/tool/prefab/sheets.py

@@ -0,0 +1,385 @@
+"""Google Sheets manipulation prefab tool."""
+
+import json
+import os
+from typing import Any
+
+from lm_deluge.tool import Tool
+
+
+class SheetsManager:
+    """
+    A prefab tool for manipulating Google Sheets.
+
+    Provides tools to read ranges and update individual cells in a Google Sheet.
+    Outputs are formatted as LLM-friendly HTML tables with row/col attributes.
+
+    Args:
+        sheet_id: The ID of the Google Sheet to manipulate
+        credentials_json: Optional. JSON string or dict containing Google service account credentials.
+                         If not provided, will look for GOOGLE_SHEETS_CREDENTIALS env variable.
+        credentials_file: Optional. Path to a JSON file containing credentials.
+                         Only used if credentials_json is not provided.
+        read_tool_name: Name for the read range tool (default: "sheets_read_range")
+        update_tool_name: Name for the update cell tool (default: "sheets_update_cell")
+
+    Example:
+        ```python
+        # Using credentials from environment
+        manager = SheetsManager(sheet_id="your-sheet-id-here")
+
+        # Using credentials directly
+        manager = SheetsManager(
+            sheet_id="your-sheet-id-here",
+            credentials_json={"type": "service_account", ...}
+        )
+
+        # Get tools
+        tools = manager.get_tools()
+        ```
+    """
+
+    def __init__(
+        self,
+        sheet_id: str,
+        *,
+        credentials_json: str | dict[str, Any] | None = None,
+        credentials_file: str | None = None,
+        list_sheets_tool_name: str = "sheets_list_sheets",
+        get_used_range_tool_name: str = "sheets_get_used_range",
+        read_tool_name: str = "sheets_read_range",
+        update_tool_name: str = "sheets_update_cell",
+    ):
+        self.sheet_id = sheet_id
+        self.list_sheets_tool_name = list_sheets_tool_name
+        self.get_used_range_tool_name = get_used_range_tool_name
+        self.read_tool_name = read_tool_name
+        self.update_tool_name = update_tool_name
+
+        # Handle credentials
+        if credentials_json is not None:
+            if isinstance(credentials_json, str):
+                self.credentials = json.loads(credentials_json)
+            else:
+                self.credentials = credentials_json
+        elif credentials_file is not None:
+            with open(credentials_file, "r") as f:
+                self.credentials = json.load(f)
+        else:
+            # Try to load from environment
+            env_creds = os.environ.get("GOOGLE_SHEETS_CREDENTIALS")
+            if env_creds:
+                self.credentials = json.loads(env_creds)
+            else:
+                raise ValueError(
+                    "No credentials provided. Please provide credentials_json, "
+                    "credentials_file, or set GOOGLE_SHEETS_CREDENTIALS environment variable."
+                )
+
+        self._service = None
+        self._tools: list[Tool] | None = None
+
+    def _get_service(self):
+        """Lazily initialize the Google Sheets API service."""
+        if self._service is not None:
+            return self._service
+
+        try:
+            from google.oauth2 import service_account
+            from googleapiclient.discovery import build
+        except ImportError:
+            raise ImportError(
+                "Google Sheets API dependencies not installed. "
+                "Please install with: pip install google-api-python-client google-auth"
+            )
+
+        # Create credentials from service account info
+        creds = service_account.Credentials.from_service_account_info(
+            self.credentials, scopes=["https://www.googleapis.com/auth/spreadsheets"]
+        )
+
+        # Build the service
+        self._service = build("sheets", "v4", credentials=creds)
+        return self._service
+
+    def _list_sheets(self) -> str:
+        """
+        List all sheets (tabs) in the spreadsheet.
+
+        Returns:
+            JSON string with status and list of sheet names
+        """
+        try:
+            service = self._get_service()
+            spreadsheet = (
+                service.spreadsheets()
+                .get(spreadsheetId=self.sheet_id, fields="sheets.properties")
+                .execute()
+            )
+
+            sheets = []
+            for sheet in spreadsheet.get("sheets", []):
+                props = sheet.get("properties", {})
+                sheets.append(
+                    {
+                        "name": props.get("title", ""),
+                        "index": props.get("index", 0),
+                        "sheetId": props.get("sheetId", 0),
+                    }
+                )
+
+            return json.dumps({"status": "success", "sheets": sheets})
+
+        except Exception as e:
+            return json.dumps({"status": "error", "error": str(e)})
+
+    def _get_used_range(self, sheet_name: str | None = None) -> str:
+        """
+        Get the used range of a sheet (the bounding box of all non-empty cells).
+
+        Args:
+            sheet_name: Name of the sheet to check. If None, uses the first sheet.
+
+        Returns:
+            JSON string with status and the used range in A1 notation
+        """
+        try:
+            service = self._get_service()
+
+            # If no sheet name provided, get the first sheet's name
+            if not sheet_name:
+                spreadsheet = (
+                    service.spreadsheets()
+                    .get(spreadsheetId=self.sheet_id, fields="sheets.properties.title")
+                    .execute()
+                )
+                sheets = spreadsheet.get("sheets", [])
+                if not sheets:
+                    return json.dumps(
+                        {"status": "error", "error": "No sheets found in spreadsheet"}
+                    )
+                sheet_name = sheets[0]["properties"]["title"]
+
+            # Get all values from the sheet to determine the used range
+            result = (
+                service.spreadsheets()
+                .values()
+                .get(spreadsheetId=self.sheet_id, range=f"'{sheet_name}'")
+                .execute()
+            )
+
+            values = result.get("values", [])
+
+            if not values:
+                return json.dumps(
+                    {
+                        "status": "success",
+                        "sheet_name": sheet_name,
+                        "used_range": None,
+                        "message": "Sheet is empty",
+                    }
+                )
+
+            # Find the max column across all rows
+            max_col = 0
+            for row in values:
+                if row:  # Skip empty rows
+                    max_col = max(max_col, len(row))
+
+            num_rows = len(values)
+            end_col = self._col_num_to_letter(max_col)
+            used_range = f"A1:{end_col}{num_rows}"
+
+            return json.dumps(
+                {
+                    "status": "success",
+                    "sheet_name": sheet_name,
+                    "used_range": f"'{sheet_name}'!{used_range}",
+                    "rows": num_rows,
+                    "cols": max_col,
+                }
+            )
+
+        except Exception as e:
+            return json.dumps({"status": "error", "error": str(e)})
+
+    def _read_range(self, range_spec: str) -> str:
+        """
+        Read a range from the Google Sheet and return as HTML table.
+
+        Args:
+            range_spec: A1 notation range (e.g., "Sheet1!A1:C10" or just "A1:C10")
+
+        Returns:
+            JSON string with status and HTML table data
+        """
+        try:
+            service = self._get_service()
+            sheet = service.spreadsheets()
+
+            result = (
+                sheet.values()
+                .get(spreadsheetId=self.sheet_id, range=range_spec)
+                .execute()
+            )
+
+            values = result.get("values", [])
+
+            if not values:
+                return json.dumps(
+                    {
+                        "status": "success",
+                        "message": "No data found in range",
+                        "html": "<p>No data found</p>",
+                    }
+                )
+
+            # Convert to HTML table with cell attribute for reference
+            html_parts = ["<table>"]
+
+            for row_idx, row in enumerate(values, start=1):
+                html_parts.append("<tr>")
+                for col_idx, cell_value in enumerate(row, start=1):
+                    # Convert column index to letter (1=A, 2=B, etc.)
+                    col_letter = self._col_num_to_letter(col_idx)
+                    cell_ref = f"{col_letter}{row_idx}"
+
+                    html_parts.append(f'<td cell="{cell_ref}">{cell_value}</td>')
+                html_parts.append("</tr>")
+
+            html_parts.append("</table>")
+            html_table = "\n".join(html_parts)
+
+            return json.dumps(
+                {"status": "success", "rows": len(values), "html": html_table}
+            )
+
+        except Exception as e:
+            return json.dumps({"status": "error", "error": str(e)})
+
+    def _update_cell(self, cell: str, value: str) -> str:
+        """
+        Update a single cell in the Google Sheet.
+
+        Args:
+            cell: Cell reference in A1 notation (e.g., "A1", "B5", "Sheet1!C3")
+            value: The value to set in the cell
+
+        Returns:
+            JSON string with status and result
+        """
+        try:
+            service = self._get_service()
+            sheet = service.spreadsheets()
+
+            body = {"values": [[value]]}
+
+            result = (
+                sheet.values()
+                .update(
+                    spreadsheetId=self.sheet_id,
+                    range=cell,
+                    valueInputOption="USER_ENTERED",  # Parse values like formulas, numbers, dates
+                    body=body,
+                )
+                .execute()
+            )
+
+            return json.dumps(
+                {
+                    "status": "success",
+                    "updated_cells": result.get("updatedCells", 0),
+                    "updated_range": result.get("updatedRange", ""),
+                    "message": f"Successfully updated {cell} to '{value}'",
+                }
+            )
+
+        except Exception as e:
+            return json.dumps({"status": "error", "error": str(e)})
+
+    @staticmethod
+    def _col_num_to_letter(n: int) -> str:
+        """Convert column number to letter (1=A, 2=B, ..., 26=Z, 27=AA, etc.)."""
+        result = ""
+        while n > 0:
+            n -= 1
+            result = chr(65 + (n % 26)) + result
+            n //= 26
+        return result
+
+    def get_tools(self) -> list[Tool]:
+        """Return the list of Google Sheets tools."""
+        if self._tools is not None:
+            return self._tools
+
+        self._tools = [
+            Tool(
+                name=self.list_sheets_tool_name,
+                description=(
+                    "List all sheets (tabs) in the spreadsheet. Returns the name, index, "
+                    "and ID of each sheet."
+                ),
+                run=self._list_sheets,
+                parameters={},
+                required=[],
+            ),
+            Tool(
+                name=self.get_used_range_tool_name,
+                description=(
+                    "Get the used range of a sheet (the bounding box containing all non-empty cells). "
+                    "Returns the range in A1 notation (e.g., 'Sheet1'!A1:C4)."
+                ),
+                run=self._get_used_range,
+                parameters={
+                    "sheet_name": {
+                        "type": "string",
+                        "description": (
+                            "Name of the sheet to check. If not provided, uses the first sheet."
+                        ),
+                    }
+                },
+                required=[],
+            ),
+            Tool(
+                name=self.read_tool_name,
+                description=(
+                    "Read a range of cells from the Google Sheet and return as an HTML table. "
+                    "Each cell has a 'cell' attribute with its A1 reference (e.g., cell='A1'). "
+                    "Use A1 notation for the range (e.g., 'A1:C10' or 'Sheet1!A1:C10')."
+                ),
+                run=self._read_range,
+                parameters={
+                    "range_spec": {
+                        "type": "string",
+                        "description": (
+                            "The range to read in A1 notation. Examples: 'A1:C10', 'Sheet1!A1:C10', "
+                            "'A:A' (entire column A), '1:1' (entire row 1)"
+                        ),
+                    }
+                },
+                required=["range_spec"],
+            ),
+            Tool(
+                name=self.update_tool_name,
+                description=(
+                    "Update a single cell in the Google Sheet. The value will be parsed "
+                    "automatically (formulas, numbers, dates, etc.)."
+                ),
+                run=self._update_cell,
+                parameters={
+                    "cell": {
+                        "type": "string",
+                        "description": (
+                            "The cell to update in A1 notation. Examples: 'A1', 'B5', 'Sheet1!C3'"
+                        ),
+                    },
+                    "value": {
+                        "type": "string",
+                        "description": "The value to set in the cell",
+                    },
+                },
+                required=["cell", "value"],
+            ),
+        ]
+
+        return self._tools

lm_deluge/tool/prefab/subagents.py

@@ -0,0 +1,233 @@
+from lm_deluge.api_requests.base import APIResponse
+from lm_deluge.client import AgentLoopResponse, _LLMClient
+from lm_deluge.prompt import Conversation, prompts_to_conversations
+from lm_deluge.tool import Tool
+
+
+class SubAgentManager:
+    """Manages subagent tasks that can be spawned by a main LLM via tool calls.
+
+    The SubAgentManager exposes tools that allow a main LLM to delegate subtasks
+    to specialized or cheaper subagent models, saving context and improving efficiency.
+
+    Example:
+        >>> manager = SubAgentManager(
+        ...     client=LLMClient("gpt-4o-mini"),  # Subagent model
+        ...     tools=[search_tool, calculator_tool]  # Tools available to subagents
+        ... )
+        >>> main_client = LLMClient("gpt-4o")  # More expensive main model
+        >>> conv = Conversation.user("Research AI and calculate market size")
+        >>> # Main model can now call manager tools to spawn subagents
+        >>> conv, resp = await main_client.run_agent_loop(
+        ...     conv,
+        ...     tools=manager.get_tools()
+        ... )
+    """
+
+    def __init__(
+        self,
+        client: _LLMClient,
+        tools: list[Tool] | None = None,
+        max_rounds: int = 5,
+    ):
+        """Initialize the SubAgentManager.
+
+        Args:
+            client: LLMClient to use for subagent tasks
+            tools: Tools available to subagents (optional)
+            max_rounds: Maximum rounds for each subagent's agent loop
+        """
+        self.client = client
+        self.tools = tools or []
+        self.max_rounds = max_rounds
+        self.subagents: dict[int, dict] = {}
+
+    async def _start_subagent(self, task: str) -> int:
+        """Start a subagent with the given task.
+
+        Args:
+            task: The task description for the subagent
+
+        Returns:
+            Subagent task ID
+        """
+        conversation = prompts_to_conversations([task])[0]
+        assert isinstance(conversation, Conversation)
+
+        # Use agent loop nowait API to start the subagent
+        task_id = self.client.start_agent_loop_nowait(
+            conversation,
+            tools=self.tools,  # type: ignore
+            max_rounds=self.max_rounds,
+        )
+
+        # Track the subagent
+        self.subagents[task_id] = {
+            "status": "running",
+            "conversation": None,
+            "response": None,
+            "error": None,
+        }
+
+        return task_id
+
+    def _finalize_subagent_result(
+        self, agent_id: int, result: AgentLoopResponse
+    ) -> str:
+        """Update subagent tracking state from a finished agent loop."""
+        agent = self.subagents[agent_id]
+        agent["conversation"] = result.conversation
+        agent["response"] = result.final_response
+
+        if result.final_response.is_error:
+            agent["status"] = "error"
+            agent["error"] = result.final_response.error_message
+            return f"Error: {agent['error']}"
+
+        agent["status"] = "finished"
+        return result.final_response.completion or "Subagent finished with no output"
+
+    async def _check_subagent(self, agent_id: int) -> str:
+        """Check the status of a subagent.
+
+        Args:
+            agent_id: The subagent task ID
+
+        Returns:
+            Status string describing the subagent's state
+        """
+        if agent_id not in self.subagents:
+            return f"Error: Subagent {agent_id} not found"
+
+        agent = self.subagents[agent_id]
+        status = agent["status"]
+
+        if status == "finished":
+            response: APIResponse = agent["response"]
+            return response.completion or "Subagent finished with no output"
+        elif status == "error":
+            return f"Error: {agent['error']}"
+        else:
+            # Try to check if it's done
+            try:
+                # Check if the task exists in client's results
+                stored_result = self.client._results.get(agent_id)
+                if isinstance(stored_result, AgentLoopResponse):
+                    return self._finalize_subagent_result(agent_id, stored_result)
+
+                task = self.client._tasks.get(agent_id)
+                if task and task.done():
+                    try:
+                        task_result = task.result()
+                    except Exception as e:
+                        agent["status"] = "error"
+                        agent["error"] = str(e)
+                        return f"Error: {agent['error']}"
+
+                    if isinstance(task_result, AgentLoopResponse):
+                        return self._finalize_subagent_result(agent_id, task_result)
+
+                    agent["status"] = "error"
+                    agent["error"] = (
+                        f"Unexpected task result type: {type(task_result).__name__}"
+                    )
+                    return f"Error: {agent['error']}"
+
+                # Still running
+                return f"Subagent {agent_id} is still running. Call this tool again to check status."
+            except Exception as e:
+                agent["status"] = "error"
+                agent["error"] = str(e)
+                return f"Error checking subagent: {e}"
+
+    async def _wait_for_subagent(self, agent_id: int) -> str:
+        """Wait for a subagent to complete and return its output.
+
+        Args:
+            agent_id: The subagent task ID
+
+        Returns:
+            The subagent's final output
+        """
+        if agent_id not in self.subagents:
+            return f"Error: Subagent {agent_id} not found"
+
+        try:
+            # Use the wait_for_agent_loop API
+            conversation, response = await self.client.wait_for_agent_loop(agent_id)
+
+            agent = self.subagents[agent_id]
+            agent["conversation"] = conversation
+            agent["response"] = response
+
+            if response.is_error:
+                agent["status"] = "error"
+                agent["error"] = response.error_message
+                return f"Error: {response.error_message}"
+            else:
+                agent["status"] = "finished"
+                return response.completion or "Subagent finished with no output"
+        except Exception as e:
+            agent = self.subagents[agent_id]
+            agent["status"] = "error"
+            agent["error"] = str(e)
+            return f"Error waiting for subagent: {e}"
+
+    def get_tools(self) -> list[Tool]:
+        """Get the tools that allow a main LLM to control subagents.
+
+        Returns:
+            List of Tool objects for starting, checking, and waiting for subagents
+        """
+        start_tool = Tool(
+            name="start_subagent",
+            description=(
+                "Start a subagent to work on a subtask independently. "
+                "Use this to delegate complex subtasks or when you need to save context. "
+                "Returns the subagent's task ID which can be used to check its status."
+            ),
+            run=self._start_subagent,
+            parameters={
+                "task": {
+                    "type": "string",
+                    "description": "The task description for the subagent to work on",
+                }
+            },
+            required=["task"],
+        )
+
+        check_tool = Tool(
+            name="check_subagent",
+            description=(
+                "Check the status and output of a running subagent. "
+                "If the subagent is still running, you'll be told to check again later. "
+                "If finished, returns the subagent's final output."
+            ),
+            run=self._check_subagent,
+            parameters={
+                "agent_id": {
+                    "type": "integer",
+                    "description": "The task ID of the subagent to check",
+                }
+            },
+            required=["agent_id"],
+        )
+
+        wait_tool = Tool(
+            name="wait_for_subagent",
+            description=(
+                "Wait for a subagent to complete and return its output. "
+                "This will block until the subagent finishes. "
+                "Use check_subagent if you want to do other work while waiting."
+            ),
+            run=self._wait_for_subagent,
+            parameters={
+                "agent_id": {
+                    "type": "integer",
+                    "description": "The task ID of the subagent to wait for",
+                }
+            },
+            required=["agent_id"],
+        )
+
+        return [start_tool, check_tool, wait_tool]