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

control_plane_api/app/services/litellm_service.py

@@ -0,0 +1,394 @@
+"""
+LiteLLM Service
+
+This service provides a wrapper around LiteLLM for agent execution.
+"""
+
+import os
+from typing import Dict, List, Optional, Any
+import litellm
+from litellm import completion, get_valid_models
+import logging
+import httpx
+
+from control_plane_api.app.config import settings
+
+logger = logging.getLogger(__name__)
+
+
+class LiteLLMService:
+    """Service for interacting with LiteLLM"""
+
+    def __init__(self):
+        """Initialize LiteLLM service with configuration"""
+        # Set LiteLLM configuration
+        if settings.litellm_api_key:
+            os.environ["LITELLM_API_KEY"] = settings.litellm_api_key
+        litellm.api_base = settings.litellm_api_base
+        litellm.drop_params = True  # Drop unsupported params instead of failing
+
+        # Configure timeout
+        litellm.request_timeout = settings.litellm_timeout
+
+        logger.info(f"LiteLLM Service initialized with base URL: {settings.litellm_api_base}")
+
+    async def fetch_available_models(self, base_url: Optional[str] = None) -> List[Dict[str, Any]]:
+        """
+        Fetch available models from the LiteLLM proxy server.
+
+        This method tries multiple approaches:
+        1. HTTP API call to /model/info endpoint (includes mode field)
+        2. Fallback to /v1/models endpoint with mode detection
+        3. Fallback to using LiteLLM SDK's get_valid_models() if available
+
+        Args:
+            base_url: Optional override for the LiteLLM base URL
+
+        Returns:
+            List of model objects with mode field from the LiteLLM server
+
+        Raises:
+            Exception if all methods fail
+        """
+        # Use provided base_url or fall back to settings
+        api_base = base_url or settings.litellm_api_base
+
+        # Try Method 1: /model/info endpoint (includes mode)
+        try:
+            model_info_url = f"{api_base.rstrip('/')}/model/info"
+            logger.info(f"Fetching models from LiteLLM /model/info endpoint: {model_info_url}")
+
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                headers = {}
+                if settings.litellm_api_key:
+                    headers["Authorization"] = f"Bearer {settings.litellm_api_key}"
+
+                response = await client.get(model_info_url, headers=headers)
+                response.raise_for_status()
+
+                data = response.json()
+
+                # Parse model_info response
+                models = []
+                for model_data in data.get("data", []):
+                    models.append({
+                        "id": model_data.get("model_name"),
+                        "object": "model",
+                        "created": 0,
+                        "owned_by": model_data.get("litellm_provider", "unknown"),
+                        "mode": model_data.get("mode", "completion")  # Include mode from LiteLLM
+                    })
+
+                logger.info(f"Successfully fetched {len(models)} models from LiteLLM /model/info with mode information")
+                return models
+
+        except httpx.HTTPStatusError as e:
+            logger.warning(
+                f"/model/info endpoint failed: {e.response.status_code}. "
+                f"Trying fallback to /v1/models..."
+            )
+
+        except Exception as e:
+            logger.warning(f"/model/info request failed: {str(e)}. Trying fallback to /v1/models...")
+
+        # Try Method 2: /v1/models endpoint (detect mode from model name)
+        try:
+            models_url = f"{api_base.rstrip('/')}/v1/models"
+            logger.info(f"Fetching models from LiteLLM server: {models_url}")
+
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                headers = {}
+                if settings.litellm_api_key:
+                    headers["Authorization"] = f"Bearer {settings.litellm_api_key}"
+
+                response = await client.get(models_url, headers=headers)
+                response.raise_for_status()
+
+                data = response.json()
+
+                # LiteLLM returns models in OpenAI format: {"data": [...], "object": "list"}
+                raw_models = data.get("data", [])
+
+                # Add mode detection from model name
+                models = []
+                for model in raw_models:
+                    model_id = model.get("id", "")
+                    model_lower = model_id.lower()
+
+                    # Detect mode from model name - expanded detection for embedding models
+                    # Check for embedding keywords (including common embedding model patterns)
+                    embedding_keywords = [
+                        "embedding", "embed", "embeddings",
+                        "ada-002", "text-embedding",
+                        "bge-", "gte-", "e5-",  # Common embedding model prefixes
+                        "voyage-", "cohere-embed",  # Provider-specific embedding models
+                        "-embed-", "_embed_", "/embed",  # Pattern matching
+                    ]
+
+                    if any(keyword in model_lower for keyword in embedding_keywords):
+                        mode = "embedding"
+                    elif any(keyword in model_lower for keyword in ["gpt", "claude", "llama", "mistral", "gemini", "deepseek", "qwen"]):
+                        mode = "chat"
+                    else:
+                        mode = "completion"
+
+                    model["mode"] = mode
+                    models.append(model)
+
+                logger.info(f"Successfully fetched {len(models)} models from LiteLLM server via /v1/models with mode detection")
+                return models
+
+        except httpx.HTTPStatusError as e:
+            logger.warning(
+                f"HTTP endpoint failed: {e.response.status_code}. "
+                f"Trying SDK fallback method..."
+            )
+
+        except Exception as e:
+            logger.warning(f"HTTP request failed: {str(e)}. Trying SDK fallback method...")
+
+        # Try Method 3: Use LiteLLM Python SDK (synchronous, needs to be run in executor)
+        try:
+            import asyncio
+            from concurrent.futures import ThreadPoolExecutor
+
+            def _get_models_sync():
+                """Synchronous function to get models using LiteLLM SDK"""
+                try:
+                    # Set environment for LiteLLM
+                    if settings.litellm_api_key:
+                        os.environ["LITELLM_API_KEY"] = settings.litellm_api_key
+
+                    # Use get_valid_models with provider endpoint check
+                    model_names = get_valid_models(check_provider_endpoint=True)
+
+                    # Convert model strings to OpenAI-like format with mode detection
+                    models = []
+                    for model_name in model_names:
+                        model_lower = model_name.lower()
+
+                        # Detect mode from model name - expanded detection for embedding models
+                        # Check for embedding keywords (including common embedding model patterns)
+                        embedding_keywords = [
+                            "embedding", "embed", "embeddings",
+                            "ada-002", "text-embedding",
+                            "bge-", "gte-", "e5-",  # Common embedding model prefixes
+                            "voyage-", "cohere-embed",  # Provider-specific embedding models
+                            "-embed-", "_embed_", "/embed",  # Pattern matching
+                        ]
+
+                        if any(keyword in model_lower for keyword in embedding_keywords):
+                            mode = "embedding"
+                        elif any(keyword in model_lower for keyword in ["gpt", "claude", "llama", "mistral", "gemini", "deepseek", "qwen"]):
+                            mode = "chat"
+                        else:
+                            mode = "completion"
+
+                        models.append({
+                            "id": model_name,
+                            "object": "model",
+                            "created": 0,
+                            "owned_by": "litellm",
+                            "mode": mode
+                        })
+
+                    return models
+                except Exception as e:
+                    logger.error(f"LiteLLM SDK get_valid_models failed: {str(e)}")
+                    return []
+
+            # Run synchronous function in thread pool
+            loop = asyncio.get_event_loop()
+            with ThreadPoolExecutor() as executor:
+                models = await loop.run_in_executor(executor, _get_models_sync)
+
+            if models:
+                logger.info(f"Successfully fetched {len(models)} models using LiteLLM SDK with mode detection")
+                return models
+
+        except Exception as e:
+            logger.error(f"LiteLLM SDK method failed: {str(e)}")
+
+        # If all methods failed, raise exception
+        raise Exception(
+            f"Failed to fetch models from LiteLLM server. "
+            f"Tried HTTP endpoint and SDK method. "
+            f"Please ensure LiteLLM proxy is running at {api_base}"
+        )
+
+    def validate_model(self, model_id: str) -> bool:
+        """
+        Validate if a model ID is supported by the LiteLLM server.
+
+        Args:
+            model_id: The model identifier to validate
+
+        Returns:
+            True if the model is valid, False otherwise
+        """
+        try:
+            # For now, we'll accept any model ID that follows the provider/model format
+            # More sophisticated validation can be added later by checking against
+            # the models list from the server
+            return bool(model_id and "/" in model_id)
+        except Exception as e:
+            logger.error(f"Error validating model {model_id}: {str(e)}")
+            return False
+
+
+    def execute_agent(
+        self,
+        prompt: str,
+        model: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Execute an agent with LiteLLM
+
+        Args:
+            prompt: The user prompt
+            model: Model identifier (defaults to configured default)
+            system_prompt: System prompt for the agent
+            temperature: Temperature for response generation
+            max_tokens: Maximum tokens to generate
+            top_p: Top-p sampling parameter
+            **kwargs: Additional parameters to pass to LiteLLM
+
+        Returns:
+            Dict containing the response and metadata
+        """
+        try:
+            # Use default model if not specified
+            if not model:
+                model = settings.litellm_default_model
+
+            # Build messages
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.append({"role": "user", "content": prompt})
+
+            # Prepare completion parameters
+            # For custom proxies, use openai/ prefix to force OpenAI-compatible mode
+            # This tells LiteLLM to use the base_url as an OpenAI-compatible endpoint
+            completion_params = {
+                "model": f"openai/{model}",  # Use openai/ prefix for custom proxy
+                "messages": messages,
+                "temperature": temperature,
+                "api_key": settings.litellm_api_key or "dummy-key",  # Fallback for when key is not set
+                "base_url": settings.litellm_api_base,
+            }
+
+            if max_tokens:
+                completion_params["max_tokens"] = max_tokens
+            if top_p:
+                completion_params["top_p"] = top_p
+
+            # Add any additional kwargs
+            completion_params.update(kwargs)
+
+            logger.info(f"Executing agent with model: {model} (using openai/{model})")
+
+            # Make the completion request
+            response = completion(**completion_params)
+
+            # Extract response content
+            result = {
+                "success": True,
+                "response": response.choices[0].message.content,
+                "model": model,
+                "usage": {
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "completion_tokens": response.usage.completion_tokens,
+                    "total_tokens": response.usage.total_tokens,
+                },
+                "finish_reason": response.choices[0].finish_reason,
+            }
+
+            logger.info(f"Agent execution successful. Tokens used: {result['usage']['total_tokens']}")
+            return result
+
+        except Exception as e:
+            logger.error(f"Error executing agent: {str(e)}")
+            return {
+                "success": False,
+                "error": str(e),
+                "model": model or settings.litellm_default_model,
+            }
+
+    def execute_agent_stream(
+        self,
+        prompt: str,
+        model: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        top_p: Optional[float] = None,
+        **kwargs: Any,
+    ):
+        """
+        Execute an agent with streaming response
+
+        Args:
+            prompt: The user prompt
+            model: Model identifier (defaults to configured default)
+            system_prompt: System prompt for the agent
+            temperature: Temperature for response generation
+            max_tokens: Maximum tokens to generate
+            top_p: Top-p sampling parameter
+            **kwargs: Additional parameters to pass to LiteLLM
+
+        Yields:
+            Response chunks as they arrive
+        """
+        try:
+            # Use default model if not specified
+            if not model:
+                model = settings.litellm_default_model
+
+            # Build messages
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.append({"role": "user", "content": prompt})
+
+            # Prepare completion parameters
+            # For custom proxies, use openai/ prefix to force OpenAI-compatible mode
+            # This tells LiteLLM to use the base_url as an OpenAI-compatible endpoint
+            completion_params = {
+                "model": f"openai/{model}",  # Use openai/ prefix for custom proxy
+                "messages": messages,
+                "temperature": temperature,
+                "stream": True,
+                "api_key": settings.litellm_api_key or "dummy-key",  # Fallback for when key is not set
+                "base_url": settings.litellm_api_base,
+            }
+
+            if max_tokens:
+                completion_params["max_tokens"] = max_tokens
+            if top_p:
+                completion_params["top_p"] = top_p
+
+            # Add any additional kwargs
+            completion_params.update(kwargs)
+
+            logger.info(f"Executing agent (streaming) with model: {model} (using openai/{model})")
+
+            # Make the streaming completion request
+            response = completion(**completion_params)
+
+            for chunk in response:
+                if chunk.choices[0].delta.content:
+                    yield chunk.choices[0].delta.content
+
+        except Exception as e:
+            logger.error(f"Error executing agent (streaming): {str(e)}")
+            yield f"Error: {str(e)}"
+
+
+# Singleton instance
+litellm_service = LiteLLMService()

control_plane_api/app/services/plan_generator.py

@@ -0,0 +1,79 @@
+"""
+Shared plan generation service - used by both sync and async endpoints.
+Clean code with zero duplication.
+"""
+from typing import Dict, Any, Callable
+from sqlalchemy.orm import Session
+import structlog
+
+from control_plane_api.app.models.task_planning import TaskPlanRequest, TaskPlanResponse
+from control_plane_api.app.lib.task_planning.planning_workflow import (
+    create_planning_workflow as create_multistep_workflow,
+    run_planning_workflow_stream,
+)
+from control_plane_api.app.lib.task_planning.entity_resolver import resolve_plan_entities
+
+logger = structlog.get_logger()
+
+
+async def generate_plan(
+    request: TaskPlanRequest,
+    organization_id: str,
+    api_token: str,
+    db: Session,
+    event_callback: Callable[[Dict[str, Any]], None] = None,
+) -> TaskPlanResponse:
+    """
+    Generate a plan using the Agno workflow.
+
+    This is the SINGLE source of truth for plan generation,
+    used by both /tasks/plan/stream and async plan generation.
+
+    Args:
+        request: Task plan request
+        organization_id: Organization ID
+        api_token: API token for authentication
+        db: Database session
+        event_callback: Optional callback for streaming events
+
+    Returns:
+        Generated TaskPlanResponse with resolved entities
+    """
+    logger.info(
+        "generating_plan",
+        description=request.description[:100],
+        organization_id=organization_id,
+        quick_mode=request.quick_mode,
+    )
+
+    # Create the 2-step planning workflow
+    workflow = create_multistep_workflow(
+        db=db,
+        organization_id=organization_id,
+        api_token=api_token,
+        quick_mode=request.quick_mode,
+        outer_context=None,
+    )
+
+    logger.info("planning_workflow_created")
+
+    # Run the workflow with event streaming
+    plan = run_planning_workflow_stream(
+        workflow,
+        request,
+        event_callback or (lambda x: None),  # No-op callback if not provided
+        request.quick_mode,
+    )
+
+    logger.info("plan_generated", plan_title=plan.title)
+
+    # Resolve entity names to UUIDs
+    await resolve_plan_entities(
+        plan_response=plan,
+        organization_id=organization_id,
+        db=db,
+    )
+
+    logger.info("plan_entities_resolved", plan_title=plan.title)
+
+    return plan

control_plane_api/app/services/planning_strategy.py

@@ -0,0 +1,66 @@
+"""
+Planning Strategy Pattern
+
+Defines the interface for different task planning implementations (Agno, Claude Code SDK, etc.)
+Like choosing transportation: train, walk, or flight - same destination, different approaches.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Any, AsyncIterator
+from sqlalchemy.orm import Session
+
+from control_plane_api.app.models.task_planning import TaskPlanResponse
+
+
+class PlanningStrategy(ABC):
+    """
+    Abstract base class for task planning strategies.
+
+    Each strategy implements the same interface but uses different underlying
+    technology (Agno, Claude Code SDK, etc.)
+    """
+
+    def __init__(self, db: Session, organization_id: str = None, api_token: str = None):
+        """
+        Initialize strategy.
+
+        Args:
+            db: Database session
+            organization_id: Organization ID
+            api_token: API token for authentication
+        """
+        self.db = db
+        self.organization_id = organization_id
+        self.api_token = api_token
+
+    @abstractmethod
+    async def plan_task(self, planning_prompt: str) -> TaskPlanResponse:
+        """
+        Generate a task plan (non-streaming).
+
+        Args:
+            planning_prompt: The complete planning prompt
+
+        Returns:
+            TaskPlanResponse with the generated plan
+        """
+        pass
+
+    @abstractmethod
+    async def plan_task_stream(self, planning_prompt: str) -> AsyncIterator[Dict[str, Any]]:
+        """
+        Generate a task plan with streaming events.
+
+        Args:
+            planning_prompt: The complete planning prompt
+
+        Yields:
+            Dict with event type and data for SSE streaming
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Return the strategy name for logging"""
+        pass

control_plane_api/app/services/planning_strategy_factory.py

@@ -0,0 +1,118 @@
+"""
+Planning Strategy Factory
+
+Factory for creating the appropriate planning strategy based on configuration.
+Like a travel agency - chooses the best transportation mode for you!
+"""
+
+import os
+import structlog
+from typing import Literal
+from sqlalchemy.orm import Session
+
+from control_plane_api.app.services.planning_strategy import PlanningStrategy
+from control_plane_api.app.services.claude_code_planning_service import ClaudeCodePlanningStrategy
+from control_plane_api.app.services.agno_planning_strategy import AgnoPlanningStrategy
+
+logger = structlog.get_logger(__name__)
+
+# Supported strategy types
+StrategyType = Literal["claude_code_sdk", "agno"]
+
+
+def get_planning_strategy(
+    strategy_type: StrategyType = None,
+    db: Session = None,
+    organization_id: str = None,
+    api_token: str = None
+) -> PlanningStrategy:
+    """
+    Factory function to create the appropriate planning strategy.
+
+    Args:
+        strategy_type: Which strategy to use ("claude_code_sdk" or "agno")
+                      If None, uses PLANNING_STRATEGY env var (defaults to "agno")
+        db: Database session
+        organization_id: Organization ID
+        api_token: API token
+
+    Returns:
+        Concrete PlanningStrategy implementation
+
+    Example:
+        # Use Claude Code SDK (default)
+        strategy = get_planning_strategy(db=db)
+
+        # Explicitly use Agno
+        strategy = get_planning_strategy(strategy_type="agno", db=db)
+
+        # Use both in your code
+        plan = await strategy.plan_task(prompt)  # Same interface!
+    """
+    # Determine strategy from parameter or env var
+    if strategy_type is None:
+        strategy_type = os.getenv("PLANNING_STRATEGY", "agno").lower()
+
+    logger.info(
+        "creating_planning_strategy",
+        strategy_type=strategy_type,
+        has_db=db is not None,
+        has_org_id=organization_id is not None,
+    )
+
+    # Validate strategy availability and detect CLI
+    if strategy_type == "claude_code_sdk":
+        # Check if SDK is available
+        try:
+            from claude_agent_sdk import ClaudeSDKClient
+        except ImportError:
+            logger.error(
+                "strategy_unavailable",
+                strategy_type=strategy_type,
+                message="claude-agent-sdk not available, falling back to Agno"
+            )
+            strategy_type = "agno"  # Fallback
+        else:
+            # Check if Claude CLI binary is available
+            import shutil
+            claude_cli = shutil.which("claude")
+            if not claude_cli:
+                logger.warning(
+                    "claude_cli_not_found",
+                    message="Claude Code CLI binary not found in PATH. Falling back to Agno.",
+                    hint="Install @anthropic-ai/claude-code npm package or use PLANNING_STRATEGY=agno"
+                )
+                strategy_type = "agno"  # Auto-fallback
+            else:
+                logger.info(
+                    "claude_cli_found",
+                    cli_path=claude_cli,
+                    message="Claude Code CLI available, using claude_code_sdk strategy"
+                )
+
+    # Create the appropriate strategy
+    if strategy_type == "claude_code_sdk":
+        return ClaudeCodePlanningStrategy(
+            db=db,
+            organization_id=organization_id,
+            api_token=api_token
+        )
+    elif strategy_type == "agno":
+        return AgnoPlanningStrategy(
+            db=db,
+            organization_id=organization_id,
+            api_token=api_token
+        )
+    else:
+        logger.error("unknown_strategy_type", strategy_type=strategy_type)
+        # Default to Claude Code SDK
+        logger.warning("defaulting_to_claude_code_sdk")
+        return ClaudeCodePlanningStrategy(
+            db=db,
+            organization_id=organization_id,
+            api_token=api_token
+        )
+
+
+# Convenience alias for backward compatibility
+get_claude_code_planning_service = get_planning_strategy