sky-dev 0.0.5__py3-none-any.whl

sky/core/router.py

@@ -0,0 +1,264 @@
+"""Model Router with F11 Batching Verification."""
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import httpx
+from openai import AsyncOpenAI
+
+from sky.config import ModelAssignmentConfig, ModelRoutingConfig
+from sky.storage import DatabaseManager
+from sky.errors import ProviderError, ConfigurationError
+
+class ModelRouter:
+    """Routes requests to appropriate LLMs with F11 verification."""
+
+    def __init__(self, config: ModelRoutingConfig, db: DatabaseManager, session_id: str):
+        import os
+        self.config = config
+        self.db = db
+        self.session_id = session_id
+        
+        self.clients = {}
+        self.provider_keys = {}
+        self.current_key_idx = {}
+        
+        for provider_name, provider_config in self.config.providers.items():
+            if provider_config.requires_api_key:
+                # Find keys for this provider, e.g., GROQ_API_KEY, NIM_API_KEY
+                env_prefix = provider_name.upper()
+                keys = []
+                for key_name in [f"{env_prefix}_API_KEY", f"{env_prefix}_API_KEY_2", f"{env_prefix}_API_KEY_3"]:
+                    key = os.getenv(key_name)
+                    if key:
+                        keys.append(key)
+                if not keys:
+                    # Log a warning, but don't crash unless they try to use it
+                    pass
+                else:
+                    self.provider_keys[provider_name] = keys
+                    self.current_key_idx[provider_name] = 0
+            else:
+                self.provider_keys[provider_name] = ["dummy"]
+                self.current_key_idx[provider_name] = 0
+                
+            self._init_client(provider_name)
+        
+    def _init_client(self, provider_name: str):
+        """Initialize the client for a specific provider."""
+        provider_config = self.config.providers.get(provider_name)
+        if not provider_config:
+            return
+            
+        keys = self.provider_keys.get(provider_name, [])
+        if not keys:
+            return
+            
+        idx = self.current_key_idx.get(provider_name, 0)
+        api_key = keys[idx]
+        
+        kwargs = {
+            "api_key": api_key,
+            "timeout": httpx.Timeout(provider_config.timeout)
+        }
+        if provider_config.base_url:
+            kwargs["base_url"] = provider_config.base_url
+            
+        self.clients[provider_name] = AsyncOpenAI(**kwargs)
+
+    def get_model_for_role(self, role: str) -> ModelAssignmentConfig:
+        """Get model config for a specific role."""
+        if role in self.config.roles:
+            return self.config.roles[role]
+        return self.get_fallback_model()
+
+    def get_fallback_model(self) -> ModelAssignmentConfig:
+        """Get the fallback model config."""
+        if self.config.fallback:
+            return self.config.fallback
+        return ModelAssignmentConfig(provider="nim", model_id="meta/llama-3.1-8b-instruct", reasoning_effort="none")
+
+    def build_completion_params(
+        self, 
+        assignment: ModelAssignmentConfig, 
+        messages: List[Dict[str, Any]], 
+        tools: Optional[List[Dict[str, Any]]], 
+        tool_choice: str, 
+        parallel_tool_calls: bool
+    ) -> Dict[str, Any]:
+        """Build parameters for Groq API call."""
+        params: Dict[str, Any] = {
+            "model": assignment.model_id,
+            "messages": messages,
+        }
+        if assignment.temperature is not None:
+            params["temperature"] = assignment.temperature
+        if assignment.max_tokens is not None:
+            params["max_tokens"] = assignment.max_tokens
+            
+        if tools:
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+            params["parallel_tool_calls"] = parallel_tool_calls
+            
+        return params
+
+    async def _call_with_rotation(self, params: Dict[str, Any], provider_name: str) -> Any:
+        """Call API, rotating keys if rate limit is hit."""
+        from rich.console import Console
+        
+        client = self.clients.get(provider_name)
+        if not client:
+            raise ConfigurationError(
+                f"No active client for provider: {provider_name}.",
+                code="SKY-001",
+                suggestion=f"Add {provider_name.upper()}_API_KEY to .env or run `sky init`"
+            )
+            
+        keys = self.provider_keys.get(provider_name, [])
+        last_error = None
+        
+        for _ in range(max(1, len(keys))):
+            try:
+                return await client.chat.completions.create(**params)
+            except Exception as e:
+                error_str = str(e).lower()
+                if any(err in error_str for err in ["413", "429", "rate limit", "rate_limit_exceeded"]):
+                    last_error = e
+                    if len(keys) > 1:
+                        next_idx = (self.current_key_idx[provider_name] + 1) % len(keys)
+                        Console().print(f"[yellow]{provider_name.upper()} API key rate limited. Rotating to next key (slot {next_idx + 1}/{len(keys)})...[/yellow]")
+                        self.current_key_idx[provider_name] = next_idx
+                        self._init_client(provider_name)
+                        client = self.clients[provider_name]
+                    else:
+                        break
+                else:
+                    raise ProviderError(
+                        f"Provider {provider_name} request failed: {e}",
+                        code="SKY-002",
+                        suggestion="Check your internet connection or run `sky check-providers` to verify availability."
+                    )
+                    
+        raise ProviderError(
+            f"Rate limit exceeded after exhausting all {len(keys)} keys.",
+            code="SKY-003",
+            suggestion="Wait 60 seconds. Consider using a different model or provider."
+        )
+
+    async def route(
+        self, 
+        role: str, 
+        messages: List[Dict[str, Any]], 
+        tools: Optional[List[Dict[str, Any]]] = None, 
+        expected_tool_calls: Optional[int] = None, 
+        tool_choice: str = "auto"
+    ) -> Tuple[Dict[str, Any], bool, str]:
+        """Route request to model, returning (response_message_dict, was_fallback, model_used)."""
+        if expected_tool_calls and expected_tool_calls > 1 and tools:
+            resp, was_fallback = await self.verify_and_dispatch(role, messages, tools, expected_tool_calls, tool_choice)
+            model_used = self.get_fallback_model().model_id if was_fallback else self.get_model_for_role(role).model_id
+            return resp, was_fallback, model_used
+            
+        assignment = self.get_model_for_role(role)
+        params = self.build_completion_params(assignment, messages, tools, tool_choice, self.config.parallel_tool_calls)
+        
+        chat_completion = await self._call_with_rotation(params, assignment.provider)
+        resp = chat_completion.choices[0].message
+        
+        if chat_completion.usage:
+            cost = (chat_completion.usage.prompt_tokens + chat_completion.usage.completion_tokens) * 0.000001
+            self.db.log_usage(
+                self.session_id, 
+                assignment.model_id, 
+                role, 
+                chat_completion.usage.prompt_tokens, 
+                chat_completion.usage.completion_tokens, 
+                cost, 
+                0.0
+            )
+        
+        resp_dict: Dict[str, Any] = {
+            "role": "assistant",
+            "content": resp.content,
+        }
+        if resp.tool_calls:
+            resp_dict["tool_calls"] = [
+                {
+                    "id": tc.id,
+                    "type": "function",
+                    "function": {
+                        "name": tc.function.name,
+                        "arguments": tc.function.arguments
+                    }
+                } for tc in resp.tool_calls
+            ]
+        return resp_dict, False, assignment.model_id
+
+    async def verify_and_dispatch(
+        self, 
+        role: str, 
+        messages: List[Dict[str, Any]], 
+        tools: List[Dict[str, Any]], 
+        expected_tool_calls: int = 1, 
+        tool_choice: str = "auto"
+    ) -> Tuple[Dict[str, Any], bool]:
+        """F11 PRIMARY FUNCTION: Verify tool batching count and fallback if under-calling."""
+        assignment = self.get_model_for_role(role)
+        params = self.build_completion_params(assignment, messages, tools, tool_choice, True)
+        
+        chat_completion = await self._call_with_rotation(params, assignment.provider)
+        resp = chat_completion.choices[0].message
+        
+        tool_call_count = len(resp.tool_calls) if resp.tool_calls else 0
+        
+        if tool_call_count < expected_tool_calls:
+            self._log_fallback_event(assignment.model_id, tool_call_count, expected_tool_calls, "under_calling")
+            
+            fallback = self.get_fallback_model()
+            fb_params = self.build_completion_params(fallback, messages, tools, tool_choice, True)
+            fb_completion = await self._call_with_rotation(fb_params, fallback.provider)
+            fb_resp = fb_completion.choices[0].message
+            
+            resp_dict: Dict[str, Any] = {
+                "role": "assistant",
+                "content": fb_resp.content,
+            }
+            if fb_resp.tool_calls:
+                resp_dict["tool_calls"] = [
+                    {
+                        "id": tc.id,
+                        "type": "function",
+                        "function": {
+                            "name": tc.function.name,
+                            "arguments": tc.function.arguments
+                        }
+                    } for tc in fb_resp.tool_calls
+                ]
+            return resp_dict, True
+            
+        resp_dict = {
+            "role": "assistant",
+            "content": resp.content,
+        }
+        if resp.tool_calls:
+            resp_dict["tool_calls"] = [
+                {
+                    "id": tc.id,
+                    "type": "function",
+                    "function": {
+                        "name": tc.function.name,
+                        "arguments": tc.function.arguments
+                    }
+                } for tc in resp.tool_calls
+            ]
+        return resp_dict, False
+
+    def _log_fallback_event(self, model: str, actual: int, expected: int, reason: str) -> None:
+        """Log fallback event to audit log."""
+        self.db._write_audit_log(self.session_id, {
+            "event": "fallback_triggered",
+            "model": model,
+            "actual_tool_calls": actual,
+            "expected_tool_calls": expected,
+            "reason": reason
+        })

sky/core/subagent.py

@@ -0,0 +1,302 @@
+"""Subagent system for SKY workflow engine."""
+
+from enum import Enum
+from typing import Dict, List, Any, Optional, Literal
+from dataclasses import dataclass, field
+import logging
+import json
+
+from sky.config.schema import DexProjectConfig
+from sky.storage.db import DatabaseManager
+from sky.core.router import ModelRouter
+from sky.core.approval import ApprovalGate
+from sky.core.fast_loop import FastLoopEngine
+from sky.memory.indexer import RepoIndexer
+
+logger = logging.getLogger(__name__)
+
+
+class SubagentRole(str, Enum):
+    """Roles for subagents."""
+    PLANNER = "planner"
+    CODER = "coder"
+    TESTER = "tester"
+    REVIEWER = "reviewer"
+
+
+@dataclass
+class SubagentConfig:
+    """Configuration for a subagent."""
+    role: SubagentRole
+    model_role: str  # Corresponds to models.yaml role
+    tools: List[str] = field(default_factory=list)
+    system_prompt_template: str = ""
+    max_turns: int = 10
+    temperature: float = 0.3
+
+
+# Role-specific configurations
+ROLE_CONFIGS = {
+    SubagentRole.PLANNER: SubagentConfig(
+        role=SubagentRole.PLANNER,
+        model_role="planning",
+        tools=["read_file", "grep", "glob", "git_diff", "git_log"],
+        system_prompt_template="""You are a planning subagent for SKY.
+
+Your role is to analyze the codebase and create a detailed plan.
+
+Available tools: read_file, grep, glob, git_diff, git_log
+
+Output a plan with:
+1. What needs to be done
+2. What files will be affected
+3. Potential risks or dependencies
+4. Estimated effort
+
+Be thorough and specific.""",
+        max_turns=5,
+        temperature=0.3,
+    ),
+    SubagentRole.CODER: SubagentConfig(
+        role=SubagentRole.CODER,
+        model_role="fast_loop",
+        tools=["read_file", "write_file", "edit_file", "grep", "glob", "git_diff", "git_commit", "git_branch"],
+        system_prompt_template="""You are a coding subagent for SKY.
+
+Your role is to implement code changes.
+
+Available tools: read_file, write_file, edit_file, grep, glob, git_diff, git_commit, git_branch
+
+Guidelines:
+1. Read relevant files first
+2. Make targeted changes
+3. Use git_commit to commit changes
+4. Explain what you changed and why
+
+Be precise and focused.""",
+        max_turns=10,
+        temperature=0.1,
+    ),
+    SubagentRole.TESTER: SubagentConfig(
+        role=SubagentRole.TESTER,
+        model_role="fast_loop",
+        tools=["read_file", "run_tests", "lint", "grep", "glob"],
+        system_prompt_template="""You are a testing subagent for SKY.
+
+Your role is to run tests and report results.
+
+Available tools: read_file, run_tests, lint, grep, glob
+
+Guidelines:
+1. Run tests after code changes
+2. Report pass/fail results
+3. If tests fail, provide error details
+4. Suggest fixes for failing tests
+
+Be thorough and accurate.""",
+        max_turns=5,
+        temperature=0.1,
+    ),
+    SubagentRole.REVIEWER: SubagentConfig(
+        role=SubagentRole.REVIEWER,
+        model_role="planning",
+        tools=["read_file", "grep", "glob", "git_diff", "git_log"],
+        system_prompt_template="""You are a code reviewer subagent for SKY.
+
+Your role is to review changes and provide feedback.
+
+Available tools: read_file, grep, glob, git_diff, git_log
+
+Guidelines:
+1. Review the diff carefully
+2. Check for code quality issues
+3. Look for potential bugs or edge cases
+4. Provide constructive feedback
+5. Suggest improvements
+
+Be critical but constructive.""",
+        max_turns=5,
+        temperature=0.3,
+    ),
+}
+
+
+class SubagentRunner:
+    """Runner for subagent execution."""
+    
+    def __init__(
+        self,
+        config: DexProjectConfig,
+        db: DatabaseManager,
+        router: ModelRouter,
+        approval_gate: ApprovalGate,
+        indexer: Optional[RepoIndexer] = None,
+    ):
+        self.config = config
+        self.db = db
+        self.router = router
+        self.approval_gate = approval_gate
+        self.indexer = indexer
+    
+    async def run(
+        self,
+        task: Dict[str, Any],
+        role: SubagentRole,
+        parent_state: Optional[Dict[str, Any]] = None,
+        event_queue: Optional[Any] = None,
+    ) -> Dict[str, Any]:
+        """Run a subagent for a specific task."""
+        logger.info(f"Running subagent {role} for task {task.get('id', 'unknown')}")
+        
+        if event_queue:
+            import asyncio
+            if isinstance(event_queue, asyncio.Queue):
+                await event_queue.put({
+                    "type": "subagent_start", 
+                    "role": role.value, 
+                    "task": task.get("description", "Unknown task")
+                })
+        
+        # Get role configuration
+        role_config = ROLE_CONFIGS.get(role)
+        if not role_config:
+            raise ValueError(f"Unknown role: {role}")
+        
+        # Build system prompt
+        system_prompt = role_config.system_prompt_template
+        
+        # Add task-specific context
+        if parent_state:
+            context = self._build_context(parent_state)
+            system_prompt = f"{system_prompt}\n\nContext from parent:\n{context}"
+        
+        # Build messages
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"Task: {task.get('description', '')}\n\nDependencies: {task.get('dependencies', [])}\n\nRequired tools: {task.get('tools_needed', [])}"}
+        ]
+        
+        # Create a fresh session for subagent
+        session_id = self.db.create_session(f"subagent_{role.value}", task.get("id", "unknown"))
+        
+        # Initialize FastLoopEngine for this subagent
+        engine = FastLoopEngine(
+            config=self.config,
+            db=self.db,
+            router=self.router,
+            approval_gate=self.approval_gate,
+            session_id=session_id,
+            indexer=self.indexer,
+        )
+        
+        # Run the fast loop
+        tool_filter = role_config.tools if role_config.tools else None
+        
+        responses = []
+        try:
+            async for event in engine.run(
+                messages=messages,
+                mode="agent",
+                max_turns=role_config.max_turns,
+                tools_filter=tool_filter,
+                inject_context=False,  # Subagents don't auto-inject context
+            ):
+                if event_queue:
+                    import asyncio
+                    if isinstance(event_queue, asyncio.Queue):
+                        if event.get("type") in ("tool_call", "tool_result", "error"):
+                            await event_queue.put(event)
+                            
+                if event.get("type") == "final_answer":
+                    responses.append(event.get("content", ""))
+                elif event.get("type") == "error":
+                    error_msg = event.get("content", event.get("error", "Unknown error"))
+                    if event_queue:
+                        import asyncio
+                        if isinstance(event_queue, asyncio.Queue):
+                            await event_queue.put({
+                                "type": "subagent_complete", 
+                                "role": role.value, 
+                                "summary": f"Error: {error_msg}"
+                            })
+                    return {
+                        "task_id": task.get("id"),
+                        "success": False,
+                        "error": error_msg,
+                        "summary": "",
+                    }
+        except Exception as e:
+            logger.error(f"Subagent {role} failed: {e}")
+            if event_queue:
+                import asyncio
+                if isinstance(event_queue, asyncio.Queue):
+                    await event_queue.put({
+                        "type": "subagent_complete", 
+                        "role": role.value, 
+                        "summary": f"Failed: {str(e)}"
+                    })
+            return {
+                "task_id": task.get("id"),
+                "success": False,
+                "error": str(e),
+                "summary": "",
+            }
+        
+        # Compress results
+        summary = self._summarize_results(responses, role)
+        
+        if event_queue:
+            import asyncio
+            if isinstance(event_queue, asyncio.Queue):
+                await event_queue.put({
+                    "type": "subagent_complete", 
+                    "role": role.value, 
+                    "summary": summary
+                })
+        
+        return {
+            "task_id": task.get("id"),
+            "role": role.value,
+            "success": True,
+            "responses": responses,
+            "summary": summary,
+            "session_id": session_id,
+            "tool_usage": engine._tool_usage if hasattr(engine, "_tool_usage") else [],
+        }
+    
+    def _build_context(self, parent_state: Dict[str, Any]) -> str:
+        """Build context from parent state."""
+        context_parts = []
+        
+        if parent_state.get("refined_goal"):
+            context_parts.append(f"Goal: {parent_state['refined_goal']}")
+        
+        if parent_state.get("task_results"):
+            completed = [
+                f"{tid}: {res.get('summary', 'completed')}"
+                for tid, res in parent_state["task_results"].items()
+                if res.get("success")
+            ]
+            if completed:
+                context_parts.append(f"Completed tasks:\n- " + "\n- ".join(completed))
+        
+        return "\n".join(context_parts) if context_parts else "No additional context"
+    
+    def _summarize_results(self, responses: List[str], role: SubagentRole) -> str:
+        """Summarize subagent results."""
+        if not responses:
+            return "No output from subagent"
+        
+        # For testers, summarize test results
+        if role == SubagentRole.TESTER:
+            # Look for test results
+            for response in responses:
+                if "passed" in response.lower() or "failed" in response.lower():
+                    # Extract test summary
+                    lines = response.split("\n")
+                    test_lines = [l for l in lines if "pass" in l.lower() or "fail" in l.lower() or "test" in l.lower()]
+                    if test_lines:
+                        return "\n".join(test_lines[:5])
+        
+        # Default: return first response truncated
+        return responses[-1][:500] + ("..." if len(responses[-1]) > 500 else "")