External/securelens-backend
1
0
Fork 0
mirror of https://github.com/Rarebuffalo/securelens-backend.git synced 2026-06-19 07:00:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

updated model for suitable testing

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

16
app/routers/code_scan.py
View File

@@ -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,

2
app/services/ai.py
View File

@@ -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:

75
app/services/code_scanner/orchestrator.py
View File

@@ -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:

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

2
models.txt
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