Update agent.py

agent.py

@@ -1,74 +1,114 @@
+# agent.py
 import os
-from typing import Optional
+from typing import Optional, List
 from crewai import Agent, Task, Crew, Process
 from crewai_tools import GithubSearchTool
-from google import genai  # Google's Gemini client
+from google import genai  # Gemini client (google-genai package)
+from dotenv import load_dotenv
+
+load_dotenv()
 
 # ---------------------------
 # 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"]
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 
-# Gemini client setup
 if not GOOGLE_API_KEY:
-    raise ValueError("❌ Missing GOOGLE_API_KEY (get one from https://aistudio.google.com)")
+    raise RuntimeError("❌ Missing GOOGLE_API_KEY (get one from https://aistudio.google.com)")
 
+# Gemini Client
 client = genai.Client(api_key=GOOGLE_API_KEY)
-MODEL_NAME = os.getenv("GEMINI_MODEL", "gemini-1.5-flash")  # fast & cheap
+MODEL_NAME = os.getenv("GEMINI_MODEL", "gemini-1.5-flash")
+
+DEFAULT_CONTENT_TYPES = ["code", "pr", "issue", "repo"]
 
 # ---------------------------
-# Custom Gemini-backed LLM wrapper
+# Gemini Embedding Adapter (Free Embeddings)
+# ---------------------------
+class GeminiEmbedding:
+    """Uses Google Gemini text-embedding-004 model (free tier)"""
+
+    def __init__(self, model="text-embedding-004", api_key=None):
+        self.model = model
+        self.client = genai.Client(api_key=api_key or GOOGLE_API_KEY)
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        vectors = []
+        for text in texts:
+            try:
+                res = self.client.models.embed_content(model=self.model, contents=text)
+                vectors.append(res.embedding.values)
+            except Exception as e:
+                print(f"⚠️ Embedding error: {e}")
+                vectors.append([])
+        return vectors
+
+    def embed_query(self, text: str) -> List[float]:
+        try:
+            res = self.client.models.embed_content(model=self.model, contents=text)
+            return res.embedding.values
+        except Exception as e:
+            print(f"⚠️ Embedding query error: {e}")
+            return []
+
+
+# ---------------------------
+# Gemini LLM Wrapper
 # ---------------------------
 class GeminiLLM:
     def __init__(self, model: str):
         self.model = model
 
     def generate(self, prompt: str) -> str:
-        """CrewAI-compatible .generate() interface."""
+        """CrewAI-compatible LLM generate method."""
         try:
             response = client.models.generate_content(
                 model=self.model,
                 contents=prompt,
-                generation_config={"temperature": 0.7}
+                generation_config={"temperature": 0.7, "max_output_tokens": 1024}
             )
             return response.text
         except Exception as e:
             return f"⚠️ Gemini API error: {e}"
 
-# Instantiate a Gemini LLM
+
+# Instantiate LLM + embedder
 gemini_llm = GeminiLLM(MODEL_NAME)
+embedder = GeminiEmbedding(api_key=GOOGLE_API_KEY)
 
 # ---------------------------
-# GithubSearchTool
+# GitHub Tool using free embeddings
 # ---------------------------
 def github_tool(repo_url: Optional[str] = None) -> GithubSearchTool:
+    """Create a GitHub Search Tool with free Gemini embeddings (no OpenAI key)."""
     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,
+            embedder=embedder,
         )
     return GithubSearchTool(
         gh_token=GITHUB_TOKEN,
         content_types=DEFAULT_CONTENT_TYPES,
+        embedder=embedder,
     )
 
+
 # ---------------------------
-# Agents
+# 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.",
+        goal="Map repo structure, dependencies, and frameworks.",
+        backstory="Understands directory trees, tech stacks, and configuration files.",
         tools=[repo_search],
         llm=gemini_llm,
         verbose=True,
@@ -76,8 +116,8 @@ def make_agents(repo_url: str):
 
     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.",
+        goal="Perform code review for quality, structure, and clarity.",
+        backstory="A senior engineer giving actionable review comments with examples.",
         tools=[repo_search],
         llm=gemini_llm,
         verbose=True,
@@ -85,8 +125,8 @@ def make_agents(repo_url: str):
 
     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.",
+        goal="Identify secrets, vulnerabilities, unsafe APIs, and dependencies.",
+        backstory="A white-hat hacker finding issues and giving fixes.",
         tools=[repo_search],
         llm=gemini_llm,
         verbose=True,
@@ -94,8 +134,8 @@ def make_agents(repo_url: str):
 
     doc_explainer = Agent(
         role="Documentation Explainer",
-        goal="Explain project architecture, flow, and quickstart guide.",
-        backstory="A great teacher simplifying complex systems.",
+        goal="Summarize architecture, data flow, and how to run the project.",
+        backstory="Explains tech systems simply and clearly with examples.",
         tools=[repo_search],
         llm=gemini_llm,
         verbose=True,
@@ -103,8 +143,8 @@ def make_agents(repo_url: str):
 
     manager = Agent(
         role="Engineering Manager",
-        goal="Coordinate all other agents and compile a cohesive final report.",
-        backstory="A senior manager ensuring thoroughness and clarity.",
+        goal="Coordinate all agents and compile a clear, cohesive final report.",
+        backstory="Ensures a professional, well-structured final document.",
         allow_delegation=True,
         llm=gemini_llm,
         verbose=True,
@@ -112,46 +152,48 @@ def make_agents(repo_url: str):
 
     return repo_mapper, code_reviewer, security_auditor, doc_explainer, manager
 
+
 # ---------------------------
-# Tasks
+# TASKS
 # ---------------------------
 def make_tasks(repo_url: str, brief: str = ""):
-    preface = f"Target repository: {repo_url}\n{brief}\nPrefer file references and clear sections."
+    prefix = f"Target Repository: {repo_url}\nBrief: {brief}\nInclude file paths where relevant."
 
     t_map = Task(
-        description=f"{preface}\nMap repo structure and dependencies.",
-        expected_output="Markdown summary of repo structure.",
+        description=f"{prefix}\nMap the repository structure, dependencies, languages, and build tools.",
+        expected_output="Markdown-formatted repository map with bullets and paths.",
         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.",
+        description=f"{prefix}\nPerform detailed code review (readability, refactors, testing, etc.).",
+        expected_output="Actionable review bullets grouped by issue type.",
         agent_role="Code Reviewer",
     )
 
     t_sec = Task(
-        description=f"{preface}\nSecurity audit (secrets, unsafe libs, injections, etc.).",
-        expected_output="Security findings in table format.",
+        description=f"{prefix}\nPerform security audit (secrets, vulnerabilities, dependencies).",
+        expected_output="Security findings table (Issue | Evidence | Risk | Fix).",
         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.",
+        description=f"{prefix}\nExplain architecture, modules, data flow, setup, and usage.",
+        expected_output="Readable explanation with Quickstart and architecture overview.",
         agent_role="Documentation Explainer",
     )
 
     t_merge = Task(
-        description="Merge all findings into one cohesive Markdown report with clear sections.",
-        expected_output="Final Markdown report.",
+        description="Merge all outputs into a clean, single Markdown report with clear sections and TOC.",
+        expected_output="Final cohesive Markdown report.",
         agent_role="Engineering Manager",
     )
 
     return t_map, t_review, t_sec, t_doc, t_merge
 
+
 # ---------------------------
-# Runner
+# RUNNER
 # ---------------------------
 def run_repo_review(repo_url: str, brief: str = "") -> str:
     repo_mapper, reviewer, auditor, explainer, manager = make_agents(repo_url)
@@ -164,5 +206,6 @@ def run_repo_review(repo_url: str, brief: str = "") -> str:
         manager_agent=manager,
         verbose=True,
     )
+
     result = crew.kickoff()
     return str(result)