updated the model

.env.example

@@ -23,5 +23,5 @@ PATH_CHECK_TIMEOUT=3
 DATABASE_URL=postgresql+asyncpg://securelens:securelens@localhost:5433/securelens
 
 # AI Integration
-OPENAI_API_KEY=
+GEMINI_API_KEY=
 

app/config.py

@@ -22,7 +22,7 @@ class Settings(BaseSettings):
     jwt_algorithm: str = "HS256"
     jwt_expiry_minutes: int = 1440
 
-    openai_api_key: str | None = None
+    gemini_api_key: str | None = None
 
     model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
 

app/routers/__init__.py

@@ -1,17 +1 @@
-from .auth import router as auth
+# Empty init file to allow correct module importing
-from .health import router as health
-from .history import router as history
-from .scan import router as scan
-from .apikey import router as apikey
-from .report import router as report
-from .code_scan import router as code_scan
-
-__all__ = [
-    "auth",
-    "health",
-    "history",
-    "scan",
-    "apikey",
-    "report",
-    "code_scan"
-]

app/routers/code_scan.py

@@ -5,7 +5,7 @@ from fastapi import APIRouter, HTTPException
 from typing import Dict, Any
 
 from app.schemas.code_scan import CodeScanRequest, CodeScanResponse, CodeChatRequest, CodeChatResponse
-from app.services.code_scanner.orchestrator import CodeScanOrchestrator, client
+from app.services.code_scanner.orchestrator import CodeScanOrchestrator
 from app.config import settings
 
 logger = logging.getLogger(__name__)
@@ -59,32 +59,41 @@ async def analyze_codebase(request: CodeScanRequest):
 
 @router.post("/code-scan/chat", response_model=CodeChatResponse)
 async def chat_with_scan(request: CodeChatRequest):
-    if not settings.openai_api_key:
+    if not settings.gemini_api_key:
-        raise HTTPException(status_code=400, detail="AI Chat is disabled because OPENAI_API_KEY is not configured.")
+        raise HTTPException(status_code=400, detail="AI Chat is disabled because GEMINI_API_KEY is not configured.")
+        
+    from google import genai
+    with open("/app/models.txt", "w") as f:
+        # Just writing a placeholder, list_models is different in new SDK
+        f.write("AVAILABLE MODELS: migrated to new SDK")
         
     scan_data = scan_store.get(request.scan_id)
     if not scan_data:
         raise HTTPException(status_code=404, detail="Scan ID not found or expired.")
         
-    system_prompt = (
+    prompt = (
         "You are SecureLens AI, an expert application security assistant. "
         "You are helping a developer understand a security scan report for their codebase. "
         f"Here is the context of the scan for the repository {scan_data.repo_url}:\n"
         f"Summary: {scan_data.summary}\n"
         f"Vulnerabilities: {json.dumps([v.model_dump() for v in scan_data.issues])}\n\n"
+        f"User Message: {request.message}\n\n"
         "Answer the user's questions clearly, concisely, and professionally. Provide code fixes if requested."
     )
 
     try:
-        response = await client.chat.completions.create(
+        from google import genai
-            model="gpt-3.5-turbo",
+        from google.genai import types
-            messages=[
+        
-                {"role": "system", "content": system_prompt},
+        client = genai.Client(api_key=settings.gemini_api_key)
-                {"role": "user", "content": request.message}
+        response = await client.aio.models.generate_content(
-            ],
+            model='gemini-2.5-flash',
-            temperature=0.5,
+            contents=prompt,
+            config=types.GenerateContentConfig(
+                temperature=0.5,
+            )
         )
-        reply = response.choices[0].message.content or "No response from AI."
+        reply = response.text or "No response from AI."
         return CodeChatResponse(reply=reply)
     except Exception as e:
         logger.error(f"AI Chat Error: {str(e)}")

app/routers/scan.py

@@ -99,7 +99,7 @@ async def scan_website(
         score = calculate_score(all_issues)
         layers = calculate_layer_statuses(all_issues)
 
-        if settings.openai_api_key and all_issues:
+        if settings.gemini_api_key and all_issues:
             issues_dict_list = [i.model_dump() for i in all_issues]
             ai_data = await enhance_security_issues(issues_dict_list)
             enhanced_list = ai_data.get("enhanced_issues", [])

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"
+if settings.gemini_api_key:
-client = AsyncOpenAI(api_key=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:
-    """
+    if not settings.gemini_api_key:
-    Takes a list of basic security issues and uses an LLM to provide:
+        logger.warning("GEMINI_API_KEY is not set. AI enhancements are skipped.")
-    - 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.")
         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_name = await get_gemini_model()
-            model="gpt-3.5-turbo",
+        response = await ai_client.aio.models.generate_content(
-            messages=[
+            model=model_name,
-                {"role": "system", "content": "You are a senior cybersecurity automation agent. Always respond with valid JSON."},
+            contents=prompt,
-                {"role": "user", "content": prompt}
+            config=types.GenerateContentConfig(
-            ],
+                response_mime_type="application/json",
-            response_format={"type": "json_object"},
+                temperature=0.2,
-            temperature=0.2,
+            )
         )
-        content = response.choices[0].message.content
+        if response.text:
-        if not content:
+            return json.loads(response.text)
-            return {"enhanced_issues": issues, "ai_error": "Empty response"}
+        return {"enhanced_issues": issues, "ai_error": "Empty response"}
-            
-        return json.loads(content)
     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:
-    """
+    if not settings.gemini_api_key:
-    Allows a user to ask a question about a specific scan's results.
+        return "AI Chat is disabled because GEMINI_API_KEY is not configured."
-    """
-    if not settings.openai_api_key:
-        return "AI Chat is disabled because OPENAI_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_name = await get_gemini_model()
-            model="gpt-3.5-turbo",
+        response = await ai_client.aio.models.generate_content(
-            messages=[
+            model=model_name,
-                {"role": "system", "content": system_prompt},
+            contents=prompt,
-                {"role": "user", "content": user_message}
+            config=types.GenerateContentConfig(
-            ],
+                temperature=0.5,
-            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:
-    """
+    if not settings.gemini_api_key:
-    Weaves multiple scan issues into a cohesive attack sequence.
+        return "AI Threat Narrative is disabled because GEMINI_API_KEY is not configured."
-    """
-    if not settings.openai_api_key:
-        return "AI Threat Narrative is disabled because OPENAI_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_name = await get_gemini_model()
-            model="gpt-3.5-turbo",
+        response = await ai_client.aio.models.generate_content(
-            messages=[
+            model=model_name,
-                {"role": "system", "content": system_prompt},
+            contents=prompt,
-                {"role": "user", "content": json.dumps(context_data)}
+            config=types.GenerateContentConfig(
-            ],
+                temperature=0.7,
-            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"
+if settings.gemini_api_key:
-client = AsyncOpenAI(api_key=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:
+        if not settings.gemini_api_key:
-            logger.warning("OPENAI_API_KEY is not set. Triaging all files up to a limit.")
+            logger.warning("GEMINI_API_KEY is not set. Triaging all files up to a limit.")
-            return all_files[:10] # Hard limit for testing
+            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(
+            response = await ai_client.aio.models.generate_content(
-                model="gpt-3.5-turbo",
+                model=self.model_name,
-                messages=[
+                contents=prompt,
-                    {"role": "system", "content": "You always respond with valid JSON."},
+                config=types.GenerateContentConfig(
-                    {"role": "user", "content": prompt}
+                    response_mime_type="application/json",
-                ],
+                    temperature=0.1,
-                response_format={"type": "json_object"},
+                )
-                temperature=0.1,
             )
-            content = response.choices[0].message.content
+            if response.text:
-            if content:
+                data = json.loads(response.text)
-                data = json.loads(content)
                 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) > 30000:
-            if len(content) > 20000:
+                content = content[:30000]
-                content = content[:20000]
 
             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(
+                response = await ai_client.aio.models.generate_content(
-                    model="gpt-3.5-turbo",
+                    model=self.model_name,
-                    messages=[
+                    contents=prompt,
-                        {"role": "system", "content": "You are a SAST security agent. Always respond with valid JSON."},
+                    config=types.GenerateContentConfig(
-                        {"role": "user", "content": prompt}
+                        response_mime_type="application/json",
-                    ],
+                        temperature=0.2,
-                    response_format={"type": "json_object"},
+                    )
-                    temperature=0.2,
                 )
-                
+                if response.text:
-                resp_content = response.choices[0].message.content
+                    data = json.loads(response.text)
-                if resp_content:
-                    data = json.loads(resp_content)
                     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(
+            response = await ai_client.aio.models.generate_content(
-                model="gpt-3.5-turbo",
+                model=self.model_name,
-                messages=[
+                contents=prompt,
-                    {"role": "system", "content": "You are a cybersecurity expert."},
+                config=types.GenerateContentConfig(
-                    {"role": "user", "content": prompt}
+                    temperature=0.4,
-                ],
+                )
-                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."

models.txt

@@ -0,0 +1 @@
+AVAILABLE MODELS: models/gemini-2.5-flash, models/gemini-2.5-pro, models/gemini-2.0-flash, models/gemini-2.0-flash-001, models/gemini-2.0-flash-lite-001, models/gemini-2.0-flash-lite, models/gemini-2.5-flash-preview-tts, models/gemini-2.5-pro-preview-tts, models/gemma-3-1b-it, models/gemma-3-4b-it, models/gemma-3-12b-it, models/gemma-3-27b-it, models/gemma-3n-e4b-it, models/gemma-3n-e2b-it, models/gemma-4-26b-a4b-it, models/gemma-4-31b-it, models/gemini-flash-latest, models/gemini-flash-lite-latest, models/gemini-pro-latest, models/gemini-2.5-flash-lite, models/gemini-2.5-flash-image, models/gemini-3-pro-preview, models/gemini-3-flash-preview, models/gemini-3.1-pro-preview, models/gemini-3.1-pro-preview-customtools, models/gemini-3.1-flash-lite-preview, models/gemini-3-pro-image-preview, models/nano-banana-pro-preview, models/gemini-3.1-flash-image-preview, models/lyria-3-clip-preview, models/lyria-3-pro-preview, models/gemini-3.1-flash-tts-preview, models/gemini-robotics-er-1.5-preview, models/gemini-robotics-er-1.6-preview, models/gemini-2.5-computer-use-preview-10-2025, models/deep-research-max-preview-04-2026, models/deep-research-preview-04-2026, models/deep-research-pro-preview-12-2025, models/gemini-embedding-001, models/gemini-embedding-2-preview, models/gemini-embedding-2, models/aqa, models/imagen-4.0-generate-001, models/imagen-4.0-ultra-generate-001, models/imagen-4.0-fast-generate-001, models/veo-2.0-generate-001, models/veo-3.0-generate-001, models/veo-3.0-fast-generate-001, models/veo-3.1-generate-preview, models/veo-3.1-fast-generate-preview, models/veo-3.1-lite-generate-preview, models/gemini-2.5-flash-native-audio-latest, models/gemini-2.5-flash-native-audio-preview-09-2025, models/gemini-2.5-flash-native-audio-preview-12-2025, models/gemini-3.1-flash-live-preview

requirements.txt

@@ -14,6 +14,6 @@ pydantic[email]
 pytest
 pytest-asyncio
 alembic
-openai
+google-genai
 aiodns
 fpdf2