updated model for suitable testing

This commit is contained in:
rarebuffalo
2026-04-25 21:29:19 +05:30
parent 139c8d982b
commit ae4ed3062a
5 changed files with 72 additions and 34 deletions

View File

@@ -56,6 +56,20 @@ async def analyze_codebase(request: CodeScanRequest):
except Exception as e: except Exception as e:
logger.error(f"Code scan failed: {str(e)}") logger.error(f"Code scan failed: {str(e)}")
raise HTTPException(status_code=500, detail=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) @router.post("/code-scan/chat", response_model=CodeChatResponse)
async def chat_with_scan(request: CodeChatRequest): 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) client = genai.Client(api_key=settings.gemini_api_key)
response = await client.aio.models.generate_content( response = await client.aio.models.generate_content(
model='gemini-2.5-flash', model='gemini-2.0-flash',
contents=prompt, contents=prompt,
config=types.GenerateContentConfig( config=types.GenerateContentConfig(
temperature=0.5, temperature=0.5,

View File

@@ -14,7 +14,7 @@ else:
ai_client = None ai_client = None
async def get_gemini_model(): async def get_gemini_model():
return 'gemini-2.5-flash' return 'gemini-2.0-flash'
async def enhance_security_issues(issues: list[dict]) -> dict: async def enhance_security_issues(issues: list[dict]) -> dict:
if not settings.gemini_api_key: if not settings.gemini_api_key:

View File

@@ -3,6 +3,7 @@ import logging
from typing import List, Dict, Any from typing import List, Dict, Any
from google import genai from google import genai
from google.genai import types from google.genai import types
import asyncio
from app.config import settings from app.config import settings
from app.services.code_scanner.github_client import GitHubClient from app.services.code_scanner.github_client import GitHubClient
@@ -21,8 +22,8 @@ class CodeScanOrchestrator:
self.repo_url = repo_url self.repo_url = repo_url
self.branch = branch self.branch = branch
self.github = GitHubClient(token=github_token) self.github = GitHubClient(token=github_token)
# We use gemini-2.5-flash for fast and cost-effective analysis # We use gemini-2.0-flash for high rate limits and stability
self.model_name = 'gemini-2.5-flash' self.model_name = 'gemini-2.0-flash'
async def triage_files(self, all_files: List[str]) -> List[str]: async def triage_files(self, all_files: List[str]) -> List[str]:
""" """
@@ -30,7 +31,7 @@ class CodeScanOrchestrator:
""" """
if not settings.gemini_api_key: if not settings.gemini_api_key:
logger.warning("GEMINI_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] return all_files[:5]
files_str = "\n".join(all_files) files_str = "\n".join(all_files)
if len(files_str) > 15000: if len(files_str) > 15000:
@@ -41,7 +42,7 @@ class CodeScanOrchestrator:
f"{files_str}\n\n" f"{files_str}\n\n"
"Select the most critical files to review for security vulnerabilities (e.g., SAST, hardcoded secrets, SQLi, Auth bypass). " "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. " "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: try:
@@ -59,18 +60,23 @@ class CodeScanOrchestrator:
except Exception as e: except Exception as e:
logger.error(f"Error triaging files: {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]: async def analyze_files(self, triaged_files: List[str]) -> List[VulnerabilityIssue]:
vulnerabilities = []
if not settings.gemini_api_key: if not settings.gemini_api_key:
return [] 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) content = await self.github.get_file_content(self.repo_url, file_path, self.branch)
if not content: if not content:
continue return []
if len(content) > 30000: if len(content) > 30000:
content = content[:30000] content = content[:30000]
@@ -88,30 +94,37 @@ class CodeScanOrchestrator:
"'line_number' (integer or null if general)." "'line_number' (integer or null if general)."
) )
try: file_vulns = []
response = await ai_client.aio.models.generate_content( async with semaphore:
model=self.model_name, try:
contents=prompt, response = await ai_client.aio.models.generate_content(
config=types.GenerateContentConfig( model=self.model_name,
response_mime_type="application/json", contents=prompt,
temperature=0.2, config=types.GenerateContentConfig(
response_mime_type="application/json",
temperature=0.2,
)
) )
) if response.text:
if response.text: data = json.loads(response.text)
data = json.loads(response.text) vulns = data.get("vulnerabilities", [])
vulns = data.get("vulnerabilities", []) for v in vulns:
for v in vulns: file_vulns.append(VulnerabilityIssue(
vulnerabilities.append(VulnerabilityIssue( file_path=file_path,
file_path=file_path, severity=v.get("severity", "Medium"),
severity=v.get("severity", "Medium"), issue=v.get("issue", "Unknown Issue"),
issue=v.get("issue", "Unknown Issue"), explanation=v.get("explanation", ""),
explanation=v.get("explanation", ""), suggested_fix=v.get("suggested_fix"),
suggested_fix=v.get("suggested_fix"), line_number=v.get("line_number")
line_number=v.get("line_number") ))
)) except Exception as e:
except Exception as e: logger.error(f"Error analyzing file {file_path}: {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 return vulnerabilities
async def generate_summary(self, vulnerabilities: List[VulnerabilityIssue]) -> str: async def generate_summary(self, vulnerabilities: List[VulnerabilityIssue]) -> str:

11
list_models.py Normal file
View File

@@ -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}")

View File

@@ -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 AVAILABLE MODELS: migrated to new SDK