kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/app/lib/planning_tools/environments.py

@@ -0,0 +1,218 @@
+"""
+Environment Context Tools - Fetch environment and infrastructure information
+"""
+
+from typing import Optional
+from sqlalchemy.orm import Session
+from control_plane_api.app.lib.planning_tools.base import BasePlanningTools
+from control_plane_api.app.lib.supabase import get_supabase
+
+
+class EnvironmentsContextTools(BasePlanningTools):
+    """
+    Tools for fetching environment and infrastructure context
+
+    Provides methods to:
+    - List available execution environments
+    - Get worker queue information
+    - Check resource availability
+    - Query runtime configurations
+    """
+
+    def __init__(self, db: Optional[Session] = None, organization_id: Optional[str] = None):
+        super().__init__(db=db, organization_id=organization_id)
+        self.name = "environment_context_tools"
+
+    async def list_environments(self, limit: int = 50) -> str:
+        """
+        List all available execution environments
+
+        Args:
+            limit: Maximum number of environments to return
+
+        Returns:
+            Formatted string with environment information including:
+            - Environment name and ID
+            - Status (active/inactive)
+            - Available resources
+            - Configuration details
+        """
+        try:
+            # Use Supabase for environments
+            client = get_supabase()
+
+            query = client.table("environments").select("*")
+            if self.organization_id:
+                query = query.eq("organization_id", self.organization_id)
+
+            result = query.limit(limit).execute()
+            environments = result.data or []
+
+            return self._format_list_response(
+                items=environments,
+                title="Available Environments",
+                key_fields=["status", "type", "region"],
+            )
+
+        except Exception as e:
+            return f"Error listing environments: {str(e)}"
+
+    async def get_environment_details(self, environment_id: str) -> str:
+        """
+        Get detailed information about a specific environment
+
+        Args:
+            environment_id: ID of the environment to fetch
+
+        Returns:
+            Detailed environment information including:
+            - Full configuration
+            - Resource limits
+            - Network settings
+            - Security policies
+        """
+        try:
+            response = await self._make_request("GET", f"/environments/{environment_id}")
+
+            return self._format_detail_response(
+                item=response,
+                title=f"Environment Details: {response.get('name', 'Unknown')}",
+            )
+
+        except Exception as e:
+            return f"Error fetching environment {environment_id}: {str(e)}"
+
+    async def list_worker_queues(self, environment_id: Optional[str] = None) -> str:
+        """
+        List available worker queues for task execution
+
+        Args:
+            environment_id: Optional environment ID to filter queues
+
+        Returns:
+            List of worker queues with:
+            - Queue name and ID
+            - Active worker count
+            - Queue status
+            - Processing capacity
+        """
+        try:
+            if environment_id:
+                endpoint = f"/environments/{environment_id}/worker-queues"
+            else:
+                endpoint = "/worker-queues"
+
+            params = {}
+            if self.organization_id:
+                params["organization_id"] = self.organization_id
+
+            response = await self._make_request("GET", endpoint, params=params)
+
+            queues = response if isinstance(response, list) else response.get("queues", [])
+
+            return self._format_list_response(
+                items=queues,
+                title="Available Worker Queues",
+                key_fields=["status", "active_workers", "pending_tasks", "environment_id"],
+            )
+
+        except Exception as e:
+            return f"Error listing worker queues: {str(e)}"
+
+    async def get_worker_queue_details(self, queue_id: str) -> str:
+        """
+        Get detailed information about a specific worker queue
+
+        Args:
+            queue_id: ID of the worker queue
+
+        Returns:
+            Detailed queue information including:
+            - Worker capacity and utilization
+            - Queue statistics
+            - Configuration
+        """
+        try:
+            response = await self._make_request("GET", f"/worker-queues/{queue_id}")
+
+            return self._format_detail_response(
+                item=response,
+                title=f"Worker Queue Details: {response.get('name', 'Unknown')}",
+            )
+
+        except Exception as e:
+            return f"Error fetching worker queue {queue_id}: {str(e)}"
+
+    async def check_resource_availability(self, environment_id: str) -> str:
+        """
+        Check resource availability in a specific environment
+
+        Args:
+            environment_id: ID of the environment to check
+
+        Returns:
+            Resource availability status including:
+            - CPU/Memory availability
+            - Active workers
+            - Queue capacity
+        """
+        try:
+            # Get environment details
+            env_response = await self._make_request("GET", f"/environments/{environment_id}")
+
+            # Get worker queues for this environment
+            queues_response = await self._make_request(
+                "GET",
+                f"/environments/{environment_id}/worker-queues"
+            )
+
+            queues = queues_response if isinstance(queues_response, list) else queues_response.get("queues", [])
+
+            total_workers = sum(q.get("active_workers", 0) for q in queues)
+            total_pending = sum(q.get("pending_tasks", 0) for q in queues)
+
+            output = [
+                f"Resource Availability for Environment: {env_response.get('name', 'Unknown')}",
+                f"  Status: {env_response.get('status', 'unknown')}",
+                f"  Worker Queues: {len(queues)}",
+                f"  Total Active Workers: {total_workers}",
+                f"  Total Pending Tasks: {total_pending}",
+                "",
+                "Queue Details:",
+            ]
+
+            for queue in queues:
+                output.append(
+                    f"  - {queue.get('name')}: "
+                    f"{queue.get('active_workers', 0)} workers, "
+                    f"{queue.get('pending_tasks', 0)} pending"
+                )
+
+            return "\n".join(output)
+
+        except Exception as e:
+            return f"Error checking resource availability: {str(e)}"
+
+    async def list_runtime_configurations(self) -> str:
+        """
+        List available runtime configurations
+
+        Returns:
+            Available runtime configurations including:
+            - Runtime types (docker, k8s, etc.)
+            - Resource templates
+            - Default configurations
+        """
+        try:
+            response = await self._make_request("GET", "/runtimes")
+
+            runtimes = response if isinstance(response, list) else response.get("runtimes", [])
+
+            return self._format_list_response(
+                items=runtimes,
+                title="Available Runtime Configurations",
+                key_fields=["type", "version", "resource_limits"],
+            )
+
+        except Exception as e:
+            return f"Error listing runtime configurations: {str(e)}"

control_plane_api/app/lib/planning_tools/knowledge.py

@@ -0,0 +1,204 @@
+"""
+Knowledge Context Tools - Fetch organizational knowledge for task planning
+"""
+
+import os
+from typing import Optional, List
+import structlog
+import httpx
+from control_plane_api.app.lib.planning_tools.base import BasePlanningTools
+
+logger = structlog.get_logger()
+
+KNOWLEDGE_API_BASE_URL = os.environ.get("KUBIYA_API_BASE", "https://api.kubiya.ai")
+ORCHESTRATOR_API_BASE_URL = os.environ.get("ORCHESTRATOR_API_BASE", "https://orchestrator.kubiya.ai")
+
+class KnowledgeContextTools(BasePlanningTools):
+    """
+    Tools for fetching organizational knowledge and documentation
+
+    Provides methods to:
+    - Search manual knowledge base (user-entered private docs)
+    - Query organizational knowledge (Slack, Confluence, GitHub, etc.)
+    - Get relevant documentation
+    - Query best practices and guidelines
+    """
+
+    def __init__(
+        self,
+        db=None,
+        organization_id: Optional[str] = None,
+        api_token: Optional[str] = None,
+        api_base_url: str = f"{KNOWLEDGE_API_BASE_URL}/api/v1",
+        orchestrator_base_url: str = ORCHESTRATOR_API_BASE_URL
+    ):
+        super().__init__(db=db, organization_id=organization_id)
+        self.name = "knowledge_context_tools"
+        self.api_token = api_token
+        self.api_base_url = api_base_url.rstrip("/")
+        self.orchestrator_base_url = orchestrator_base_url.rstrip("/")
+
+    async def get_all_knowledge(self) -> List[dict]:
+        """
+        Get all organizational knowledge and documentation
+
+        Endpoint: GET /api/v1/knowledge
+        Returns all knowledge items for the organization
+
+        Returns:
+            List of all knowledge items with complete data
+        """
+        if not self.api_token:
+            logger.warning("knowledge_api_token_not_provided")
+            return []
+
+        try:
+            headers = {
+                "Authorization": f"Bearer {self.api_token}",
+                "Content-Type": "application/json"
+            }
+
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                response = await client.get(
+                    f"{self.api_base_url}/knowledge",
+                    headers=headers
+                )
+
+                if response.status_code == 200:
+                    data = response.json()
+                    knowledge_items = data if isinstance(data, list) else data.get("data", [])
+
+                    logger.info(
+                        "knowledge_fetch_success",
+                        count=len(knowledge_items)
+                    )
+                    return knowledge_items
+                else:
+                    logger.error(
+                        "knowledge_fetch_failed",
+                        status_code=response.status_code,
+                        error=response.text
+                    )
+                    return []
+
+        except Exception as e:
+            logger.error("knowledge_fetch_error", error=str(e))
+            return []
+
+    async def query_knowledge(self, query: str) -> List[str]:
+        """
+        Query the knowledge base with a specific question or search term
+
+        Endpoint: POST /api/v1/knowledge/query
+        Searches through organizational knowledge and returns relevant results
+
+        Args:
+            query: Search query or question (e.g., "how to deploy to AWS?", "raz==superman ? omer==?")
+
+        Returns:
+            List of relevant knowledge items/answers as strings
+        """
+        if not self.api_token:
+            logger.warning("knowledge_api_token_not_provided")
+            return []
+
+        try:
+            headers = {
+                "Authorization": f"Bearer {self.api_token}",
+                "Content-Type": "application/json"
+            }
+
+            payload = {"query": query}
+
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                response = await client.post(
+                    f"{self.api_base_url}/knowledge/query",
+                    headers=headers,
+                    json=payload
+                )
+
+                if response.status_code == 200:
+                    data = response.json()
+                    # API returns a list of strings
+                    results = data if isinstance(data, list) else []
+
+                    logger.info(
+                        "knowledge_query_success",
+                        query=query,
+                        results_count=len(results)
+                    )
+                    return results
+                else:
+                    logger.error(
+                        "knowledge_query_failed",
+                        query=query,
+                        status_code=response.status_code,
+                        error=response.text
+                    )
+                    return []
+
+        except Exception as e:
+            logger.error("knowledge_query_error", query=query, error=str(e))
+            return []
+
+    async def query_organizational_knowledge(self, query: str, limit: int = 5) -> List[dict]:
+        """
+        Query organizational knowledge from integrated sources (Slack, Confluence, GitHub, etc.)
+
+        Endpoint: POST https://orchestrator.kubiya.ai/api/query
+        Searches through Slack messages, Confluence pages, GitHub discussions, and other integrated sources
+
+        Args:
+            query: Search query or question (e.g., "I seems that version 0.3.37 is broken")
+            limit: Maximum number of results to return (default: 5)
+
+        Returns:
+            List of relevant knowledge items from organizational sources
+        """
+        if not self.api_token:
+            logger.warning("orchestrator_api_token_not_provided")
+            return []
+
+        try:
+            headers = {
+                "Authorization": f"Bearer {self.api_token}",
+                "Content-Type": "application/json",
+                "Accept": "application/json"
+            }
+
+            payload = {
+                "query": query,
+                "limit": limit
+            }
+
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                response = await client.post(
+                    f"{self.orchestrator_base_url}/api/query",
+                    headers=headers,
+                    json=payload
+                )
+
+                if response.status_code == 200:
+                    data = response.json()
+                    # API returns results as a list or dict with results key
+                    results = data if isinstance(data, list) else data.get("results", data.get("data", []))
+
+                    logger.info(
+                        "organizational_knowledge_query_success",
+                        query=query,
+                        results_count=len(results) if isinstance(results, list) else 0
+                    )
+                    return results if isinstance(results, list) else []
+                else:
+                    logger.error(
+                        "organizational_knowledge_query_failed",
+                        query=query,
+                        status_code=response.status_code,
+                        error=response.text
+                    )
+                    return []
+
+        except Exception as e:
+            logger.error("organizational_knowledge_query_error", query=query, error=str(e))
+            return []
+

control_plane_api/app/lib/planning_tools/models.py

@@ -0,0 +1,93 @@
+"""
+Pydantic models for planning tool responses
+
+These models ensure type safety and consistent data structures
+when tools return information to the planning agent.
+"""
+
+from pydantic import BaseModel, Field
+from typing import List, Optional, Dict, Any
+from datetime import datetime
+
+
+class AgentModel(BaseModel):
+    """Agent information model with full context for planning"""
+    id: str
+    name: str
+    model_id: Optional[str] = Field(default="default")
+    description: Optional[str] = None
+    status: str
+    capabilities: Optional[List[str]] = None
+    configuration: Dict[str, Any] = Field(default_factory=dict)
+    llm_config: Dict[str, Any] = Field(default_factory=dict)
+    runtime: Optional[str] = None
+    team_id: Optional[str] = None
+    skills: List[Dict[str, Any]] = Field(default_factory=list)  # Full skill details
+    skill_ids: List[str] = Field(default_factory=list)
+    projects: List[Dict[str, Any]] = Field(default_factory=list)
+    environments: List[Dict[str, Any]] = Field(default_factory=list)
+    execution_environment: Optional[Dict[str, Any]] = None  # env_vars, secrets, integration_ids
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+    last_active_at: Optional[datetime] = None
+    error_message: Optional[str] = None
+
+
+class TeamMemberModel(BaseModel):
+    """Team member information model"""
+    id: str
+    name: str
+    model_id: str = Field(default="default")
+    description: Optional[str] = None
+
+
+class TeamModel(BaseModel):
+    """Team information model"""
+    id: str
+    name: str
+    description: Optional[str] = None
+    status: str
+    agent_count: int = 0
+    agents: Optional[List[TeamMemberModel]] = None
+
+
+class EnvironmentModel(BaseModel):
+    """Environment information model"""
+    id: str
+    name: str
+    display_name: Optional[str] = None
+    description: Optional[str] = None
+    status: str
+    tags: List[str] = Field(default_factory=list)
+    settings: Dict[str, Any] = Field(default_factory=dict)
+
+
+class WorkerQueueModel(BaseModel):
+    """Worker queue information model"""
+    id: str
+    name: str
+    display_name: Optional[str] = None
+    environment_id: Optional[str] = None
+    active_workers: int = 0
+    status: str
+
+
+class ExecutionHistoryModel(BaseModel):
+    """Execution history summary model"""
+    entity_id: str
+    entity_type: str  # "agent" or "team"
+    total_executions: int
+    completed: int
+    failed: int
+    success_rate: float
+    recent_executions: List[Dict[str, Any]] = Field(default_factory=list)
+
+
+class SkillModel(BaseModel):
+    """Skill/tool information model"""
+    id: str
+    name: str
+    skill_type: str
+    description: Optional[str] = None
+    enabled: bool = True
+    configuration: Dict[str, Any] = Field(default_factory=dict)