From 139c8d982bc4c1729f2336a1b7e78db2a1da0181 Mon Sep 17 00:00:00 2001 From: rarebuffalo Date: Sat, 25 Apr 2026 20:49:17 +0530 Subject: [PATCH] updated the model --- .env.example | 2 +- app/config.py | 2 +- app/routers/__init__.py | 18 +--- app/routers/code_scan.py | 33 ++++--- app/routers/scan.py | 2 +- app/services/ai.py | 105 +++++++++++----------- app/services/code_scanner/orchestrator.py | 92 +++++++++---------- models.txt | 1 + requirements.txt | 2 +- 9 files changed, 120 insertions(+), 137 deletions(-) create mode 100644 models.txt diff --git a/.env.example b/.env.example index cad8b80..eb251a5 100644 --- a/.env.example +++ b/.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= diff --git a/app/config.py b/app/config.py index ac17e7c..7246f91 100644 --- a/app/config.py +++ b/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") diff --git a/app/routers/__init__.py b/app/routers/__init__.py index 23d41e2..032dabf 100644 --- a/app/routers/__init__.py +++ b/app/routers/__init__.py @@ -1,17 +1 @@ -from .auth import router as auth -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" -] +# Empty init file to allow correct module importing diff --git a/app/routers/code_scan.py b/app/routers/code_scan.py index c4a9057..0331dfd 100644 --- a/app/routers/code_scan.py +++ b/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: - raise HTTPException(status_code=400, detail="AI Chat is disabled because OPENAI_API_KEY is not configured.") + if not settings.gemini_api_key: + 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( - model="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": request.message} - ], - temperature=0.5, + from google import genai + from google.genai import types + + client = genai.Client(api_key=settings.gemini_api_key) + response = await client.aio.models.generate_content( + model='gemini-2.5-flash', + 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)}") diff --git a/app/routers/scan.py b/app/routers/scan.py index a461eef..b4a8d1d 100644 --- a/app/routers/scan.py +++ b/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", []) diff --git a/app/services/ai.py b/app/services/ai.py index 6b93b78..8c9bc4b 100644 --- a/app/services/ai.py +++ b/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" -client = AsyncOpenAI(api_key=api_key) +if settings.gemini_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: - """ - Takes a list of basic security issues and uses an LLM to provide: - - 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.") + if not settings.gemini_api_key: + logger.warning("GEMINI_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="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": "You are a senior cybersecurity automation agent. Always respond with valid JSON."}, - {"role": "user", "content": prompt} - ], - response_format={"type": "json_object"}, - temperature=0.2, + model_name = await get_gemini_model() + response = await ai_client.aio.models.generate_content( + model=model_name, + contents=prompt, + config=types.GenerateContentConfig( + response_mime_type="application/json", + temperature=0.2, + ) ) - content = response.choices[0].message.content - if not content: - return {"enhanced_issues": issues, "ai_error": "Empty response"} - - return json.loads(content) + if response.text: + return json.loads(response.text) + return {"enhanced_issues": issues, "ai_error": "Empty response"} 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: - """ - Allows a user to ask a question about a specific scan's results. - """ - if not settings.openai_api_key: - return "AI Chat is disabled because OPENAI_API_KEY is not configured." + if not settings.gemini_api_key: + return "AI Chat is disabled because GEMINI_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="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_message} - ], - temperature=0.5, + model_name = await get_gemini_model() + response = await ai_client.aio.models.generate_content( + model=model_name, + contents=prompt, + config=types.GenerateContentConfig( + 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: - """ - Weaves multiple scan issues into a cohesive attack sequence. - """ - if not settings.openai_api_key: - return "AI Threat Narrative is disabled because OPENAI_API_KEY is not configured." + if not settings.gemini_api_key: + return "AI Threat Narrative is disabled because GEMINI_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="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": json.dumps(context_data)} - ], - temperature=0.7, + model_name = await get_gemini_model() + response = await ai_client.aio.models.generate_content( + model=model_name, + contents=prompt, + config=types.GenerateContentConfig( + 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." diff --git a/app/services/code_scanner/orchestrator.py b/app/services/code_scanner/orchestrator.py index 4c297ce..c26d9ea 100644 --- a/app/services/code_scanner/orchestrator.py +++ b/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" -client = AsyncOpenAI(api_key=api_key) +if settings.gemini_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: - logger.warning("OPENAI_API_KEY is not set. Triaging all files up to a limit.") - return all_files[:10] # Hard limit for testing + 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] - # 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( - model="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": "You always respond with valid JSON."}, - {"role": "user", "content": prompt} - ], - response_format={"type": "json_object"}, - temperature=0.1, + response = await ai_client.aio.models.generate_content( + model=self.model_name, + contents=prompt, + config=types.GenerateContentConfig( + response_mime_type="application/json", + temperature=0.1, + ) ) - content = response.choices[0].message.content - if content: - data = json.loads(content) + if response.text: + data = json.loads(response.text) 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) > 20000: - content = content[:20000] + if len(content) > 30000: + content = content[:30000] 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( - model="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": "You are a SAST security agent. Always respond with valid JSON."}, - {"role": "user", "content": prompt} - ], - response_format={"type": "json_object"}, - temperature=0.2, + 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, + ) ) - - resp_content = response.choices[0].message.content - if resp_content: - data = json.loads(resp_content) + if response.text: + data = json.loads(response.text) 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( - model="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": "You are a cybersecurity expert."}, - {"role": "user", "content": prompt} - ], - temperature=0.4, + response = await ai_client.aio.models.generate_content( + model=self.model_name, + contents=prompt, + config=types.GenerateContentConfig( + 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." diff --git a/models.txt b/models.txt new file mode 100644 index 0000000..eeb0611 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index f59b38b..ebe8ffb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,6 +14,6 @@ pydantic[email] pytest pytest-asyncio alembic -openai +google-genai aiodns fpdf2