planner.py

from utils_llm import get_llm
from langchain_core.messages import SystemMessage, HumanMessage
import json
import re

PLANNER_PROMPT = """
You are a senior codebase architect. Your job is to break down a developer's query into a step-by-step investigation plan.
You can't see the code yet, but you need to guide the 'Execution Agent' on where to look.

Return ONLY a JSON list of strings (steps).
Example:
[
  "Search for payment gateway integration in the codebase",
  "List files in the 'routes' directory",
  "Check how state is managed in the frontend"
]

No explanation, no markdown, JUST the JSON list.
"""

def create_plan(query: str, openai_api_key: str):
    llm = get_llm(openai_api_key)
    
    messages = [
        SystemMessage(content=PLANNER_PROMPT),
        HumanMessage(content=query)
    ]
    
    response = llm.invoke(messages)
    content = response.content.strip()
    
    # Try to extract JSON if the LLM wraps it in markdown
    if "```" in content:
        match = re.search(r"```(?:json)?\s*(\[.*\])\s*```", content, re.DOTALL)
        if match:
            content = match.group(1)
        else:
            # Maybe it just has backticks but no json label
            content = content.replace("```json", "").replace("```", "").strip()
            
    try:
        steps = json.loads(content)
        if isinstance(steps, list):
            return steps
    except Exception as e:
        print(f"Error parsing plan: {e}")
        # Return a simple plan as fallback
        return [f"Search for '{query}' in the codebase", f"Analyze the findings related to '{query}'"]
    
    return [f"Search for '{query}' in the codebase"]