mirror of
https://github.com/Rarebuffalo/securelens-backend.git
synced 2026-06-19 07:00:30 +00:00
updated model for suitable testing
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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
11
list_models.py
Normal 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}")
|
||||||
@@ -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
|
||||||
Reference in New Issue
Block a user