agent.py

import os
from typing import Optional
from crewai import Agent, Task, Crew, Process
from crewai_tools import GithubSearchTool
from google import genai  # Google's Gemini client

# ---------------------------
# CONFIG
# ---------------------------

# Required environment variables
GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")  # from AI Studio
DEFAULT_CONTENT_TYPES = ["code", "pr", "issue", "repo"]

# Gemini client setup
if not GOOGLE_API_KEY:
    raise ValueError("❌ Missing GOOGLE_API_KEY (get one from https://aistudio.google.com)")

client = genai.Client(api_key=GOOGLE_API_KEY)
MODEL_NAME = os.getenv("GEMINI_MODEL", "gemini-1.5-flash")  # fast & cheap

# ---------------------------
# Custom Gemini-backed LLM wrapper
# ---------------------------
class GeminiLLM:
    def __init__(self, model: str):
        self.model = model

    def generate(self, prompt: str) -> str:
        """CrewAI-compatible .generate() interface."""
        try:
            response = client.models.generate_content(
                model=self.model,
                contents=prompt,
                generation_config={"temperature": 0.7}
            )
            return response.text
        except Exception as e:
            return f"⚠️ Gemini API error: {e}"

# Instantiate a Gemini LLM
gemini_llm = GeminiLLM(MODEL_NAME)

# ---------------------------
# GithubSearchTool
# ---------------------------
def github_tool(repo_url: Optional[str] = None) -> GithubSearchTool:
    if not GITHUB_TOKEN:
        raise RuntimeError("Missing GITHUB_TOKEN in environment.")
    if repo_url:
        return GithubSearchTool(
            github_repo=repo_url,
            gh_token=GITHUB_TOKEN,
            content_types=DEFAULT_CONTENT_TYPES,
        )
    return GithubSearchTool(
        gh_token=GITHUB_TOKEN,
        content_types=DEFAULT_CONTENT_TYPES,
    )

# ---------------------------
# Agents
# ---------------------------
def make_agents(repo_url: str):
    repo_search = github_tool(repo_url)

    repo_mapper = Agent(
        role="Repository Mapper",
        goal="Map repository structure, dependencies, and build tools.",
        backstory="Understands folder trees and project metadata deeply.",
        tools=[repo_search],
        llm=gemini_llm,
        verbose=True,
    )

    code_reviewer = Agent(
        role="Code Reviewer",
        goal="Perform detailed code reviews for quality, readability, and maintainability.",
        backstory="An experienced engineer identifying refactors and anti-patterns.",
        tools=[repo_search],
        llm=gemini_llm,
        verbose=True,
    )

    security_auditor = Agent(
        role="Security Auditor",
        goal="Detect security vulnerabilities, secrets, and risky dependencies.",
        backstory="Thinks like a red team but reports like a white hat.",
        tools=[repo_search],
        llm=gemini_llm,
        verbose=True,
    )

    doc_explainer = Agent(
        role="Documentation Explainer",
        goal="Explain project architecture, flow, and quickstart guide.",
        backstory="A great teacher simplifying complex systems.",
        tools=[repo_search],
        llm=gemini_llm,
        verbose=True,
    )

    manager = Agent(
        role="Engineering Manager",
        goal="Coordinate all other agents and compile a cohesive final report.",
        backstory="A senior manager ensuring thoroughness and clarity.",
        allow_delegation=True,
        llm=gemini_llm,
        verbose=True,
    )

    return repo_mapper, code_reviewer, security_auditor, doc_explainer, manager

# ---------------------------
# Tasks
# ---------------------------
def make_tasks(repo_url: str, brief: str = ""):
    preface = f"Target repository: {repo_url}\n{brief}\nPrefer file references and clear sections."

    t_map = Task(
        description=f"{preface}\nMap repo structure and dependencies.",
        expected_output="Markdown summary of repo structure.",
        agent_role="Repository Mapper",
    )

    t_review = Task(
        description=f"{preface}\nCode review with key issues, refactors, and suggestions.",
        expected_output="Concise bullet-based code review with examples.",
        agent_role="Code Reviewer",
    )

    t_sec = Task(
        description=f"{preface}\nSecurity audit (secrets, unsafe libs, injections, etc.).",
        expected_output="Security findings in table format.",
        agent_role="Security Auditor",
    )

    t_doc = Task(
        description=f"{preface}\nExplain architecture, modules, and how to run the project.",
        expected_output="Readable architecture explanation and quickstart.",
        agent_role="Documentation Explainer",
    )

    t_merge = Task(
        description="Merge all findings into one cohesive Markdown report with clear sections.",
        expected_output="Final Markdown report.",
        agent_role="Engineering Manager",
    )

    return t_map, t_review, t_sec, t_doc, t_merge

# ---------------------------
# Runner
# ---------------------------
def run_repo_review(repo_url: str, brief: str = "") -> str:
    repo_mapper, reviewer, auditor, explainer, manager = make_agents(repo_url)
    t_map, t_review, t_sec, t_doc, t_merge = make_tasks(repo_url, brief)

    crew = Crew(
        agents=[repo_mapper, reviewer, auditor, explainer, manager],
        tasks=[t_map, t_review, t_sec, t_doc, t_merge],
        process=Process.hierarchical,
        manager_agent=manager,
        verbose=True,
    )
    result = crew.kickoff()
    return str(result)