From 0c38fe199e02ee0b2fe0526314e0231ac2c5c4e9 Mon Sep 17 00:00:00 2001 From: z5w0rd Date: Wed, 9 Sep 2026 14:56:57 +0700 Subject: [PATCH 1/2] feat(alert-triage): integrate Gemini API and update env template --- .env.example | 11 +++++++ services/alert-triage/config.py | 1 + services/alert-triage/llm_client.py | 46 +++++++++++++++++++---------- 3 files changed, 42 insertions(+), 16 deletions(-) diff --git a/.env.example b/.env.example index 95de8d5..1997102 100644 --- a/.env.example +++ b/.env.example @@ -236,5 +236,16 @@ BACKUP_SCHEDULE=0 2 * * * # - Python: python -c "import secrets; print(secrets.token_urlsafe(32))" # # ============================================================================ +# ============================================================================ +# AI LLM ENGINE / ALERT TRIAGE CONFIGURATION +# ============================================================================ +# Google Gemini API Key for alert triage and analysis +# Generate key at: https://aistudio.google.com/ +GEMINI_API_KEY=your_gemini_api_key_here + +# Local Ollama Host (fallback if Gemini is not configured) +TRIAGE_OLLAMA_HOST=http://ollama:11434 +TRIAGE_PRIMARY_MODEL=foundation-sec-8b +TRIAGE_FALLBACK_MODEL=llama3.1:8b # END OF CONFIGURATION # ============================================================================ diff --git a/services/alert-triage/config.py b/services/alert-triage/config.py index a86d72a..d072a4c 100644 --- a/services/alert-triage/config.py +++ b/services/alert-triage/config.py @@ -83,6 +83,7 @@ class Config: env_prefix = "TRIAGE_" env_file = ".env" case_sensitive = False + extra = "ignore" # Global settings instance diff --git a/services/alert-triage/llm_client.py b/services/alert-triage/llm_client.py index d45b911..e52bf8d 100644 --- a/services/alert-triage/llm_client.py +++ b/services/alert-triage/llm_client.py @@ -8,13 +8,20 @@ import json import logging +import os from typing import Optional, Dict, Any +from dotenv import load_dotenv +from google import genai import httpx + from config import settings from models import SecurityAlert, TriageResponse, SeverityLevel, AlertCategory, IOC, TriageRecommendation from ml_client import MLInferenceClient, MLPrediction, enrich_llm_prompt_with_ml from context_manager import ContextManager +# Nạp các biến từ file .env vào os.environ +load_dotenv() + logger = logging.getLogger(__name__) @@ -83,14 +90,17 @@ def __init__(self): history_limit=settings.context_history_limit, environment_context=settings.environment_context, ) + # Gemini Client Init + api_key = os.getenv("GEMINI_API_KEY") + self.gemini_client = genai.Client(api_key=api_key) if api_key else None + if self.gemini_client: + logger.info("Gemini client initialized successfully.") async def check_health(self) -> bool: - """ - Check if Ollama service is reachable. + """Check if LLM service is reachable.""" + if self.gemini_client: + return True - Returns: - bool: True if Ollama is available - """ try: async with httpx.AsyncClient(timeout=5.0) as client: response = await client.get(f"{self.base_url}/api/tags") @@ -191,17 +201,21 @@ async def _call_ollama( model: str, temperature: float = 0.1 ) -> Optional[str]: - """ - Make API call to Ollama. - - Args: - prompt: Text prompt - model: Model identifier - temperature: Sampling temperature + """Make API call to Gemini (fallback to Ollama if no key).""" + # 1. Ưu tiên gọi Gemini API nếu đã cấu hình key + if self.gemini_client: + try: + logger.info("Calling Gemini API...") + response = await self.gemini_client.aio.models.generate_content( + model="gemini-3.6-flash", + contents=prompt, + ) + return response.text + except Exception as e: + logger.error(f"Gemini API call failed: {e}") + return None - Returns: - Optional[str]: Model response or None on error - """ + # 2. Fallback sang Ollama local nếu không có Gemini Client try: async with httpx.AsyncClient(timeout=self.timeout) as client: payload = { @@ -212,7 +226,7 @@ async def _call_ollama( "temperature": temperature, "num_predict": settings.max_tokens, }, - "format": "json" # Request JSON output + "format": "json" } logger.info(f"Calling Ollama model: {model}") From 9fc1f834f8ef974ce84d8063d818980e7024dcc0 Mon Sep 17 00:00:00 2001 From: z5w0rd Date: Thu, 10 Sep 2026 10:19:40 +0700 Subject: [PATCH 2/2] docs(alert-triage): add gemini config templates to .env.example --- services/alert-triage/.env.example | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/services/alert-triage/.env.example b/services/alert-triage/.env.example index eaacdd6..82ff6d8 100644 --- a/services/alert-triage/.env.example +++ b/services/alert-triage/.env.example @@ -6,7 +6,11 @@ TRIAGE_SERVICE_NAME=alert-triage TRIAGE_SERVICE_VERSION=1.0.0 TRIAGE_LOG_LEVEL=INFO -# Ollama LLM Configuration +# Gemini LLM Configuration (Primary / Fallback) +TRIAGE_GEMINI_API_KEY=your_gemini_api_key_here +TRIAGE_GEMINI_MODEL=gemini-1.5-flash + +# Ollama LLM Configuration (Local) TRIAGE_OLLAMA_HOST=http://ollama:11434 TRIAGE_PRIMARY_MODEL=foundation-sec-8b TRIAGE_FALLBACK_MODEL=llama3.1:8b @@ -22,7 +26,7 @@ TRIAGE_MEDIUM_CONFIDENCE_THRESHOLD=0.70 TRIAGE_AUTO_ACTION_THRESHOLD=0.80 # ML Inference Configuration -TRIAGE_ML_ENABLED=true +TRIAGE_ML_ENABLED=false TRIAGE_ML_API_URL=http://ml-inference:8001 TRIAGE_ML_TIMEOUT=10 TRIAGE_ML_DEFAULT_MODEL=random_forest @@ -46,4 +50,4 @@ TRIAGE_REQUEST_TIMEOUT=120 # Security (Production) TRIAGE_API_KEY_ENABLED=false -# TRIAGE_API_KEY=your_secure_api_key_here +# TRIAGE_API_KEY=your_secure_api_key_here \ No newline at end of file