agent/agent.py

"""
NeuraPrompt Agent v8.9 — Main Core (GitHub MCP + JSON/Auth Bug Fixes)
"""

import os
import json
import time
import asyncio
import logging
import re  # FIX 1: Added missing import
from typing import Optional, Dict, Any

import requests
from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from .prompts.system_prompt import get_system_prompt
from .memory.memory import get_memory_manager
from .tools.registry import register_tool, get_tool, get_tool_descriptions, list_tools  # ← FIX 2: Added list_tools
from .tools import web_tools, code_tools, file_tools, vision_tools
from .tools.github_tools import (  # ← NEW: GitHub tools
    github_list_repos, github_create_repo, github_get_repo,
    github_list_issues, github_create_issue, github_read_file,
    github_write_file, github_create_branch, github_create_pull_request,
    github_search_code, github_get_user_profile
)
from .schemas.tool_schemas import (
    RunShellInput, RunPythonInput, CreateFileInput,
    WebSearchInput, FetchUrlInput, AnalyzeImageInput,
    # NEW: GitHub schemas
    GitHubListReposInput, GitHubCreateRepoInput, GitHubGetRepoInput,
    GitHubListIssuesInput, GitHubCreateIssueInput, GitHubReadFileInput,
    GitHubWriteFileInput, GitHubCreateBranchInput, GitHubCreatePullRequestInput,
    GitHubSearchCodeInput, GitHubGetUserProfileInput
)

log = logging.getLogger("agent.core.v8.9")

# ==================== CONFIG ====================
OPENROUTER_KEY = os.getenv("OPENROUTE_KEY", "")
MAX_STEPS = 8
MAX_TOKENS = 22000
PRIMARY_MODEL = "openrouter/owl-alpha"
TEMPERATURE = 0.15

# ==================== TOOL SCHEMAS ====================
TOOL_SCHEMAS: Dict[str, type[BaseModel]] = {
    # Existing
    "run_shell": RunShellInput,
    "run_python": RunPythonInput,
    "create_file": CreateFileInput,
    "web_search": WebSearchInput,
    "fetch_url": FetchUrlInput,
    "analyze_image": AnalyzeImageInput,
    # GitHub MCP
    "github_list_repos": GitHubListReposInput,
    "github_create_repo": GitHubCreateRepoInput,
    "github_get_repo": GitHubGetRepoInput,
    "github_list_issues": GitHubListIssuesInput,
    "github_create_issue": GitHubCreateIssueInput,
    "github_read_file": GitHubReadFileInput,
    "github_write_file": GitHubWriteFileInput,
    "github_create_branch": GitHubCreateBranchInput,
    "github_create_pull_request": GitHubCreatePullRequestInput,
    "github_search_code": GitHubSearchCodeInput,
    "github_get_user_profile": GitHubGetUserProfileInput,
}

# ==================== REGISTER TOOLS ====================
# Existing tools
register_tool("web_search", web_tools.web_search, "Search the web for current information")
register_tool("fetch_url", web_tools.fetch_url, "Fetch and extract content from a URL")
register_tool("run_python", code_tools.run_python, "Execute Python code (calculations, data processing)")
register_tool("run_shell", code_tools.run_shell, "Execute shell commands")
register_tool("create_file", file_tools.create_file, "Create files (supports batch)")
register_tool("read_file", file_tools.read_file, "Read file content")
register_tool("write_file", file_tools.write_file, "Write to file")
register_tool("list_dir", file_tools.list_directory, "List directory contents")
register_tool("analyze_image", vision_tools.analyze_image, "Analyze images with vision")

# GitHub MCP tools
register_tool("github_list_repos", github_list_repos, "List user's GitHub repositories")
register_tool("github_create_repo", github_create_repo, "Create a new GitHub repository")
register_tool("github_get_repo", github_get_repo, "Get repository details")
register_tool("github_list_issues", github_list_issues, "List repository issues")
register_tool("github_create_issue", github_create_issue, "Create a new issue")
register_tool("github_read_file", github_read_file, "Read file from repository")
register_tool("github_write_file", github_write_file, "Create or update file in repository")
register_tool("github_create_branch", github_create_branch, "Create a new branch")
register_tool("github_create_pull_request", github_create_pull_request, "Create a pull request")
register_tool("github_search_code", github_search_code, "Search code across GitHub")
register_tool("github_get_user_profile", github_get_user_profile, "Get user's GitHub profile")

# ==================== MODELS ====================
class AgentRequest(BaseModel):
    user_id: str
    goal: str
    max_steps: int = MAX_STEPS
    context: Optional[str] = None
    memory_context: Optional[str] = None

# ==================== SMALL, FOCUSED FUNCTIONS (Clean Architecture) ====================

def _build_system_prompt(memory_context: str = "") -> str:
    """Centralized prompt builder - calls the prompt module."""
    tool_descs = get_tool_descriptions()
    return get_system_prompt(memory_context=memory_context, tool_descriptions=tool_descs)

async def _call_llm(messages: list) -> Dict[str, Any]:
    """Call LLM and return parsed dict. Expects JSON output from prompt."""
    if not OPENROUTER_KEY:
        raise HTTPException(503, "OPENROUTE_KEY not configured")

    headers = {
        "Authorization": f"Bearer {OPENROUTER_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": PRIMARY_MODEL,
        "messages": messages,
        "max_tokens": MAX_TOKENS,
        "temperature": TEMPERATURE,
    }

    try:
        r = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=headers,
            json=payload,
            timeout=45
        )
        r.raise_for_status()
        content = r.json()["choices"][0]["message"]["content"].strip()

        # Original parsing — simple, was working with Owl Alpha
        if content.startswith("{"):
            try:
                return json.loads(content)
            except json.JSONDecodeError:
                pass

        # Fallback for non-JSON responses
        clean = content.strip().strip("```").strip()
        return {
            "thought": "Your message has been responded directly.",
            "action": "finish",
            "input": {"final_answer": clean}
        }
    except requests.exceptions.RequestException as e:
        log.error(f"OpenRouter request failed: {e}")
        raise HTTPException(503, f"LLM request failed: {str(e)}")
    except Exception as e:
        log.error(f"LLM error: {e}")
        raise HTTPException(503, "LLM service unavailable")


async def _dispatch_tool(name: str, raw_input: dict) -> str:
    """Async tool dispatch with Pydantic validation (modern pattern)."""
    tool_info = get_tool(name)
    if not tool_info:
        return f"Unknown tool: {name}"

    # Validate input if schema exists
    schema = TOOL_SCHEMAS.get(name)
    if schema:
        try:
            validated = schema(**raw_input)
            raw_input = validated.model_dump()
        except Exception as e:
            return f"Invalid parameters for {name}: {str(e)}"

    try:
        fn = tool_info["fn"]
        if asyncio.iscoroutinefunction(fn):
            return await fn(**raw_input)
        else:
            return await asyncio.to_thread(fn, **raw_input)
    except Exception as e:
        log.error(f"Tool error ({name}): {e}")
        return f"Tool error ({name}): {str(e)}"


def _make_sse(event_type: str, step: int, content: str, tool: str = "", elapsed: float = 0.0, reasoning: str = "") -> str:
    """Standard SSE formatter matching the working UI."""
    data = {
        "type": event_type,
        "step": step,
        "content": content,
        "tool": tool,
        "elapsed": round(elapsed, 2),
        "thought": reasoning
    }
    return f"data: {json.dumps(data)}\n\n"


async def _execute_tool_and_update_history(
    action: str,
    tool_input: dict,
    messages: list,
    thought: str  # FIX: Added thought parameter to preserve model reasoning
) -> tuple[str, float]:
    """Execute tool + update conversation history. Returns (result, elapsed_time)."""
    tool_start = time.time()
    result = await _dispatch_tool(action, tool_input)
    tool_elapsed = time.time() - tool_start

    # FIX: Preserve the model's actual thought in history instead of replacing with generic text
    messages.append({
        "role": "assistant",
        "content": json.dumps({"thought": thought, "action": action, "input": tool_input})
    })
    messages.append({
        "role": "user",
        "content": f"[Tool Result from {action}]\n{result}\n\nContinue with next JSON action or finish."
    })

    return result, tool_elapsed

# ==================== MAIN STREAMING AGENT ====================

async def stream_agent(
    goal: str,
    user_id: str,
    context: str = "",
    memory_context: str = "",
    max_steps: int = MAX_STEPS
):
    memory = get_memory_manager(user_id)
    full_memory = memory.get_context_for_agent() + (f"\n\n{memory_context}" if memory_context else "")

    system_prompt = _build_system_prompt(full_memory)
    messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": f"user_id: {user_id}\nGoal: {goal}" + (f"\nContext: {context}" if context else "")}
    ]

    total_start = time.time()

    for step in range(1, max_steps + 1):
        step_start = time.time()

        # 1. Get LLM decision
        parsed = await _call_llm(messages)
        # FIX: Use "thought" key to match system prompt (was "reasoning")
        thought = parsed.get("thought", "Thinking...")
        action = parsed.get("action", "finish")
        tool_input = parsed.get("input", {})
        
        # FIX: Ensure tool_input is a dict
        if not isinstance(tool_input, dict):
            tool_input = {}
        
        # FIX: Auto-inject user_id for all GitHub tools so the model never forgets it
        if action.startswith("github_"):
            tool_input.setdefault("user_id", user_id)

        # 2. Always emit reasoning (this is what shows the nice bubbles in UI)
        yield _make_sse("reasoning", step, thought, action, time.time() - step_start, thought)

        # 3. Finish if done
        if action == "finish":
            # FIX: Better fallback handling for final_answer
            if isinstance(tool_input, dict):
                final = tool_input.get("final_answer", "") or thought
            else:
                final = str(tool_input) if tool_input else thought
            
            # Safety: if model leaked raw JSON/tool code into final_answer, use thought instead
            if not final or final.strip().startswith("{"):
                final = thought
            
            yield _make_sse("done", step, final, "finish", time.time() - total_start, thought)
            return

        # 4. Execute tool (emit start + result events manually for clean streaming)
        yield _make_sse("tool_start", step, f"Running {action}...", action, time.time() - step_start)

        # FIX: Pass thought to preserve model reasoning in conversation history
        tool_start = time.time()
        result = await _dispatch_tool(action, tool_input)
        tool_elapsed = time.time() - tool_start

        # FIX: Only add user message — NO assistant message with JSON
        messages.append({
            "role": "user",
            "content": f"[Tool Result from {action}]\n{result}\n\nContinue with next JSON action or finish."
        })

        display = result[:1800] + "..." if len(result) > 1800 else result
        yield _make_sse("result", step, display, action, tool_elapsed)

        # Small delay to keep UI responsive and avoid rate limits
        await asyncio.sleep(0.12)

    # Max steps reached
    yield _make_sse(
        "done",
        max_steps + 1,
        "I reached the maximum number of steps. Please try breaking the goal into smaller parts or rephrase it.",
        "finish",
        time.time() - total_start
    )

# ==================== FASTAPI ROUTER ====================
agent_router = APIRouter(prefix="/agent")

@agent_router.post("/stream")
async def agent_stream(req: AgentRequest):
    if not OPENROUTER_KEY:
        raise HTTPException(503, "OPENROUTE_KEY not configured")
    return StreamingResponse(
        stream_agent(req.goal, req.user_id, req.context or "", req.memory_context or "", req.max_steps),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )

@agent_router.get("/health")
async def health_check():
    return {
        "status": "ok",
        "version": "8.9",
        "tools": list_tools(),  # ← FIX 2: Clean, uses registry directly
    }