kubiya-control-plane-api 0.1.0__py3-none-any.whl → 0.3.4__py3-none-any.whl

control_plane_api/app/lib/litellm_pricing.py

@@ -0,0 +1,166 @@
+"""
+LiteLLM Pricing Integration - Fetch and cache model pricing data
+"""
+
+import httpx
+import structlog
+from typing import Dict, Optional
+from datetime import datetime, timedelta
+import asyncio
+
+logger = structlog.get_logger()
+
+# Cache for pricing data
+_pricing_cache: Optional[Dict] = None
+_cache_timestamp: Optional[datetime] = None
+_cache_lock = asyncio.Lock()
+
+PRICING_URL = "https://raw.githubusercontent.com/BerriAI/litellm/refs/heads/main/model_prices_and_context_window.json"
+CACHE_TTL_HOURS = 24  # Refresh pricing data daily
+
+
+async def get_litellm_pricing() -> Dict:
+    """
+    Fetch LiteLLM pricing data with caching
+
+    Returns:
+        Dict containing model pricing information
+    """
+    global _pricing_cache, _cache_timestamp
+
+    async with _cache_lock:
+        # Check if cache is valid
+        if _pricing_cache and _cache_timestamp:
+            age = datetime.utcnow() - _cache_timestamp
+            if age < timedelta(hours=CACHE_TTL_HOURS):
+                logger.debug("litellm_pricing_cache_hit", age_hours=age.total_seconds() / 3600)
+                return _pricing_cache
+
+        # Fetch fresh pricing data
+        logger.info("fetching_litellm_pricing", url=PRICING_URL)
+        try:
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                response = await client.get(PRICING_URL)
+                response.raise_for_status()
+                pricing_data = response.json()
+
+                # Update cache
+                _pricing_cache = pricing_data
+                _cache_timestamp = datetime.utcnow()
+
+                logger.info(
+                    "litellm_pricing_fetched_successfully",
+                    models_count=len(pricing_data),
+                    cached_until=(_cache_timestamp + timedelta(hours=CACHE_TTL_HOURS)).isoformat()
+                )
+
+                return pricing_data
+        except Exception as e:
+            logger.error("litellm_pricing_fetch_failed", error=str(e), error_type=type(e).__name__)
+            # Return empty dict on failure
+            return {}
+
+
+def get_model_pricing(model_id: str, pricing_data: Dict) -> Optional[Dict]:
+    """
+    Get pricing information for a specific model
+
+    Args:
+        model_id: Model identifier (e.g., "claude-sonnet-4", "gpt-4o")
+        pricing_data: Full pricing data from LiteLLM
+
+    Returns:
+        Dict with pricing info or None if not found
+    """
+    # Try exact match first
+    if model_id in pricing_data:
+        return pricing_data[model_id]
+
+    # Try common variations
+    variations = [
+        model_id,
+        f"openai/{model_id}",
+        f"anthropic/{model_id}",
+        f"anthropic.{model_id}",
+        f"bedrock/{model_id}",
+    ]
+
+    for variation in variations:
+        if variation in pricing_data:
+            return pricing_data[variation]
+
+    logger.warning("model_pricing_not_found", model_id=model_id, tried_variations=variations)
+    return None
+
+
+def calculate_llm_cost(
+    model_id: str,
+    estimated_input_tokens: int,
+    estimated_output_tokens: int,
+    pricing_data: Dict
+) -> tuple[float, float, float]:
+    """
+    Calculate LLM cost for a model
+
+    Args:
+        model_id: Model identifier
+        estimated_input_tokens: Expected input tokens
+        estimated_output_tokens: Expected output tokens
+        pricing_data: Full pricing data from LiteLLM
+
+    Returns:
+        Tuple of (cost_per_1k_input, cost_per_1k_output, total_cost)
+    """
+    model_pricing = get_model_pricing(model_id, pricing_data)
+
+    if not model_pricing:
+        # Fallback defaults
+        logger.warning("using_default_pricing", model_id=model_id)
+        return (0.003, 0.015, 0.0)
+
+    input_cost_per_token = model_pricing.get("input_cost_per_token", 0.000003)
+    output_cost_per_token = model_pricing.get("output_cost_per_token", 0.000015)
+
+    # Convert to per-1k pricing for display
+    input_cost_per_1k = input_cost_per_token * 1000
+    output_cost_per_1k = output_cost_per_token * 1000
+
+    # Calculate total cost
+    total_cost = (
+        (estimated_input_tokens * input_cost_per_token) +
+        (estimated_output_tokens * output_cost_per_token)
+    )
+
+    return (input_cost_per_1k, output_cost_per_1k, total_cost)
+
+
+def get_model_display_name(model_id: str) -> str:
+    """
+    Get human-readable model name
+
+    Args:
+        model_id: Model identifier
+
+    Returns:
+        Display name
+    """
+    # Map of common model IDs to display names
+    display_names = {
+        "claude-sonnet-4": "Claude Sonnet 4",
+        "claude-3-5-sonnet-20241022": "Claude 3.5 Sonnet",
+        "claude-3-5-sonnet-20240620": "Claude 3.5 Sonnet (Legacy)",
+        "claude-3-opus-20240229": "Claude 3 Opus",
+        "claude-3-haiku-20240307": "Claude 3 Haiku",
+        "gpt-4o": "GPT-4o",
+        "gpt-4o-mini": "GPT-4o Mini",
+        "gpt-4-turbo": "GPT-4 Turbo",
+        "gpt-4": "GPT-4",
+        "gpt-3.5-turbo": "GPT-3.5 Turbo",
+        "o1": "OpenAI o1",
+        "o1-mini": "OpenAI o1-mini",
+        "gemini-2.0-flash-exp": "Gemini 2.0 Flash",
+        "gemini-1.5-pro": "Gemini 1.5 Pro",
+        "gemini-1.5-flash": "Gemini 1.5 Flash",
+    }
+
+    return display_names.get(model_id, model_id.replace("-", " ").title())

control_plane_api/app/lib/planning_tools/__init__.py

@@ -0,0 +1,22 @@
+"""
+Planning Tools - Modular tools for the task planning agent
+
+This package provides decoupled, maintainable tools organized by category:
+- agents: Agent-related context and operations
+- teams: Team-related context and operations
+- environments: Environment and infrastructure context
+- resources: General resource and capability queries
+- workflows: Workflow and execution context
+"""
+
+from control_plane_api.app.lib.planning_tools.agents import AgentsContextTools
+from control_plane_api.app.lib.planning_tools.teams import TeamsContextTools
+from control_plane_api.app.lib.planning_tools.environments import EnvironmentsContextTools
+from control_plane_api.app.lib.planning_tools.resources import ResourcesContextTools
+
+__all__ = [
+    "AgentsContextTools",
+    "TeamsContextTools",
+    "EnvironmentsContextTools",
+    "ResourcesContextTools",
+]

control_plane_api/app/lib/planning_tools/agents.py

@@ -0,0 +1,155 @@
+"""
+Agent Context Tools - Fetch agent information for task planning
+"""
+
+from typing import Optional
+from control_plane_api.app.lib.planning_tools.base import BasePlanningTools
+
+
+class AgentsContextTools(BasePlanningTools):
+    """
+    Tools for fetching agent context and capabilities
+
+    Provides methods to:
+    - List all available agents
+    - Get detailed agent information
+    - Query agent capabilities
+    - Check agent availability
+    """
+
+    def __init__(self, base_url: str = "http://localhost:8000", organization_id: Optional[str] = None):
+        super().__init__(base_url=base_url, organization_id=organization_id)
+        self.name = "agent_context_tools"
+
+    async def list_agents(self, limit: int = 50) -> str:
+        """
+        List all available agents with their basic information
+
+        Args:
+            limit: Maximum number of agents to return
+
+        Returns:
+            Formatted string with agent information including:
+            - Agent name and ID
+            - Model being used
+            - Description/capabilities
+            - Current status
+        """
+        try:
+            params = {"limit": limit}
+            if self.organization_id:
+                params["organization_id"] = self.organization_id
+
+            response = await self._make_request("GET", "/agents", params=params)
+
+            agents = response if isinstance(response, list) else response.get("agents", [])
+
+            return self._format_list_response(
+                items=agents,
+                title="Available Agents",
+                key_fields=["model_id", "description", "status", "runner_name"],
+            )
+
+        except Exception as e:
+            return f"Error listing agents: {str(e)}"
+
+    async def get_agent_details(self, agent_id: str) -> str:
+        """
+        Get detailed information about a specific agent
+
+        Args:
+            agent_id: ID of the agent to fetch
+
+        Returns:
+            Detailed agent information including:
+            - Full configuration
+            - Available tools/capabilities
+            - Model details
+            - Resource requirements
+        """
+        try:
+            response = await self._make_request("GET", f"/agents/{agent_id}")
+
+            return self._format_detail_response(
+                item=response,
+                title=f"Agent Details: {response.get('name', 'Unknown')}",
+            )
+
+        except Exception as e:
+            return f"Error fetching agent {agent_id}: {str(e)}"
+
+    async def search_agents_by_capability(self, capability: str) -> str:
+        """
+        Search for agents that have a specific capability
+
+        Args:
+            capability: Capability to search for (e.g., "kubernetes", "aws", "python")
+
+        Returns:
+            List of agents matching the capability
+        """
+        try:
+            # First get all agents
+            params = {}
+            if self.organization_id:
+                params["organization_id"] = self.organization_id
+
+            response = await self._make_request("GET", "/agents", params=params)
+            agents = response if isinstance(response, list) else response.get("agents", [])
+
+            # Filter by capability (search in description and tools)
+            matching_agents = []
+            for agent in agents:
+                agent_text = f"{agent.get('name', '')} {agent.get('description', '')}".lower()
+                if capability.lower() in agent_text:
+                    matching_agents.append(agent)
+
+            return self._format_list_response(
+                items=matching_agents,
+                title=f"Agents with '{capability}' capability",
+                key_fields=["model_id", "description"],
+            )
+
+        except Exception as e:
+            return f"Error searching agents: {str(e)}"
+
+    async def get_agent_execution_history(self, agent_id: str, limit: int = 10) -> str:
+        """
+        Get recent execution history for an agent
+
+        Args:
+            agent_id: ID of the agent
+            limit: Number of recent executions to fetch
+
+        Returns:
+            Recent execution history with success rates
+        """
+        try:
+            params = {"limit": limit, "entity_id": agent_id}
+            response = await self._make_request("GET", "/executions", params=params)
+
+            executions = response if isinstance(response, list) else response.get("executions", [])
+
+            if not executions:
+                return f"No execution history found for agent {agent_id}"
+
+            # Calculate success rate
+            completed = sum(1 for e in executions if e.get("status") == "completed")
+            total = len(executions)
+            success_rate = (completed / total * 100) if total > 0 else 0
+
+            output = [
+                f"Execution History for Agent (Last {total} runs):",
+                f"Success Rate: {success_rate:.1f}% ({completed}/{total})",
+                "\nRecent Executions:",
+            ]
+
+            for idx, execution in enumerate(executions[:5], 1):
+                status = execution.get("status", "unknown")
+                prompt = execution.get("prompt", "No description")[:50]
+                output.append(f"{idx}. Status: {status} | Task: {prompt}...")
+
+            return "\n".join(output)
+
+        except Exception as e:
+            return f"Error fetching execution history: {str(e)}"

control_plane_api/app/lib/planning_tools/base.py

@@ -0,0 +1,189 @@
+"""
+Base classes and utilities for planning tools
+"""
+
+from typing import Optional, Dict, Any, List
+from agno.tools.toolkit import Toolkit
+import structlog
+import httpx
+import json
+
+logger = structlog.get_logger()
+
+
+class BasePlanningTools(Toolkit):
+    """
+    Base class for all planning tools with common utilities
+
+    Provides:
+    - HTTP client for API calls
+    - Common error handling
+    - Response formatting
+    - Logging utilities
+    """
+
+    def __init__(
+        self,
+        base_url: str = "http://localhost:8000",
+        organization_id: Optional[str] = None,
+        timeout: int = 30,
+    ):
+        """
+        Initialize base planning tools
+
+        Args:
+            base_url: Base URL for the control plane API
+            organization_id: Organization context for filtering
+            timeout: HTTP request timeout in seconds
+        """
+        super().__init__(name="base_planning_tools")
+        self.base_url = base_url.rstrip("/")
+        self.organization_id = organization_id
+        self.timeout = timeout
+        self._client: Optional[httpx.AsyncClient] = None
+
+    async def _get_client(self) -> httpx.AsyncClient:
+        """Get or create HTTP client"""
+        if self._client is None:
+            self._client = httpx.AsyncClient(
+                base_url=self.base_url,
+                timeout=self.timeout,
+                headers={"Content-Type": "application/json"},
+            )
+        return self._client
+
+    async def _make_request(
+        self,
+        method: str,
+        endpoint: str,
+        params: Optional[Dict[str, Any]] = None,
+        data: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Make HTTP request with error handling
+
+        Args:
+            method: HTTP method (GET, POST, etc.)
+            endpoint: API endpoint (will be appended to base_url)
+            params: Query parameters
+            data: Request body data
+
+        Returns:
+            Parsed JSON response
+
+        Raises:
+            Exception: If request fails
+        """
+        client = await self._get_client()
+
+        try:
+            logger.info(
+                "planning_tool_request",
+                method=method,
+                endpoint=endpoint,
+                has_params=bool(params),
+                has_data=bool(data),
+            )
+
+            response = await client.request(
+                method=method,
+                url=endpoint,
+                params=params,
+                json=data,
+            )
+            response.raise_for_status()
+
+            result = response.json()
+            logger.info(
+                "planning_tool_response",
+                endpoint=endpoint,
+                status_code=response.status_code,
+            )
+
+            return result
+
+        except httpx.HTTPStatusError as e:
+            logger.error(
+                "planning_tool_http_error",
+                endpoint=endpoint,
+                status_code=e.response.status_code,
+                error=str(e),
+            )
+            raise Exception(f"HTTP {e.response.status_code}: {str(e)}")
+        except Exception as e:
+            logger.error(
+                "planning_tool_error",
+                endpoint=endpoint,
+                error=str(e),
+                error_type=type(e).__name__,
+            )
+            raise
+
+    def _format_list_response(
+        self,
+        items: List[Dict[str, Any]],
+        title: str,
+        key_fields: List[str],
+    ) -> str:
+        """
+        Format a list of items as a readable string
+
+        Args:
+            items: List of items to format
+            title: Title for the list
+            key_fields: Fields to include in the output
+
+        Returns:
+            Formatted string representation
+        """
+        if not items:
+            return f"{title}: None available"
+
+        output = [f"{title} ({len(items)} total):"]
+        for idx, item in enumerate(items, 1):
+            output.append(f"\n{idx}. {item.get('name', 'Unnamed')} (ID: {item.get('id', 'N/A')})")
+            for field in key_fields:
+                if field in item and item[field]:
+                    value = item[field]
+                    # Format nested objects
+                    if isinstance(value, dict):
+                        value = json.dumps(value, indent=2)
+                    elif isinstance(value, list):
+                        value = f"{len(value)} items"
+                    output.append(f"   - {field}: {value}")
+
+        return "\n".join(output)
+
+    def _format_detail_response(
+        self,
+        item: Dict[str, Any],
+        title: str,
+    ) -> str:
+        """
+        Format a single item as a readable string
+
+        Args:
+            item: Item to format
+            title: Title for the item
+
+        Returns:
+            Formatted string representation
+        """
+        if not item:
+            return f"{title}: Not found"
+
+        output = [f"{title}:"]
+        for key, value in item.items():
+            if isinstance(value, dict):
+                value = json.dumps(value, indent=2)
+            elif isinstance(value, list):
+                value = f"{len(value)} items: {', '.join([str(v) for v in value[:3]])}{'...' if len(value) > 3 else ''}"
+            output.append(f"  {key}: {value}")
+
+        return "\n".join(output)
+
+    async def close(self):
+        """Close HTTP client"""
+        if self._client:
+            await self._client.aclose()
+            self._client = None