updated the model

app/services/ai.py

@@ -1,25 +1,28 @@
 import json
 import logging
-from openai import AsyncOpenAI
+import asyncio
+from google import genai
+from google.genai import types
 from app.config import settings
 
 logger = logging.getLogger(__name__)
 
-api_key = settings.openai_api_key or "mock-key-for-testing"
-client = AsyncOpenAI(api_key=api_key)
+if settings.gemini_api_key:
+    # Initialize google-genai client
+    ai_client = genai.Client(api_key=settings.gemini_api_key)
+else:
+    ai_client = None
+
+async def get_gemini_model():
+    return 'gemini-2.5-flash'
 
 async def enhance_security_issues(issues: list[dict]) -> dict:
-    """
-    Takes a list of basic security issues and uses an LLM to provide:
-    - Contextual severity
-    - Natural language explanations
-    - Auto-generated remediation code snippets
-    """
-    if not settings.openai_api_key:
-        logger.warning("OPENAI_API_KEY is not set. AI enhancements are skipped.")
+    if not settings.gemini_api_key:
+        logger.warning("GEMINI_API_KEY is not set. AI enhancements are skipped.")
         return {"enhanced_issues": issues}
 
     prompt = (
+        "You are a senior cybersecurity automation agent. Always respond with valid JSON.\n"
         "Analyze the following security vulnerabilities:\n"
         f"{json.dumps(issues, indent=2)}\n\n"
         "Return a JSON object with a single key 'enhanced_issues' containing a list of objects. "
@@ -31,75 +34,69 @@ async def enhance_security_issues(issues: list[dict]) -> dict:
     )
 
     try:
-        response = await client.chat.completions.create(
-            model="gpt-3.5-turbo",
-            messages=[
-                {"role": "system", "content": "You are a senior cybersecurity automation agent. Always respond with valid JSON."},
-                {"role": "user", "content": prompt}
-            ],
-            response_format={"type": "json_object"},
-            temperature=0.2,
+        model_name = await get_gemini_model()
+        response = await ai_client.aio.models.generate_content(
+            model=model_name,
+            contents=prompt,
+            config=types.GenerateContentConfig(
+                response_mime_type="application/json",
+                temperature=0.2,
+            )
         )
-        content = response.choices[0].message.content
-        if not content:
-            return {"enhanced_issues": issues, "ai_error": "Empty response"}
-            
-        return json.loads(content)
+        if response.text:
+            return json.loads(response.text)
+        return {"enhanced_issues": issues, "ai_error": "Empty response"}
     except Exception as e:
         logger.error(f"AI Generation Error: {str(e)}")
         return {"enhanced_issues": issues, "ai_error": str(e)}
 
 async def chat_with_scan_context(scan_id: str, context_data: dict, user_message: str) -> str:
-    """
-    Allows a user to ask a question about a specific scan's results.
-    """
-    if not settings.openai_api_key:
-        return "AI Chat is disabled because OPENAI_API_KEY is not configured."
+    if not settings.gemini_api_key:
+        return "AI Chat is disabled because GEMINI_API_KEY is not configured."
 
-    system_prompt = (
+    prompt = (
         "You are SecureLens AI, an expert cybersecurity assistant. "
         "You are helping a developer understand a security scan report for their website. "
-        f"Here is the context of the scan: {json.dumps(context_data)}"
+        f"Here is the context of the scan: {json.dumps(context_data)}\n\n"
+        f"User Message: {user_message}"
     )
 
     try:
-        response = await client.chat.completions.create(
-            model="gpt-3.5-turbo",
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_message}
-            ],
-            temperature=0.5,
+        model_name = await get_gemini_model()
+        response = await ai_client.aio.models.generate_content(
+            model=model_name,
+            contents=prompt,
+            config=types.GenerateContentConfig(
+                temperature=0.5,
+            )
         )
-        return response.choices[0].message.content or "No response from AI."
+        return response.text or "No response from AI."
     except Exception as e:
         logger.error(f"AI Chat Error: {str(e)}")
         return "I encountered an error trying to process your request."
 
 async def generate_threat_narrative(context_data: dict) -> str:
-    """
-    Weaves multiple scan issues into a cohesive attack sequence.
-    """
-    if not settings.openai_api_key:
-        return "AI Threat Narrative is disabled because OPENAI_API_KEY is not configured."
+    if not settings.gemini_api_key:
+        return "AI Threat Narrative is disabled because GEMINI_API_KEY is not configured."
 
-    system_prompt = (
+    prompt = (
         "You are a senior cybersecurity red-teamer. Analyze the following security scan results "
         "and weave them into a single, cohesive 'Threat Narrative'. Explain how an attacker might "
         "chain these specific vulnerabilities together to compromise the system. "
-        "Keep it professional, concise (2-3 paragraphs), and actionable."
+        "Keep it professional, concise (2-3 paragraphs), and actionable.\n\n"
+        f"Context: {json.dumps(context_data)}"
     )
 
     try:
-        response = await client.chat.completions.create(
-            model="gpt-3.5-turbo",
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": json.dumps(context_data)}
-            ],
-            temperature=0.7,
+        model_name = await get_gemini_model()
+        response = await ai_client.aio.models.generate_content(
+            model=model_name,
+            contents=prompt,
+            config=types.GenerateContentConfig(
+                temperature=0.7,
+            )
         )
-        return response.choices[0].message.content or "Could not generate threat narrative."
+        return response.text or "Could not generate threat narrative."
     except Exception as e:
         logger.error(f"AI Narrative Error: {str(e)}")
         return "I encountered an error trying to generate the threat narrative."

app/services/code_scanner/orchestrator.py

@@ -1,7 +1,8 @@
 import json
 import logging
 from typing import List, Dict, Any
-from openai import AsyncOpenAI
+from google import genai
+from google.genai import types
 
 from app.config import settings
 from app.services.code_scanner.github_client import GitHubClient
@@ -9,26 +10,28 @@ from app.schemas.code_scan import VulnerabilityIssue
 
 logger = logging.getLogger(__name__)
 
-api_key = settings.openai_api_key or "mock-key-for-testing"
-client = AsyncOpenAI(api_key=api_key)
+if settings.gemini_api_key:
+    # google-genai client init
+    ai_client = genai.Client(api_key=settings.gemini_api_key)
+else:
+    ai_client = None
 
 class CodeScanOrchestrator:
     def __init__(self, repo_url: str, github_token: str, branch: str = "main"):
         self.repo_url = repo_url
         self.branch = branch
         self.github = GitHubClient(token=github_token)
+        # We use gemini-2.5-flash for fast and cost-effective analysis
+        self.model_name = 'gemini-2.5-flash'
 
     async def triage_files(self, all_files: List[str]) -> List[str]:
         """
         Uses the LLM to select which files are most likely to contain security vulnerabilities 
-        (e.g., config files, routers, auth logic).
         """
-        if not settings.openai_api_key:
-            logger.warning("OPENAI_API_KEY is not set. Triaging all files up to a limit.")
-            return all_files[:10] # Hard limit for testing
+        if not settings.gemini_api_key:
+            logger.warning("GEMINI_API_KEY is not set. Triaging all files up to a limit.")
+            return all_files[:10]
 
-        # To avoid context limit issues, we might want to chunk this, but for now we pass the list
-        # We can enforce a soft limit on the string length
         files_str = "\n".join(all_files)
         if len(files_str) > 15000:
             files_str = files_str[:15000] + "\n... (truncated)"
@@ -42,42 +45,35 @@ class CodeScanOrchestrator:
         )
 
         try:
-            response = await client.chat.completions.create(
-                model="gpt-3.5-turbo",
-                messages=[
-                    {"role": "system", "content": "You always respond with valid JSON."},
-                    {"role": "user", "content": prompt}
-                ],
-                response_format={"type": "json_object"},
-                temperature=0.1,
+            response = await ai_client.aio.models.generate_content(
+                model=self.model_name,
+                contents=prompt,
+                config=types.GenerateContentConfig(
+                    response_mime_type="application/json",
+                    temperature=0.1,
+                )
             )
-            content = response.choices[0].message.content
-            if content:
-                data = json.loads(content)
+            if response.text:
+                data = json.loads(response.text)
                 return data.get("critical_files", [])
         except Exception as e:
             logger.error(f"Error triaging files: {e}")
             
-        return all_files[:10] # Fallback
+        return all_files[:10]
 
     async def analyze_files(self, triaged_files: List[str]) -> List[VulnerabilityIssue]:
-        """
-        Fetches the contents of the triaged files and uses the LLM to find vulnerabilities.
-        """
         vulnerabilities = []
         
-        if not settings.openai_api_key:
+        if not settings.gemini_api_key:
             return []
 
-        # Analyze files sequentially or in batches (sequential to avoid rate limits for now)
         for file_path in triaged_files:
             content = await self.github.get_file_content(self.repo_url, file_path, self.branch)
             if not content:
                 continue
                 
-            # Truncate very large files
-            if len(content) > 20000:
-                content = content[:20000]
+            if len(content) > 30000:
+                content = content[:30000]
 
             prompt = (
                 f"Review the following code from the file '{file_path}' for security vulnerabilities.\n"
@@ -93,19 +89,16 @@ class CodeScanOrchestrator:
             )
 
             try:
-                response = await client.chat.completions.create(
-                    model="gpt-3.5-turbo",
-                    messages=[
-                        {"role": "system", "content": "You are a SAST security agent. Always respond with valid JSON."},
-                        {"role": "user", "content": prompt}
-                    ],
-                    response_format={"type": "json_object"},
-                    temperature=0.2,
+                response = await ai_client.aio.models.generate_content(
+                    model=self.model_name,
+                    contents=prompt,
+                    config=types.GenerateContentConfig(
+                        response_mime_type="application/json",
+                        temperature=0.2,
+                    )
                 )
-                
-                resp_content = response.choices[0].message.content
-                if resp_content:
-                    data = json.loads(resp_content)
+                if response.text:
+                    data = json.loads(response.text)
                     vulns = data.get("vulnerabilities", [])
                     for v in vulns:
                         vulnerabilities.append(VulnerabilityIssue(
@@ -125,7 +118,7 @@ class CodeScanOrchestrator:
         if not vulnerabilities:
             return "No obvious security vulnerabilities found in the scanned files."
             
-        if not settings.openai_api_key:
+        if not settings.gemini_api_key:
             return f"Found {len(vulnerabilities)} potential issues."
 
         issues_data = [v.model_dump() for v in vulnerabilities]
@@ -137,15 +130,14 @@ class CodeScanOrchestrator:
         )
 
         try:
-            response = await client.chat.completions.create(
-                model="gpt-3.5-turbo",
-                messages=[
-                    {"role": "system", "content": "You are a cybersecurity expert."},
-                    {"role": "user", "content": prompt}
-                ],
-                temperature=0.4,
+            response = await ai_client.aio.models.generate_content(
+                model=self.model_name,
+                contents=prompt,
+                config=types.GenerateContentConfig(
+                    temperature=0.4,
+                )
             )
-            return response.choices[0].message.content or "Could not generate summary."
+            return response.text or "Could not generate summary."
         except Exception as e:
             logger.error(f"Error generating summary: {e}")
             return f"Found {len(vulnerabilities)} potential issues."