app/llm.py

from typing import List, Dict, Any, Optional, Union
import httpx
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

class LLMClient:
    """
    Client để tương tác với các mô hình ngôn ngữ lớn (LLM).
    Hỗ trợ nhiều provider: OpenAI, HuggingFace, local models, etc.
    """
    
    def __init__(self, provider: str = "openai", **kwargs):
        """
        Khởi tạo LLMClient.
        
        Args:
            provider (str): Loại provider ("openai", "huggingface", "local", "custom")
            **kwargs: Các tham số cấu hình khác
        """
        self.provider = provider.lower()
        self._client = httpx.AsyncClient(timeout=60.0)
        
        # Cấu hình theo provider
        if self.provider == "openai":
            self._setup_openai(kwargs)
        elif self.provider == "huggingface":
            self._setup_huggingface(kwargs)
        elif self.provider == "local":
            self._setup_local(kwargs)
        elif self.provider == "custom":
            self._setup_custom(kwargs)
        elif self.provider == "hfs":
            self._setup_HFS(kwargs)
        elif self.provider == "gemini":
            self._setup_gemini(kwargs)
        else:
            raise ValueError(f"Unsupported provider: {provider}")
    
    def _setup_openai(self, config: Dict[str, Any]):
        """Cấu hình cho OpenAI."""
        self.api_key = config.get("api_key") or os.getenv("OPENAI_API_KEY") or ""
        self.base_url = config.get("base_url", "https://api.openai.com/v1")
        self.model = config.get("model", "gpt-3.5-turbo")
        self.max_tokens = config.get("max_tokens", 1000)
        self.temperature = config.get("temperature", 0.7)
        
        if not self.api_key:
            raise ValueError("OpenAI API key is required")
    
    def _setup_huggingface(self, config: Dict[str, Any]):
        """Cấu hình cho HuggingFace."""
        self.api_key = config.get("api_key", "")
        self.base_url = config.get("base_url", "https://api-inference.huggingface.co")
        self.model = config.get("model", "microsoft/DialoGPT-medium")
        self.max_tokens = config.get("max_tokens", 1000)
        self.temperature = config.get("temperature", 0.7)
    
    def _setup_local(self, config: Dict[str, Any]):
        """Cấu hình cho local model."""
        self.api_key = ""
        self.base_url = config.get("base_url", "http://localhost:8000")
        self.model = config.get("model", "default")
        self.max_tokens = config.get("max_tokens", 1000)
        self.temperature = config.get("temperature", 0.7)
    
    def _setup_custom(self, config: Dict[str, Any]):
        """Cấu hình cho custom provider."""
        self.api_key = config.get("api_key", "")
        self.base_url = config.get("base_url")
        self.model = config.get("model", "default")
        self.max_tokens = config.get("max_tokens", 1000)
        self.temperature = config.get("temperature", 0.7)
        if not self.base_url:
            raise ValueError("Custom provider requires base_url")

    def _setup_HFS(self, config: Dict[str, Any]):
        """Cấu hình cho custom provider."""
        self.api_key = config.get("api_key", "")
        self.base_url = config.get("base_url")
        if not self.base_url:
            raise ValueError("Custom provider requires base_url")
    
    def _setup_gemini(self, config: Dict[str, Any]):
        """Cấu hình cho Gemini."""
        # Sử dụng GeminiClient với RequestLimitManager
        self.gemini_client = GeminiClient()
        logger.info("[LLM] Initialized GeminiClient with RequestLimitManager")
    
    @timing_decorator_async
    async def generate_text(
        self, 
        prompt: str, 
        system_prompt: Optional[str] = None,
        **kwargs
    ) -> str:
        """
        Tạo text từ prompt sử dụng LLM.
        
        Args:
            prompt (str): Prompt đầu vào
            system_prompt (str, optional): System prompt
            **kwargs: Các tham số bổ sung
            
        Returns:
            str: Text được tạo ra
        """
        logger.info(f"[LLM] generate_text - provider: {self.provider} \n\t prompt: {prompt}")
        try:
            result = None
            if self.provider == "openai":
                result = await self._generate_openai(prompt, system_prompt, **kwargs)
            elif self.provider == "huggingface":
                result = await self._generate_huggingface(prompt, **kwargs)
            elif self.provider == "local":
                result = await self._generate_local(prompt, **kwargs)
            elif self.provider == "custom":
                result = await self._generate_custom(prompt, **kwargs)
            elif self.provider == "hfs":
                result = await self._generate_hfs(prompt, **kwargs)
            elif self.provider == "gemini":
                result = await self._generate_gemini(prompt, **kwargs)
            else:
                raise ValueError(f"Unsupported provider: {self.provider}")
            logger.info(f"[LLM] generate_text - provider: {self.provider}\n\t result: {result}")
            return result
        except Exception as e:
            logger.error(f"[LLM] Error generating text with {self.provider}: {e}")
            raise
    
    async def _generate_openai(self, prompt: str, system_prompt: Optional[str] = None, **kwargs) -> str:
        url = f"{self.base_url}/chat/completions"
        payload = {"model": kwargs.get("model", self.model), "messages": [{"role": "system", "content": system_prompt or ""}, {"role": "user", "content": prompt}], "max_tokens": kwargs.get("max_tokens", self.max_tokens), "temperature": kwargs.get("temperature", self.temperature), "stream": False}
        headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
        response = await call_endpoint_with_retry(self._client, url, payload, headers=headers)
        if response is not None:
            data = response.json()
            return data["choices"][0]["message"]["content"]
        else:
            logger.error("OpenAI API response is None")
            raise RuntimeError("OpenAI API response is None")
    
    async def _generate_huggingface(self, prompt: str, **kwargs) -> str:
        url = f"{self.base_url}/generate"
        payload = {"inputs": prompt}
        response = await call_endpoint_with_retry(self._client, url, payload)
        if response is not None:
            data = response.json()
            return data[0]["generated_text"]
        else:
            logger.error("HuggingFace API response is None")
            raise RuntimeError("HuggingFace API response is None")
    
    async def _generate_local(self, prompt: str, **kwargs) -> str:
        url = f"{self.base_url}/generate"
        payload = {"prompt": prompt}
        response = await call_endpoint_with_retry(self._client, url, payload)
        if response is not None:
            data = response.json()
            return data.get("text", "")
        else:
            logger.error("Local API response is None")
            raise RuntimeError("Local API response is None")
    
    async def _generate_custom(self, prompt: str, **kwargs) -> str:
        url = f"{self.base_url}/custom"
        payload = {"prompt": prompt}
        response = await call_endpoint_with_retry(self._client, url, payload)
        if response is not None:
            data = response.json()
            return data.get("text", "")
        else:
            logger.error("Custom API response is None")
            raise RuntimeError("Custom API response is None")
    
    async def _generate_hfs(self, prompt: str, **kwargs) -> str:
        endpoint = f"{self.base_url}/purechat"
        payload = {"prompt": prompt}
        headers = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        response = await call_endpoint_with_retry(self._client, endpoint, payload, 3, 500, headers=headers)
        logger.info(f"[LLM] generate_text - provider: {self.provider}\n\t response: {response}")
        try:
            import json as _json
            logger.info(f"[LLM][RAW_RESPONSE] { _json.dumps(response, ensure_ascii=False, indent=2) }")
        except Exception:
            logger.info(f"[LLM][RAW_RESPONSE] {str(response)}")
        if response is not None:
            data = response.json()
            if 'response' in data:
                return data['response']
            elif 'result' in data:
                return data['result']
            elif 'data' in data and isinstance(data['data'], list):
                return data['data'][0]
            return str(data)
        else:
            logger.error("HFS API response is None")
            raise RuntimeError("HFS API response is None")

    async def _generate_gemini(self, prompt: str, **kwargs) -> str:
        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(
        self, 
        messages: List[Dict[str, str]], 
        **kwargs
    ) -> str:
        """
        Chat với LLM sử dụng conversation history.
        
        Args:
            messages (List[Dict]): List các message với format [{"role": "user", "content": "..."}]
            **kwargs: Các tham số bổ sung
            
        Returns:
            str: Response từ LLM
        """
        logger.info(f"[LLM] chat - messages: {messages}")
        if self.provider == "openai":
            return await self._chat_openai(messages, **kwargs)
        else:
            # Với các provider khác, convert messages thành prompt
            prompt = self._messages_to_prompt(messages)
            return await self.generate_text(prompt, **kwargs)
    
    async def _chat_openai(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """Chat với OpenAI API."""
        payload = {
            "model": kwargs.get("model", self.model),
            "messages": messages,
            "max_tokens": kwargs.get("max_tokens", self.max_tokens),
            "temperature": kwargs.get("temperature", self.temperature),
            "stream": False
        }
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        response = await self._client.post(
            f"{self.base_url}/chat/completions",
            headers=headers,
            json=payload
        )
        response.raise_for_status()
        
        data = response.json()
        return data["choices"][0]["message"]["content"]
    
    def _messages_to_prompt(self, messages: List[Dict[str, str]]) -> str:
        """Convert conversation messages thành prompt string."""
        prompt = ""
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            
            if role == "system":
                prompt += f"System: {content}\n\n"
            elif role == "user":
                prompt += f"User: {content}\n"
            elif role == "assistant":
                prompt += f"Assistant: {content}\n"
        
        prompt += "Assistant: "
        return prompt
    
    @timing_decorator_async
    async def classify_text(
        self, 
        text: str, 
        categories: List[str],
        **kwargs
    ) -> Dict[str, Any]:
        """
        Phân loại text vào các categories.
        
        Args:
            text (str): Text cần phân loại
            categories (List[str]): List các categories
            **kwargs: Các tham số bổ sung
            
        Returns:
            Dict: Kết quả phân loại
        """
        prompt = f"""
        Phân loại text sau vào một trong các categories: {', '.join(categories)}
        
        Text: {text}
        
        Trả về kết quả theo format JSON:
        {{
            "category": "tên_category",
            "confidence": 0.95,
            "reasoning": "lý do phân loại"
        }}
        """
        
        response = await self.generate_text(prompt, **kwargs)
        
        try:
            # Tìm JSON trong response
            import re
            json_match = re.search(r'\{.*\}', response, re.DOTALL)
            if json_match:
                result = json.loads(json_match.group())
                return result
            else:
                return {
                    "category": "unknown",
                    "confidence": 0.0,
                    "reasoning": "Không thể parse JSON response"
                }
        except json.JSONDecodeError:
            return {
                "category": "unknown", 
                "confidence": 0.0,
                "reasoning": f"JSON parse error: {response}"
            }
    
    @timing_decorator_async
    async def extract_entities(
        self, 
        text: str,
        entity_types: Optional[List[str]] = None,
        **kwargs
    ) -> List[Dict[str, Any]]:
        """
        Trích xuất entities từ text.
        
        Args:
            text (str): Text cần trích xuất
            entity_types (List[str]): Các loại entity cần tìm
            **kwargs: Các tham số bổ sung
            
        Returns:
            List[Dict]: List các entities được tìm thấy
        """
        if entity_types is None:
            entity_types = ["PERSON", "ORGANIZATION", "LOCATION", "MONEY", "DATE"]
        
        prompt = f"""
        Trích xuất các entities từ text sau. Tìm các entities thuộc types: {', '.join(entity_types)}
        
        Text: {text}
        
        Trả về kết quả theo format JSON:
        [
            {{
                "text": "tên entity",
                "type": "loại entity",
                "start": 0,
                "end": 10
            }}
        ]
        """
        
        response = await self.generate_text(prompt, **kwargs)
        
        try:
            import re
            # Log toàn bộ response dưới dạng text dễ đọc
            try:
                import json as _json
                logger.info(f"[LLM][RAW_RESPONSE] { _json.dumps(response, ensure_ascii=False, indent=2) }")
            except Exception:
                logger.info(f"[LLM][RAW_RESPONSE] {str(response)}")
            # Ưu tiên parse object JSON nếu có
            json_match_obj = re.search(r'\{[\s\S]*?\}', response)
            json_match_list = re.search(r'\[[\s\S]*?\]', response)
            if json_match_list:
                entities = json.loads(json_match_list.group())
                return entities
            elif json_match_obj:
                entity = json.loads(json_match_obj.group())
                return [entity]
            else:
                return []
        except json.JSONDecodeError:
            logger.error(f"Error parsing entities JSON: {response}")
            return []

    @timing_decorator_async
    async def analyze(
        self, 
        text: str,
        **kwargs
    ) -> List[Dict[str, Any]]:
        """
        Trích xuất entities từ text.
        
        Args:
            text (str): Text cần trích xuất
            **kwargs: Các tham số bổ sung
            
        Returns:
            List[Dict]: List các entities được tìm thấy
        """
                
        prompt = f"""
            Bạn là một AI chuyên phân tích ngữ nghĩa câu hỏi về luật giao thông. Với mỗi câu đầu vào, hãy trích xuất các thông tin sau dưới dạng JSON:

            {{
                "muc_dich": "mục đích của câu hỏi, chỉ chọn một trong các giá trị: "hỏi về mức phạt", "hỏi về quy tắc giao thông", "hỏi về báo hiệu đường bộ", "hỏi về quy trình xử lý vi phạm giao thông", "thông tin cá nhân của AI", 'khác'",
                "phuong_tien": "loại phương tiện giao thông (xe máy, ô tô, xe tải, người đi bộ...)",
                "hanh_vi_vi_pham": "hành vi vi phạm pháp luật giao thông"
            }}

            Nếu một trường không có trong câu hỏi, hãy để giá trị là chuỗi rỗng "". Ví dụ: nếu câu không đề cập phương tiện, thì "phuong_tien": "".

            Lưu ý: Trả về đúng format JSON, và giá trị của "muc_dich" chỉ được chọn một trong các giá trị sau: "hỏi về mức phạt", "hỏi về quy tắc giao thông", "hỏi về báo hiệu đường bộ", "hỏi về quy trình xử lý vi phạm giao thông", "thông tin cá nhân của AI", "khác".

            Câu hỏi:
            \"{text}\"
            """.strip()
        
        response = await self.generate_text(prompt, **kwargs)
        
        logger.info(f"[LLM][RAW] Kết quả trả về từ generate_text: {response}")

        try:
            import re
            # Log toàn bộ response dưới dạng text dễ đọc
            try:
                import json as _json
                logger.info(f"[LLM][RAW_RESPONSE] { _json.dumps(response, ensure_ascii=False, indent=2) }")
            except Exception:
                logger.info(f"[LLM][RAW_RESPONSE] {str(response)}")
            # Ưu tiên parse object JSON nếu có
            json_match_obj = re.search(r'\{[\s\S]*?\}', response)
            json_match_list = re.search(r'\[[\s\S]*?\]', response)
            if json_match_list:
                entities = json.loads(json_match_list.group())
                return entities
            elif json_match_obj:
                entity = json.loads(json_match_obj.group())
                return [entity]
            else:
                return []
        except json.JSONDecodeError:
            logger.error(f"Error parsing entities JSON: {response}")
            return []
    
    async def close(self):
        """Đóng client connection."""
        await self._client.aclose()


# Factory function để tạo LLMClient dễ dàng
def create_llm_client(provider: str = "openai", **kwargs) -> LLMClient:
    """
    Factory function để tạo LLMClient.
    
    Args:
        provider (str): Loại provider
        **kwargs: Các tham số cấu hình
        
    Returns:
        LLMClient: Instance của LLMClient
    """
    return LLMClient(provider, **kwargs)


# Ví dụ sử dụng
if __name__ == "__main__":
    import asyncio
    
    async def test_llm():
        # Test với OpenAI
        settings = get_settings()
        llm_client = create_llm_client(
            provider=settings.llm_provider,
            model=settings.llm_model,
            # ... các config khác nếu cần ...
        )
        
        # Generate text
        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_client.chat(messages)
        print(f"Chat response: {chat_response}")
        
        await llm_client.close()
    
    asyncio.run(test_llm())