Tribh/devops-copilot

1"""2workflow_agents.py — Multi-Provider Planner & Executor3Supports: Gemini (google-genai) | Groq (llama-3.3-70b-versatile)4Set LLM_PROVIDER=gemini|groq in your .env5"""6import json7import os8import re9from abc import ABC, abstractmethod10from pydantic import BaseModel11from devops_copilot.agents.base import BaseAgent, AgentState12from devops_copilot.tools.registry import registry13from devops_copilot.utils.logger import logger14 15 16# ── Data models ────────────────────────────────────────────────────────────────17 18class PlanStep(BaseModel):19    tool_name: str20    arguments: dict21    thought: str22 23 24class Plan(BaseModel):25    steps: list[PlanStep]26 27 28# ── Shared ReAct prompt ────────────────────────────────────────────────────────29 30_PLANNER_SYSTEM = """You are the Planner Agent in a ReAct (Reason + Act) loop acting as an AI DevOps Copilot.31Your job is to resolve an incident step by step using only the available tools.32 33Rules:341. Output EXACTLY one step per response.352. After seeing tool results, decide whether to take another step or finish.363. If the goal is fully achieved, output an empty steps list: {{"steps": []}}374. If restarting or doing a destructive action, include REQUIRES_APPROVAL in the thought.38 39Output format (strict JSON, no markdown fences):40{{41  "steps": [42    {{43      "tool_name": "<tool name from available tools>",44      "arguments": {{"<arg>": "<value>"}},45      "thought": "<your reasoning>"46    }}47  ]48}}49 50Available tools:51{tools}52"""53 54 55# ── LLM Provider backends ──────────────────────────────────────────────────────56 57class _LLMBackend(ABC):58    """Abstract LLM backend. Implement `complete(prompt) -> str`."""59    @abstractmethod60    def complete(self, prompt: str) -> str: ...61 62 63class _GeminiBackend(_LLMBackend):64    def __init__(self, model: str):65        from google import genai  # lazy import66        api_key = os.getenv("GOOGLE_API_KEY")67        if not api_key:68            raise EnvironmentError("GOOGLE_API_KEY not set.")69        self._client = genai.Client(api_key=api_key)70        self._model = model71 72    def complete(self, prompt: str) -> str:73        response = self._client.models.generate_content(74            model=self._model, contents=prompt75        )76        return response.text.strip()77 78 79class _GroqBackend(_LLMBackend):80    def __init__(self, model: str):81        from groq import Groq  # lazy import82        api_key = os.getenv("GROQ_API_KEY")83        if not api_key:84            raise EnvironmentError("GROQ_API_KEY not set.")85        self._client = Groq(api_key=api_key)86        self._model = model87 88    def complete(self, prompt: str) -> str:89        response = self._client.chat.completions.create(90            model=self._model,91            messages=[{"role": "user", "content": prompt}],92            temperature=0.2,93        )94        return response.choices[0].message.content.strip()95 96 97def _build_backend(provider: str) -> _LLMBackend:98    """Factory: builds the correct LLM backend based on provider name."""99    provider = provider.lower()100    if provider == "gemini":101        model = os.getenv("GEMINI_MODEL", "gemini-2.0-flash")102        logger.info(f"Using Gemini backend: {model}")103        return _GeminiBackend(model)104    elif provider == "groq":105        model = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")106        logger.info(f"Using Groq backend: {model}")107        return _GroqBackend(model)108    else:109        raise ValueError(f"Unknown LLM_PROVIDER: '{provider}'. Use gemini or groq.")110 111 112def _mock_plan() -> Plan:113    """Fallback demo plan when no API key is configured."""114    return Plan(steps=[PlanStep(115        tool_name="get_metrics",116        arguments={"service": "payment-gateway"},117        thought="Check health metrics to detect anomalies."118    )])119 120 121# ── Agents ─────────────────────────────────────────────────────────────────────122 123class PlannerAgent(BaseAgent):124    """125    Provider-agnostic Planner Agent.126    Set LLM_PROVIDER=gemini|groq in your .env to choose backend.127    """128    def __init__(self):129        super().__init__(name="Planner", role="Strategic Planning")130 131    async def chat(self, message: str, state: AgentState) -> Plan:132        provider = os.getenv("LLM_PROVIDER", "gemini")133        try:134            backend = _build_backend(provider)135        except EnvironmentError as e:136            logger.warning(f"No API key for '{provider}' — using mock plan. ({e})")137            return _mock_plan()138        except ValueError as e:139            logger.error(str(e))140            return _mock_plan()141 142        tools_info = registry.list_tools()143        system_prompt = _PLANNER_SYSTEM.format(tools=json.dumps(tools_info, indent=2))144 145        conversation = "\n".join(146            f"{m['role'].upper()}: {m['content']}"147            for m in state.history[-6:]148        )149        full_prompt = f"{system_prompt}\n\n{conversation}\nUSER: {message}"150 151        try:152            raw = backend.complete(full_prompt)153            # Strip markdown fences if model wraps with ```json...```154            raw = re.sub(r"^```(?:json)?\n?", "", raw)155            raw = re.sub(r"\n?```$", "", raw)156 157            parsed = json.loads(raw)158            steps = [PlanStep(**s) for s in parsed.get("steps", [])]159            self._log_interaction(state, "assistant", raw)160            logger.info(f"[{provider.upper()}] Planner proposed {len(steps)} step(s).")161            return Plan(steps=steps)162 163        except json.JSONDecodeError as e:164            logger.error(f"LLM returned non-JSON: {raw[:200]} — {e}")165            return Plan(steps=[])166        except Exception as e:167            logger.error(f"LLM backend error: {e}")168            return Plan(steps=[])169 170 171class ExecutorAgent(BaseAgent):172    """Executes a single plan step using the Tool Registry."""173    def __init__(self):174        super().__init__(name="Executor", role="Task Execution")175 176    async def chat(self, step: PlanStep, state: AgentState) -> str:177        tool = registry.get_tool(step.tool_name)178        if not tool:179            msg = f"ERROR: Tool '{step.tool_name}' not found in registry."180            logger.error(msg)181            self._log_interaction(state, "tool_error", msg)182            return msg183        try:184            result = tool.execute(**step.arguments)185            result_str = json.dumps(result) if isinstance(result, dict) else str(result)186            self._log_interaction(state, "tool_result", result_str)187            return result_str188        except Exception as e:189            err = f"Execution error in '{step.tool_name}': {e}"190            logger.error(err)191            self._log_interaction(state, "tool_error", err)192            return err193