refactor llm/embedding flow

app/config.py

@@ -35,6 +35,13 @@ class Settings(BaseSettings):
     gemini_base_url: str = os.getenv("GEMINI_BASE_URL", "https://generativelanguage.googleapis.com/v1/models/gemini-2.5-flash:generateContent") or ""
     gemini_model: str = os.getenv("GEMINI_MODEL", "gemini-2.5-flash") or ""
 
+    # LLM (chat/completion) provider/model
+    llm_provider: str = os.getenv("LLM_PROVIDER", "gemini") or ""
+    llm_model: str = os.getenv("LLM_MODEL", "gemini-1.5-flash-latest") or ""
+    # Embedding provider/model
+    embedding_provider: str = os.getenv("EMBEDDING_PROVIDER", "gemini") or ""
+    embedding_model: str = os.getenv("EMBEDDING_MODEL", "models/embedding-001") or ""
+
     class Config:
         env_file = ".env"
 

app/embedding.py

@@ -2,8 +2,10 @@ from typing import List
 import numpy as np
 from loguru import logger
 import httpx
-
+from .config import get_settings
 from .utils import timing_decorator_async, timing_decorator_sync, call_endpoint_with_retry
+from .llm import LLMClient
+from .gemini_client import GeminiClient
 
 class EmbeddingClient:
     def __init__(self):
@@ -13,14 +15,38 @@ class EmbeddingClient:
         Output: EmbeddingClient instance.
         """
         self._client = httpx.AsyncClient()
+        settings = get_settings()
+        self.provider = getattr(settings, 'embedding_provider', 'default')
+        self.model = getattr(settings, 'embedding_model', 'models/embedding-001')
+        if self.provider == 'gemini':
+            self.gemini_client = GeminiClient(settings.gemini_api_key, model=self.model)
+        else:
+            self.gemini_client = None
 
     @timing_decorator_async
     async def create_embedding(self, text: str) -> List[float]:
         """
-        Tạo embedding vector từ text bằng dịch vụ embedding (ví dụ OpenAI).
+        Tạo embedding vector từ text bằng dịch vụ embedding (ví dụ OpenAI hoặc Gemini).
         Input: text (str)
         Output: list[float] embedding vector.
         """
+        if self.provider == 'gemini' and self.gemini_client:
+            try:
+                # GeminiClient.create_embedding là hàm sync, chạy trong executor
+                import asyncio
+                loop = asyncio.get_event_loop()
+                embedding = await loop.run_in_executor(None, self.gemini_client.create_embedding, text)
+                # Kiểm tra kiểu dữ liệu trả về
+                if isinstance(embedding, list):
+                    preview = f"{embedding[:10]}...{embedding[-10:]}" if len(embedding) > 20 else str(embedding)
+                    logger.info(f"[DEBUG] Embedding API response: {preview}")
+                    return embedding
+                else:
+                    logger.error(f"[DEBUG] Unknown embedding type: {type(embedding)} - value: {embedding}")
+                    raise RuntimeError(f"Embedding returned unexpected type: {type(embedding)}")
+            except Exception as e:
+                logger.error(f"Error creating embedding with Gemini: {e}")
+                raise
         url = "https://vietcat-vietnameseembeddingv2.hf.space/embed"
         payload = {"text": text}
         try:

app/llm.py

@@ -4,6 +4,8 @@ import json
 from loguru import logger
 from tenacity import retry, stop_after_attempt, wait_exponential
 import os
+from .gemini_client import GeminiClient
+from .config import get_settings
 
 from .utils import timing_decorator_async, timing_decorator_sync, call_endpoint_with_retry
 
@@ -87,10 +89,8 @@ class LLMClient:
     def _setup_gemini(self, config: Dict[str, Any]):
         """Cấu hình cho Gemini."""
         self.api_key = config.get("api_key", "")
-        self.base_url = config.get("base_url", "")
-        self.model = config.get("model", "")
-        self.max_tokens = config.get("max_tokens", 1024)
-        self.temperature = config.get("temperature", 0.7)
+        self.model = config.get("model", "gemini-1.5-flash-latest")
+        self.gemini_client = GeminiClient(self.api_key, self.model)
     
     @timing_decorator_async
     async def generate_text(
@@ -205,32 +205,9 @@ class LLMClient:
             raise RuntimeError("HFS API response is None")
 
     async def _generate_gemini(self, prompt: str, **kwargs) -> str:
-        """Gọi Gemini API để sinh text từ prompt."""
-        url = self.base_url
-        headers = {"Content-Type": "application/json"}
-        if self.api_key:
-            headers["X-Goog-Api-Key"] = f"{self.api_key}"
-        # Gemini API expects {"contents": [{"parts": [{"text": prompt}]}]}
-        payload = {"contents": [{"parts": [{"text": prompt}]}]}
-        response = await call_endpoint_with_retry(self._client, url, payload, headers=headers)
-        if response is not None and hasattr(response, 'text'):
-            logger.info(f"[LLM][GEMINI][RAW_RESPONSE] {response.text}")
-        else:
-            logger.info(f"[LLM][GEMINI][RAW_RESPONSE] {str(response)}")
-        if response is not None:
-            data = response.json()
-            # Log token usage nếu có
-            usage = data.get('usage') or data.get('usageMetadata')
-            if usage:
-                logger.info(f"[LLM][GEMINI][USAGE] {usage}")
-            # Gemini trả về: {'candidates': [{'content': {'parts': [{'text': '...'}]}}]}
-            try:
-                return data['candidates'][0]['content']['parts'][0]['text']
-            except Exception:
-                return str(data)
-        else:
-            logger.error("Gemini API response is None")
-            raise RuntimeError("Gemini API response is None")
+        import asyncio
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, self.gemini_client.generate_text, prompt)
 
     @timing_decorator_async
     async def chat(
@@ -498,19 +475,25 @@ if __name__ == "__main__":
     
     async def test_llm():
         # Test với OpenAI
-        llm = create_llm_client("openai", model="gpt-3.5-turbo")
+        settings = get_settings()
+        llm_client = create_llm_client(
+            provider=settings.llm_provider,
+            model=settings.llm_model,
+            api_key=settings.gemini_api_key,
+            # ... các config khác nếu cần ...
+        )
         
         # Generate text
-        response = await llm.generate_text("Xin chào, bạn có khỏe không?")
+        response = await llm_client.generate_text("Xin chào, bạn có khỏe không?")
         print(f"Response: {response}")
         
         # Chat
         messages = [
             {"role": "user", "content": "Bạn có thể giúp tôi không?"}
         ]
-        chat_response = await llm.chat(messages)
+        chat_response = await llm_client.chat(messages)
         print(f"Chat response: {chat_response}")
         
-        await llm.close()
+        await llm_client.close()
     
     asyncio.run(test_llm()) 

requirements.txt

@@ -2,7 +2,7 @@ fastapi==0.104.1
 uvicorn==0.24.0
 python-dotenv==1.0.0
 httpx>=0.24.0,<0.25.0
-loguru==0.7.2
+loguru>=0.7.0
 google-auth==2.23.4
 google-auth-oauthlib==1.1.0
 google-auth-httplib2==0.1.1
@@ -11,4 +11,5 @@ supabase==2.0.3
 numpy==1.26.2
 python-multipart==0.0.6
 tenacity==8.2.3 
-pydantic-settings
+pydantic-settings
+google-generativeai>=0.3.0