codeoptix 0.1.0__py3-none-any.whl

codeoptix/acp/orchestrator.py

@@ -0,0 +1,349 @@
+"""Agent Orchestration - Route to best agent for each task."""
+
+import logging
+from typing import Any
+
+from acp import text_block
+
+from codeoptix.acp.code_extractor import extract_code_from_message, extract_code_from_text
+from codeoptix.acp.registry import ACPAgentRegistry
+from codeoptix.evaluation import EvaluationEngine
+from codeoptix.utils.llm import LLMClient
+
+logger = logging.getLogger(__name__)
+
+
+class AgentOrchestrator:
+    """Orchestrates multiple ACP agents for task execution."""
+
+    def __init__(
+        self,
+        registry: ACPAgentRegistry,
+        evaluation_engine: EvaluationEngine | None = None,
+        llm_client: LLMClient | None = None,
+    ):
+        """Initialize agent orchestrator.
+
+        Args:
+            registry: ACP agent registry
+            evaluation_engine: Optional evaluation engine
+            llm_client: Optional LLM client
+        """
+        self.registry = registry
+        self.evaluation_engine = evaluation_engine
+        self.llm_client = llm_client
+
+    async def route_to_agent(
+        self,
+        prompt: str,
+        agent_name: str | None = None,
+        context: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Route a prompt to the best agent.
+
+        Args:
+            prompt: The prompt to execute
+            agent_name: Specific agent name (if None, selects best agent)
+            context: Additional context
+
+        Returns:
+            Result dictionary with agent response and metadata
+        """
+        # Select agent
+        if agent_name:
+            selected_agent = agent_name
+        else:
+            selected_agent = await self._select_best_agent(prompt, context)
+
+        if not selected_agent:
+            raise ValueError("No agent available")
+
+        # Connect to agent
+        connection = await self.registry.connect(selected_agent)
+        session_id = self.registry.get_session_id(selected_agent)
+
+        if not session_id:
+            raise RuntimeError(f"No session ID for agent {selected_agent}")
+
+        # Send prompt
+        response = await connection.prompt(
+            session_id=session_id,
+            prompt=[text_block(prompt)],
+        )
+
+        return {
+            "agent": selected_agent,
+            "response": response,
+            "session_id": session_id,
+        }
+
+    async def _select_best_agent(
+        self,
+        prompt: str,
+        context: dict[str, Any] | None = None,
+    ) -> str | None:
+        """Select the best agent for a given prompt.
+
+        Args:
+            prompt: The prompt
+            context: Additional context
+
+        Returns:
+            Agent name or None if no agents available
+        """
+        agents = self.registry.list_agents()
+        if not agents:
+            return None
+
+        # Intelligent agent selection based on:
+        # - Agent capabilities (from registry)
+        # - Task type (inferred from prompt)
+        # - Context requirements
+
+        # Check if context specifies an agent
+        if context and "preferred_agent" in context:
+            preferred = context["preferred_agent"]
+            if preferred in agents:
+                return preferred
+
+        # Infer task type from prompt
+        prompt_lower = prompt.lower()
+
+        # Security-focused tasks
+        if any(
+            keyword in prompt_lower
+            for keyword in ["security", "secure", "vulnerability", "exploit", "attack"]
+        ):
+            # Prefer agents with security capabilities
+            for agent_name in agents:
+                agent_config = self.registry.get_agent(agent_name)
+                if agent_config and "security" in [c.lower() for c in agent_config.capabilities]:
+                    return agent_name
+
+        # Code review tasks
+        if any(keyword in prompt_lower for keyword in ["review", "critique", "judge", "evaluate"]):
+            # Prefer agents with review capabilities
+            for agent_name in agents:
+                agent_config = self.registry.get_agent(agent_name)
+                if agent_config and "review" in [c.lower() for c in agent_config.capabilities]:
+                    return agent_name
+
+        # Default: use first available agent
+        return agents[0]
+
+    async def execute_multi_agent_workflow(
+        self,
+        workflow: list[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
+        """Execute a multi-agent workflow.
+
+        Args:
+            workflow: List of workflow steps, each with 'agent', 'prompt', etc.
+
+        Returns:
+            List of results from each step
+        """
+        results = []
+        for step in workflow:
+            agent_name = step.get("agent")
+            prompt = step.get("prompt", "")
+            context = step.get("context", {})
+
+            result = await self.route_to_agent(
+                prompt=prompt,
+                agent_name=agent_name,
+                context=context,
+            )
+            results.append(result)
+
+        return results
+
+
+class MultiAgentJudge:
+    """Multi-agent judge - Use different agents for generation vs. judgment."""
+
+    def __init__(
+        self,
+        registry: ACPAgentRegistry,
+        generate_agent: str,
+        judge_agent: str,
+        evaluation_engine: EvaluationEngine | None = None,
+        llm_client: LLMClient | None = None,
+    ):
+        """Initialize multi-agent judge.
+
+        Args:
+            registry: ACP agent registry
+            generate_agent: Name of agent for code generation
+            judge_agent: Name of agent for code judgment/critique
+            evaluation_engine: Optional evaluation engine
+            llm_client: Optional LLM client
+        """
+        self.registry = registry
+        self.generate_agent = generate_agent
+        self.judge_agent = judge_agent
+        self.evaluation_engine = evaluation_engine
+        self.llm_client = llm_client
+
+    async def generate_and_judge(
+        self,
+        prompt: str,
+        context: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Generate code with one agent and judge with another.
+
+        Args:
+            prompt: The prompt for code generation
+            context: Additional context
+
+        Returns:
+            Dictionary with generated code, judgment, and evaluation results
+        """
+        # Step 1: Generate code with generate agent
+        generate_conn = await self.registry.connect(self.generate_agent)
+        generate_session = self.registry.get_session_id(self.generate_agent)
+
+        if not generate_session:
+            raise RuntimeError(f"No session for generate agent {self.generate_agent}")
+
+        generate_response = await generate_conn.prompt(
+            session_id=generate_session,
+            prompt=[text_block(prompt)],
+        )
+
+        # Extract generated code (simplified - would need proper extraction)
+        generated_code = self._extract_code_from_response(generate_response)
+
+        # Step 2: Judge code with judge agent
+        judge_prompt = f"""Please review and critique the following code:
+
+```python
+{generated_code}
+```
+
+Provide a detailed critique focusing on:
+- Code quality and best practices
+- Potential bugs or issues
+- Security concerns
+- Performance considerations
+- Suggestions for improvement
+"""
+
+        judge_conn = await self.registry.connect(self.judge_agent)
+        judge_session = self.registry.get_session_id(self.judge_agent)
+
+        if not judge_session:
+            raise RuntimeError(f"No session for judge agent {self.judge_agent}")
+
+        judge_response = await judge_conn.prompt(
+            session_id=judge_session,
+            prompt=[text_block(judge_prompt)],
+        )
+
+        # Extract judgment
+        judgment = self._extract_text_from_response(judge_response)
+
+        # Step 3: Evaluate both with CodeOptiX
+        evaluation_results = None
+        if self.evaluation_engine:
+            from codeoptix.adapters.base import AgentOutput
+
+            AgentOutput(
+                code=generated_code,
+                tests="",
+                messages=[],
+                metadata={"source": "multi_agent_judge", "judgment": judgment},
+            )
+
+            # Evaluate behaviors
+            evaluation_results = await self.evaluation_engine.evaluate_behaviors(
+                behavior_names=["insecure-code", "vacuous-tests", "plan-drift"],
+                context={"code": generated_code, "judgment": judgment},
+            )
+
+        return {
+            "generated_code": generated_code,
+            "judgment": judgment,
+            "evaluation_results": evaluation_results,
+            "generate_agent": self.generate_agent,
+            "judge_agent": self.judge_agent,
+        }
+
+    def _extract_code_from_response(self, response: Any) -> str:
+        """Extract code from agent response.
+
+        Args:
+            response: ACP prompt response
+
+        Returns:
+            Extracted code as string
+        """
+        if not response:
+            return ""
+
+        # Extract from response messages
+        code_blocks = []
+
+        # Check if response has messages
+        if hasattr(response, "messages"):
+            for message in response.messages:
+                if hasattr(message, "content"):
+                    content = message.content
+                    if isinstance(content, str):
+                        code_blocks.extend(extract_code_from_text(content))
+                    elif hasattr(content, "text"):
+                        code_blocks.extend(extract_code_from_text(getattr(content, "text", "")))
+
+        # Check if response has updates
+        if hasattr(response, "updates"):
+            for update in response.updates:
+                code_blocks.extend(extract_code_from_message(update))
+
+        # Combine all code blocks
+        if code_blocks:
+            # Prefer code blocks over inline code
+            block_codes = [cb["content"] for cb in code_blocks if cb.get("type") == "block"]
+            if block_codes:
+                return "\n\n".join(block_codes)
+            # Fallback to inline code
+            inline_codes = [cb["content"] for cb in code_blocks if cb.get("type") == "inline"]
+            if inline_codes:
+                return "\n".join(inline_codes)
+
+        return ""
+
+    def _extract_text_from_response(self, response: Any) -> str:
+        """Extract text from agent response.
+
+        Args:
+            response: ACP prompt response
+
+        Returns:
+            Extracted text as string
+        """
+        if not response:
+            return ""
+
+        text_parts = []
+
+        # Check if response has messages
+        if hasattr(response, "messages"):
+            for message in response.messages:
+                if hasattr(message, "content"):
+                    content = message.content
+                    if isinstance(content, str):
+                        text_parts.append(content)
+                    elif hasattr(content, "text"):
+                        text_parts.append(getattr(content, "text", ""))
+
+        # Check if response has updates
+        if hasattr(response, "updates"):
+            for update in response.updates:
+                if hasattr(update, "content"):
+                    content = update.content
+                    if isinstance(content, str):
+                        text_parts.append(content)
+                    elif hasattr(content, "text"):
+                        text_parts.append(getattr(content, "text", ""))
+
+        return "\n".join(text_parts)

codeoptix/acp/registry.py

@@ -0,0 +1,294 @@
+"""ACP Agent Registry - Manage and connect to multiple ACP-compatible agents."""
+
+import asyncio
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+from acp import PROTOCOL_VERSION, Client, connect_to_agent
+from acp.core import ClientSideConnection
+from acp.schema import (
+    ClientCapabilities,
+    CreateTerminalResponse,
+    Implementation,
+    ReadTextFileResponse,
+    RequestPermissionResponse,
+    WriteTextFileResponse,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ACPAgentConfig:
+    """Configuration for an ACP agent."""
+
+    name: str
+    command: list[str]
+    cwd: str | None = None
+    env: dict[str, str] = field(default_factory=dict)
+    description: str = ""
+    capabilities: list[str] = field(default_factory=list)
+
+
+class ACPAgentRegistry:
+    """Registry for managing ACP-compatible agents."""
+
+    def __init__(self):
+        """Initialize the agent registry."""
+        self._agents: dict[str, ACPAgentConfig] = {}
+        self._connections: dict[str, ClientSideConnection] = {}
+        self._sessions: dict[str, str] = {}  # agent_name -> session_id
+
+    def register(
+        self,
+        name: str,
+        command: list[str],
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        description: str = "",
+        capabilities: list[str] | None = None,
+    ) -> None:
+        """Register an ACP agent.
+
+        Args:
+            name: Unique name for the agent
+            command: Command to spawn the agent (e.g., ["python", "agent.py"])
+            cwd: Working directory for the agent
+            env: Environment variables
+            description: Description of the agent
+            capabilities: List of agent capabilities
+        """
+        config = ACPAgentConfig(
+            name=name,
+            command=command,
+            cwd=cwd,
+            env=env or {},
+            description=description,
+            capabilities=capabilities or [],
+        )
+        self._agents[name] = config
+        logger.info(f"Registered ACP agent: {name}")
+
+    def unregister(self, name: str) -> None:
+        """Unregister an ACP agent.
+
+        Args:
+            name: Name of the agent to unregister
+        """
+        if name in self._agents:
+            # Close connection if open
+            if name in self._connections:
+                self._connections.pop(name)
+            if name in self._sessions:
+                self._sessions.pop(name)
+            del self._agents[name]
+            logger.info(f"Unregistered ACP agent: {name}")
+
+    def list_agents(self) -> list[str]:
+        """List all registered agent names."""
+        return list(self._agents.keys())
+
+    def get_agent(self, name: str) -> ACPAgentConfig | None:
+        """Get agent configuration.
+
+        Args:
+            name: Agent name
+
+        Returns:
+            Agent configuration or None if not found
+        """
+        return self._agents.get(name)
+
+    async def connect(self, name: str) -> ClientSideConnection:
+        """Connect to a registered ACP agent.
+
+        Args:
+            name: Agent name
+
+        Returns:
+            ClientSideConnection to the agent
+
+        Raises:
+            ValueError: If agent not found
+            RuntimeError: If connection fails
+        """
+        if name not in self._agents:
+            raise ValueError(f"Agent '{name}' not found in registry")
+
+        # Return existing connection if available
+        if name in self._connections:
+            return self._connections[name]
+
+        config = self._agents[name]
+
+        # Spawn agent process
+        process = await asyncio.create_subprocess_exec(
+            *config.command,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            cwd=config.cwd,
+            env={**config.env, **dict(asyncio.get_event_loop().get_environ())},
+        )
+
+        if process.stdin is None or process.stdout is None:
+            raise RuntimeError("Agent process does not expose stdio pipes")
+
+        # Create client implementation
+        client_impl = _RegistryClientImpl()
+        connection = connect_to_agent(client_impl, process.stdin, process.stdout)
+
+        # Initialize connection
+        await connection.initialize(
+            protocol_version=PROTOCOL_VERSION,
+            client_capabilities=ClientCapabilities(),
+            client_info=Implementation(
+                name="codeoptix-registry",
+                title="CodeOptiX Agent Registry",
+                version="0.1.0",
+            ),
+        )
+
+        # Create new session
+        session = await connection.new_session(mcp_servers=[], cwd=config.cwd or ".")
+        self._sessions[name] = session.session_id
+
+        # Store connection
+        self._connections[name] = connection
+
+        logger.info(f"Connected to ACP agent: {name}")
+        return connection
+
+    async def disconnect(self, name: str) -> None:
+        """Disconnect from an agent.
+
+        Args:
+            name: Agent name
+        """
+        if name in self._connections:
+            connection = self._connections[name]
+            try:
+                # Try to gracefully close the connection
+                if hasattr(connection, "close"):
+                    await connection.close()
+            except Exception as e:
+                logger.warning(f"Error closing connection to agent {name}: {e}")
+            finally:
+                self._connections.pop(name)
+        if name in self._sessions:
+            self._sessions.pop(name)
+        logger.info(f"Disconnected from ACP agent: {name}")
+
+    def get_session_id(self, name: str) -> str | None:
+        """Get session ID for an agent.
+
+        Args:
+            name: Agent name
+
+        Returns:
+            Session ID or None if not connected
+        """
+        return self._sessions.get(name)
+
+
+class _RegistryClientImpl(Client):
+    """Internal client implementation for registry connections."""
+
+    async def request_permission(
+        self,
+        options: list,
+        session_id: str,
+        tool_call: Any,
+        **kwargs: Any,
+    ) -> RequestPermissionResponse:
+        """Handle permission requests."""
+        return RequestPermissionResponse(granted=True)
+
+    async def write_text_file(
+        self,
+        content: str,
+        path: str,
+        session_id: str,
+        **kwargs: Any,
+    ) -> WriteTextFileResponse | None:
+        """Handle file write requests."""
+        return WriteTextFileResponse()
+
+    async def read_text_file(
+        self,
+        path: str,
+        session_id: str,
+        limit: int | None = None,
+        line: int | None = None,
+        **kwargs: Any,
+    ) -> ReadTextFileResponse:
+        """Handle file read requests."""
+        try:
+            with open(path, encoding="utf-8") as f:
+                if line is not None:
+                    lines = f.readlines()
+                    if 0 <= line < len(lines):
+                        content = lines[line]
+                    else:
+                        content = ""
+                elif limit is not None:
+                    content = f.read(limit)
+                else:
+                    content = f.read()
+            return ReadTextFileResponse(content=content)
+        except Exception as e:
+            logger.error(f"Error reading file {path}: {e}")
+            raise
+
+    async def create_terminal(self, *args: Any, **kwargs: Any) -> CreateTerminalResponse:
+        """Handle terminal creation."""
+        from acp.exceptions import RequestError
+
+        raise RequestError.method_not_found("terminal/create")
+
+    async def terminal_output(self, *args: Any, **kwargs: Any) -> Any:
+        """Handle terminal output."""
+        from acp.exceptions import RequestError
+
+        raise RequestError.method_not_found("terminal/output")
+
+    async def release_terminal(self, *args: Any, **kwargs: Any) -> Any:
+        """Handle terminal release."""
+        from acp.exceptions import RequestError
+
+        raise RequestError.method_not_found("terminal/release")
+
+    async def wait_for_terminal_exit(self, *args: Any, **kwargs: Any) -> Any:
+        """Handle terminal exit wait."""
+        from acp.exceptions import RequestError
+
+        raise RequestError.method_not_found("terminal/wait_for_exit")
+
+    async def kill_terminal(self, *args: Any, **kwargs: Any) -> Any:
+        """Handle terminal kill."""
+        from acp.exceptions import RequestError
+
+        raise RequestError.method_not_found("terminal/kill")
+
+    async def session_update(
+        self,
+        session_id: str,
+        update: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Handle session updates from agent."""
+        # Registry doesn't need to handle updates
+
+    async def ext_method(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
+        """Handle extension methods."""
+        from acp.exceptions import RequestError
+
+        raise RequestError.method_not_found(method)
+
+    async def ext_notification(self, method: str, params: dict[str, Any]) -> None:
+        """Handle extension notifications."""
+        logger.debug(f"Extension notification: {method}")
+
+    def on_connect(self, conn: Any) -> None:
+        """Called when client connects to agent."""
+        logger.debug("Registry connected to ACP agent")

codeoptix/adapters/__init__.py

@@ -0,0 +1,18 @@
+"""Agent adapters for CodeOptix."""
+
+from codeoptix.adapters.base import AgentAdapter, AgentOutput
+from codeoptix.adapters.basic import BasicAdapter
+from codeoptix.adapters.claude_code import ClaudeCodeAdapter
+from codeoptix.adapters.codex import CodexAdapter
+from codeoptix.adapters.factory import create_adapter
+from codeoptix.adapters.gemini_cli import GeminiCLIAdapter
+
+__all__ = [
+    "AgentAdapter",
+    "AgentOutput",
+    "BasicAdapter",
+    "ClaudeCodeAdapter",
+    "CodexAdapter",
+    "GeminiCLIAdapter",
+    "create_adapter",
+]