Merge pull request #48 from DUT-Team-21TCLC-DT3/feat/tester_ai

Feat/tester ai

Author: thanhNgan13

ai_service/app/config.py

@@ -28,6 +28,20 @@ def base_url(self) -> str:
 
     # LLMs
     gemini_api_key: Optional[str] = Field(None, alias="GEMINI_API_KEY")
+    gemini_api_keys_str: Optional[str] = Field(None, alias="GEMINI_API_KEYS")
+    
+    @property
+    def gemini_api_keys(self) -> list[str]:
+        """
+        Returns a list of API keys. 
+        Prioritizes GEMINI_API_KEYS (comma separated), falls back to GEMINI_API_KEY.
+        """
+        if self.gemini_api_keys_str:
+            return [k.strip() for k in self.gemini_api_keys_str.split(",") if k.strip()]
+        if self.gemini_api_key:
+            return [self.gemini_api_key]
+        return []
+
     cypher_model: str = Field("gpt-4o-mini", alias="CYPHER_MODEL")
     qa_model: str = Field("gpt-4o-mini", alias="QA_MODEL")
     cypher_temperature: float = Field(0.0, alias="CYPHER_TEMPERATURE")

ai_service/app/gemini_constants.py

@@ -0,0 +1,6 @@
+from app.dependencies import get_settings
+
+settings = get_settings()
+
+GEMINI_MODEL = "gemini-2.5-flash"
+GEMINI_API_KEYS = settings.gemini_api_keys

ai_service/app/pipelines/answer_composer.py

@@ -1,18 +1,15 @@
 import logging
-import google.generativeai as genai
 from typing import Dict, Any, List
 from ..dependencies import get_settings
 from ..utils.stream_utils import create_metadata_chunk
+from ..utils.gemini_utils import GeminiClient
 
 log = logging.getLogger(__name__)
 
 class AnswerComposerService:
     def __init__(self):
         settings = get_settings()
-        if not settings.gemini_api_key:
-             raise ValueError("GEMINI_API_KEY not found")
-        genai.configure(api_key=settings.gemini_api_key)
-        self.model_name = 'gemini-2.5-pro'
+        self.client = GeminiClient(generation_config={"temperature": 0.2})
         self.base_url = settings.base_url  # Get dynamic base URL
 
     def _build_citation_from_node_id(self, node_id: str) -> str:
@@ -191,13 +188,8 @@ def compose(
 
 HÃY VIẾT CÂU TRẢ LỜI:
 """
-        model = genai.GenerativeModel(
-            self.model_name,
-            generation_config={"temperature": 0.2} # Hơi sáng tạo một chút để viết văn mượt mà
-        )
-
         try:
-            response = model.generate_content(prompt)
+            response = self.client.generate_content(prompt)
             return response.text.strip()
         except Exception as e:
             log.error(f"Lỗi compose answer: {e}")
@@ -349,14 +341,9 @@ def compose_stream(
 
 HÃY VIẾT CÂU TRẢ LỜI:
 """
-        model = genai.GenerativeModel(
-            self.model_name,
-            generation_config={"temperature": 0.2}
-        )
-
         try:
             # Stream content directly, LLM includes citations in answer
-            response = model.generate_content(prompt, stream=True)
+            response = self.client.generate_content_stream(prompt)
             for chunk in response:
                 if chunk.text:
                     yield chunk.text

ai_service/app/pipelines/graph_search.py

@@ -1,12 +1,12 @@
 import logging
 import json
 import re
-import google.generativeai as genai
 from typing import Dict, Any, List, Optional
 from dataclasses import dataclass
 from neo4j import GraphDatabase
 from langchain_core.prompts import PromptTemplate
 from ..dependencies import get_settings
+from ..utils.gemini_utils import GeminiClient
 
 log = logging.getLogger(__name__)
 
@@ -216,15 +216,7 @@ def __init__(self):
         self.driver = GraphSearchService._driver
 
         # Configure Gemini
-        if not settings.gemini_api_key:
-            raise ValueError("GEMINI_API_KEY not found")
-        genai.configure(api_key=settings.gemini_api_key)
-
-        self.gemini_model_name = 'gemini-2.5-pro'
-        self.model = genai.GenerativeModel(
-            self.gemini_model_name,
-            generation_config={"temperature": 0.0},
-        )
+        self.client = GeminiClient(generation_config={"temperature": 0.0})
 
         # Build schema snapshot if not exists
         if GraphSearchService._schema_snapshot is None:
@@ -363,11 +355,7 @@ def run_direct_graph_query(
 """
 
         try:
-            model = genai.GenerativeModel(
-                self.gemini_model_name,
-                generation_config={"temperature": 0.0},
-            )
-            response = model.generate_content(prompt)
+            response = self.client.generate_content(prompt)
             cypher_query = self._clean(response.text)
 
             log.info(f"Generated direct Cypher query: {cypher_query}...")
@@ -537,7 +525,7 @@ def generate_cypher(self, question: str, node_context: Dict[str, Any]) -> Cypher
         )
 
         try:
-            response = self.model.generate_content(prompt)
+            response = self.client.generate_content(prompt)
             cypher_query = self._clean(response.text)
             reasoning = f"Graph traversal từ node {node_id}"
 
@@ -693,11 +681,7 @@ def decide_context(self, question: str, node_text: str) -> Dict[str, str]:
   "reason": "giải thích ngắn gọn, tiếng Việt"
 }}
 """
-        model = genai.GenerativeModel(
-            self.gemini_model_name,
-            generation_config={"temperature": 0.1},
-        )
-        resp = model.generate_content(prompt)
+        resp = self.client.generate_content(prompt)
         try:
             text = resp.text.strip()
             if text.startswith("```"):
@@ -781,11 +765,7 @@ def filter_relevant_nodes(
 """
 
         try:
-            model = genai.GenerativeModel(
-                self.gemini_model_name,
-                generation_config={"temperature": 0.0},
-            )
-            response = model.generate_content(prompt)
+            response = self.client.generate_content(prompt)
             raw_text = response.text.strip()
 
             if raw_text.startswith("```"):
@@ -1143,7 +1123,7 @@ def run_graph_search_for_node_ids(
 
         # Step 6: Final Gemini reranking to select best candidate
         log.info("Step 6: Final Gemini reranking...")
-        rerank_model = genai.GenerativeModel(self.gemini_model_name)
+        # rerank_model = genai.GenerativeModel(self.gemini_model_name) # Removed
 
         numbered_candidates = []
         for idx, c in enumerate(candidates_for_llm, start=1):
@@ -1187,7 +1167,7 @@ def run_graph_search_for_node_ids(
 """
 
         try:
-            response = rerank_model.generate_content(selection_prompt)
+            response = self.client.generate_content(selection_prompt)
             raw_text = response.text.strip()
             if raw_text.startswith("```"):
                 raw_text = self._clean(raw_text)

ai_service/app/pipelines/query_processor.py

@@ -1,8 +1,8 @@
 import logging
 import json
-import google.generativeai as genai
 from typing import Dict, Any
 from ..dependencies import get_settings
+from ..utils.gemini_utils import GeminiClient
 
 log = logging.getLogger(__name__)
 
@@ -129,28 +129,19 @@
 
 class QueryPreprocessor:
     def __init__(self):
-        settings = get_settings()
-        api_key = settings.gemini_api_key
-        if not api_key:
-            raise ValueError("GEMINI_API_KEY not found in settings")
-        
-        genai.configure(api_key=api_key)
-        
         # Cấu hình Model
         generation_config = {
             "response_mime_type": "application/json",
             "temperature": 0.0, 
         }
 
         # Khởi tạo model với System Instruction riêng biệt
-        gemini_model = 'gemini-2.5-flash'
-        self.model = genai.GenerativeModel(
-            gemini_model, 
-            generation_config=generation_config,
-            system_instruction=SYSTEM_INSTRUCTION
+        self.client = GeminiClient(
+            system_instruction=SYSTEM_INSTRUCTION,
+            generation_config=generation_config
         )
 
-        log.info(f"QueryRewriter initialized successfully ({gemini_model})")
+        log.info(f"QueryRewriter initialized successfully")
 
     def _clean_json_string(self, text: str) -> str:
         """Làm sạch chuỗi JSON nếu model trả về markdown"""
@@ -172,7 +163,7 @@ def rewrite(self, question: str) -> Dict[str, Any]:
         try:
             user_prompt = f"Phân tích câu hỏi sau: \"{question}\""
 
-            response = self.model.generate_content(user_prompt)
+            response = self.client.generate_content(user_prompt)
             json_text = self._clean_json_string(response.text)
 
             try:

ai_service/app/pipelines/reranker.py

@@ -1,18 +1,16 @@
 import logging
 import json
-import google.generativeai as genai
 from typing import List, Dict, Any
 from ..dependencies import get_settings
+from ..utils.gemini_utils import GeminiClient
 
 log = logging.getLogger(__name__)
 
 class RerankerService:
     def __init__(self):
-        settings = get_settings()
-        if not settings.gemini_api_key:
-             raise ValueError("GEMINI_API_KEY not found")
-        genai.configure(api_key=settings.gemini_api_key)
-        self.model_name = 'gemini-2.5-pro' # Or settings.GEMINI_RERANK_MODEL
+        # settings = get_settings()
+        # Initializing GeminiClient for potential future use (reranking currently disabled)
+        self.client = GeminiClient()
 
     def _clean(self, text: str) -> str:
         cleaned = text.strip()

ai_service/app/pipelines/social_chat.py

@@ -1,26 +1,19 @@
 import logging
-import google.generativeai as genai
 from ..dependencies import get_settings
+from ..utils.gemini_utils import GeminiClient
 
 log = logging.getLogger(__name__)
 
 class SocialChatService:
     def __init__(self):
-        settings = get_settings()
-        self.api_key = settings.gemini_api_key
-        self.model_name = 'gemini-2.5-flash'
-
-        if not self.api_key:
-            log.warning("⚠️ GEMINI_API_KEY not found. Social Chat will fail.")
-            self.client = None
-        else:
-            try:
-                genai.configure(api_key=self.api_key)
-                self.client = genai.GenerativeModel(self.model_name)
-                log.info("Social Chat Service (Gemini Client) initialized.")
-            except Exception as e:
-                log.error(f"Failed to init Gemini Client: {e}")
-                self.client = None
+        # settings = get_settings()
+        try:
+             # Initialize GeminiClient with higher temperature for social chat
+             self.client = GeminiClient(generation_config={"temperature": 0.7})
+             log.info("Social Chat Service (Gemini Client) initialized.")
+        except Exception as e:
+             log.error(f"Failed to init Gemini Client: {e}")
+             self.client = None
 
     def chat(self, intent: str, question: str):
         """
@@ -55,12 +48,7 @@ def chat(self, intent: str, question: str):
 
 HÃY VIẾT CÂU TRẢ LỜI SÁNG TẠO VÀ PHÙ HỢP NHẤT:
 """
-            model = genai.GenerativeModel(
-                self.model_name,
-                generation_config={"temperature": 0.7}  # Hơi sáng tạo để trả lời tự nhiên hơn
-            )
-
-            response = model.generate_content(prompt, stream=True)
+            response = self.client.generate_content_stream(prompt)
             for chunk in response:
                 if chunk.text:
                     yield chunk.text