feat(alert-triage): integrate Gemini API and update env template

.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
 # ============================================================================

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

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

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