From ae4ed3062af1a62a3f867f2a39e09eaf0de8a992 Mon Sep 17 00:00:00 2001 From: rarebuffalo Date: Sat, 25 Apr 2026 21:29:19 +0530 Subject: [PATCH] updated model for suitable testing --- app/routers/code_scan.py | 16 ++++- app/services/ai.py | 2 +- app/services/code_scanner/orchestrator.py | 75 +++++++++++++---------- list_models.py | 11 ++++ models.txt | 2 +- 5 files changed, 72 insertions(+), 34 deletions(-) create mode 100644 list_models.py diff --git a/app/routers/code_scan.py b/app/routers/code_scan.py index 0331dfd..cae5fa8 100644 --- a/app/routers/code_scan.py +++ b/app/routers/code_scan.py @@ -56,6 +56,20 @@ async def analyze_codebase(request: CodeScanRequest): except Exception as e: logger.error(f"Code scan failed: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) +@router.get("/code-scan/models") +async def list_available_models(): + if not settings.gemini_api_key: + raise HTTPException(status_code=500, detail="GEMINI_API_KEY is not set.") + try: + from google import genai + client = genai.Client(api_key=settings.gemini_api_key) + models = [] + for model in client.models.list(): + if 'generateContent' in model.supported_actions: + models.append(model.name) + return {"models": models} + except Exception as e: + raise HTTPException(status_code=500, detail=f"Error fetching models: {e}") @router.post("/code-scan/chat", response_model=CodeChatResponse) async def chat_with_scan(request: CodeChatRequest): @@ -87,7 +101,7 @@ async def chat_with_scan(request: CodeChatRequest): client = genai.Client(api_key=settings.gemini_api_key) response = await client.aio.models.generate_content( - model='gemini-2.5-flash', + model='gemini-2.0-flash', contents=prompt, config=types.GenerateContentConfig( temperature=0.5, diff --git a/app/services/ai.py b/app/services/ai.py index 8c9bc4b..125c56a 100644 --- a/app/services/ai.py +++ b/app/services/ai.py @@ -14,7 +14,7 @@ else: ai_client = None async def get_gemini_model(): - return 'gemini-2.5-flash' + return 'gemini-2.0-flash' async def enhance_security_issues(issues: list[dict]) -> dict: if not settings.gemini_api_key: diff --git a/app/services/code_scanner/orchestrator.py b/app/services/code_scanner/orchestrator.py index c26d9ea..fd8f754 100644 --- a/app/services/code_scanner/orchestrator.py +++ b/app/services/code_scanner/orchestrator.py @@ -3,6 +3,7 @@ import logging from typing import List, Dict, Any from google import genai from google.genai import types +import asyncio from app.config import settings from app.services.code_scanner.github_client import GitHubClient @@ -21,8 +22,8 @@ class CodeScanOrchestrator: 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' + # We use gemini-2.0-flash for high rate limits and stability + self.model_name = 'gemini-2.0-flash' async def triage_files(self, all_files: List[str]) -> List[str]: """ @@ -30,7 +31,7 @@ class CodeScanOrchestrator: """ 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] + return all_files[:5] files_str = "\n".join(all_files) if len(files_str) > 15000: @@ -41,7 +42,7 @@ class CodeScanOrchestrator: f"{files_str}\n\n" "Select the most critical files to review for security vulnerabilities (e.g., SAST, hardcoded secrets, SQLi, Auth bypass). " "Return a JSON object with a single key 'critical_files' containing a list of the exact file paths. " - "Do not select more than 15 files." + "Do not select more than 5 files." ) try: @@ -59,18 +60,23 @@ class CodeScanOrchestrator: except Exception as e: logger.error(f"Error triaging files: {e}") - return all_files[:10] + return all_files[:5] async def analyze_files(self, triaged_files: List[str]) -> List[VulnerabilityIssue]: - vulnerabilities = [] - if not settings.gemini_api_key: return [] - for file_path in triaged_files: + vulnerabilities = [] + semaphore = asyncio.Semaphore(5) # Max 5 concurrent requests to avoid rate limits + + async def process_file(file_path: str): + # Skip massive dependency lock files as they are too slow and unhelpful for SAST + if file_path.endswith('package-lock.json') or file_path.endswith('yarn.lock'): + return [] + content = await self.github.get_file_content(self.repo_url, file_path, self.branch) if not content: - continue + return [] if len(content) > 30000: content = content[:30000] @@ -88,30 +94,37 @@ class CodeScanOrchestrator: "'line_number' (integer or null if general)." ) - try: - 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, + file_vulns = [] + async with semaphore: + try: + 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, + ) ) - ) - if response.text: - data = json.loads(response.text) - vulns = data.get("vulnerabilities", []) - for v in vulns: - vulnerabilities.append(VulnerabilityIssue( - file_path=file_path, - severity=v.get("severity", "Medium"), - issue=v.get("issue", "Unknown Issue"), - explanation=v.get("explanation", ""), - suggested_fix=v.get("suggested_fix"), - line_number=v.get("line_number") - )) - except Exception as e: - logger.error(f"Error analyzing file {file_path}: {e}") + if response.text: + data = json.loads(response.text) + vulns = data.get("vulnerabilities", []) + for v in vulns: + file_vulns.append(VulnerabilityIssue( + file_path=file_path, + severity=v.get("severity", "Medium"), + issue=v.get("issue", "Unknown Issue"), + explanation=v.get("explanation", ""), + suggested_fix=v.get("suggested_fix"), + line_number=v.get("line_number") + )) + except Exception as e: + logger.error(f"Error analyzing file {file_path}: {e}") + return file_vulns + results = await asyncio.gather(*(process_file(f) for f in triaged_files)) + for res in results: + vulnerabilities.extend(res) + return vulnerabilities async def generate_summary(self, vulnerabilities: List[VulnerabilityIssue]) -> str: diff --git a/list_models.py b/list_models.py new file mode 100644 index 0000000..6fac0e3 --- /dev/null +++ b/list_models.py @@ -0,0 +1,11 @@ +import os +from google import genai +from dotenv import load_dotenv + +load_dotenv('/home/Krishna-Singh/securelens-backend/.env') +client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY")) + +print("Supported Models:") +for model in client.models.list(): + if 'generateContent' in model.supported_actions: + print(f"- {model.name}") diff --git a/models.txt b/models.txt index eeb0611..fc051bc 100644 --- a/models.txt +++ b/models.txt @@ -1 +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 \ No newline at end of file +AVAILABLE MODELS: migrated to new SDK \ No newline at end of file