backend/models.py

# models.py

import os
import ast
import re
import logging
import json
import asyncio
from typing import List, Dict, Any, Optional, Union, Tuple, AsyncGenerator
from dotenv import load_dotenv
from openai import AsyncOpenAI, RateLimitError, APIError, OpenAI
# from sentence_transformers import SentenceTransformer
from langfuse.decorators import langfuse_context, observe
from tools import TOOL_DEFINITIONS, execute_tool

from systemprompt import (
    get_rag_classification_prompt,
    get_subquery_prompt,
    get_normal_prompt,
    get_non_rag_prompt,
)
from utils import get_device
if get_device() == "mps":
    load_dotenv(override=True)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
ConversationHistory = List[Dict[str, str]]

# --- Constants ---
CLASSIFICATION_MODEL = "jai-chat-1-3-2"
RERANKER_MODEL = "typhoon-gemma-12b"
SUBQUERY_MODEL = "gemini-2.0-flash"
NORMAL_RAG_MODEL = 'gemini-2.5-flash'
NON_RAG_MODEL = "gemini-2.5-flash"

# --- Embedding Setup (Global Scope) ---
# BGE = SentenceTransformer("BAAI/bge-m3")

class Embedder:
    def __init__(self):
        """Initializes the Embedder with a local BGE model."""
        logger.info("Embedder initialized with BGE SentenceTransformer.")

    async def embed(self, text: Union[str, List[str]], input_type: str) -> Optional[List[List[float]]]:
        """
        Generate embeddings using a local BGE model asynchronously.
        The 'input_type' parameter is kept for signature consistency but is not used by this BGE implementation.
        """
        try:
            # BGE.encode is synchronous and CPU-bound, so run it in a thread to avoid blocking the event loop.
            # loop = asyncio.get_running_loop()
            # response = await loop.run_in_executor(None, BGE.encode, text)
            # print(response)
            # print(len(response))
            # return response.tolist()
        
            client = OpenAI(base_url="https://bai-ap.jts.co.th:10629/v1")
            response = client.embeddings.create(
                input=text,
                model="bge-m3"
            )

            # print(len(response.data[0].embedding))
            # print(response.data[0].embedding)
            return response.data[0].embedding
        except Exception as e:
            logger.error(f"Error during BGE embedding: {e}", exc_info=True)
            return None

class LLMFinanceAnalyzer:
    def __init__(self):
        self.gemini_api_key = os.getenv("GEMINI_API_KEY")

        
        
        self.client_gemini = None
        if self.gemini_api_key:
            try:
                self.client_gemini = AsyncOpenAI(api_key=self.gemini_api_key, base_url="https://generativelanguage.googleapis.com/v1beta/openai/")
                logger.info("LLMFinanceAnalyzer initialized with Gemini client.")
            except Exception as e:
                logger.error(f"Failed to initialize Gemini client: {e}")
        else:
            logger.warning("GEMINI_API_KEY not found, Gemini client not initialized.")

        

    def _get_client_for_model(self, model_name: str) -> Optional[AsyncOpenAI]:
        """Selects the appropriate client based on the model name."""
        if model_name.startswith("gpt-"):
            return self.client_openai
        elif model_name.startswith("gemini-"):
            return self.client_gemini
        elif model_name.startswith("typhoon-"):
            return self.client_typhoon
        elif model_name.startswith("gemma3-"):
            return self.client_gemma
        else:
            return self.client_jai

    @observe()
    async def _call_llm(
        self,
        model: str,
        messages: List[Dict[str, str]],
        temperature: float,
        max_tokens: int = 2048,
        seed: int = 66,
        max_retries: int = 2,
        stream: bool = False,
        tools: Optional[List[Dict[str, Any]]] = None,
    ) -> Union[Optional[str], AsyncGenerator[str, None]]:
        """Internal helper to call the appropriate LLM client with retries."""
        client = self._get_client_for_model(model)
        if not client:
            logger.error(f"No async client available for model {model}.")
            return None if not stream else (x for x in [])
        attempt = 0
        while attempt <= max_retries:
            try:
                if stream:
                    if model.startswith("gemini-"):
                        response_stream = await client.chat.completions.create(
                        model=model, messages=messages, stream=True, reasoning_effort="none", tools= tools
                    )
                    else:
                        response_stream = await client.chat.completions.create(
                            model=model, messages=messages, stream=True, tools= tools
                        )
                    async def _async_stream_generator():
                        full_tool_calls = None

                        try:
                            async for chunk in response_stream:
                                # delta_content = chunk.choices[0].delta.content.replace("•", "\n•")
                                # print(1)
                                if chunk:
                                    delta = chunk.choices[0].delta
                                    content = delta.content
                                    if delta.content:
                                        # Clean up content by removing unwanted characters
                                        
                                        delta_content = content.replace("•", "\n•").replace("!","")
                                        delta_content = re.sub(r'(?<=[\u0E00-\u0E7F]) +(?=[\u0E00-\u0E7F])', '', delta_content)
                                        
                                        yield delta_content
                                    
                                    if delta.tool_calls:
                                        tool_call  = delta.tool_calls[0]
                                        full_tool_calls = [
                                            {
                                                "id":tool_call.id,
                                                "type":"function",
                                                "function": {"name": tool_call.function.name, "arguments": tool_call.function.arguments}
                                            }
                                            ] 
                            i = 0            
                            while full_tool_calls and i<7:
                                assistant_tool_call_msg = {
                                    "role": "assistant",
                                    "content": None,
                                    "tool_calls": full_tool_calls
                                }
                                # Yield this message to be added to the main history
                                yield assistant_tool_call_msg

                                messages_for_next_call = messages + [assistant_tool_call_msg]
                                
                                # Execute tools and create/yield tool result messages
                        
                                fn_name = full_tool_calls[0]["function"]["name"]
                                fn_args_str = full_tool_calls[0]["function"]["arguments"]
                                try:
                                    fn_args = json.loads(fn_args_str)
                                    if fn_name == "call_admin":
                                    # Add the chat history to the function arguments
                                        fn_args['chat_history'] = messages[1:]
                                        # print(f"call_admin fn_args: {fn_args}")
                                        


                                    result_json = execute_tool(fn_name, fn_args)
                                except Exception as e:
                                    result_json = f"Error executing tool {fn_name}: {e}"

                                tool_result_msg = {
                                    "role": "tool",
                                    "tool_call_id": full_tool_calls[0]["id"],
                                    "content": result_json
                                }
                                    # Yield this message for the history as well
                                yield tool_result_msg
                                messages_for_next_call.append(tool_result_msg)

                                i += 1
                                follow_stream = await client.chat.completions.create(
                                                                model=model,
                                                                messages=messages_for_next_call,
                                                                stream=True,
                                                                tools=tools
                                                            )

                                async for follow_chunk in follow_stream:
                                    delta  = follow_chunk.choices[0].delta
                                    if delta.content:
                                        full_tool_calls = None #set to None to break the loop
                                        delta_content = delta.content.replace("•", "\n•").replace("!","")
                                        delta_content = re.sub(r'(?<=[\u0E00-\u0E7F]) +(?=[\u0E00-\u0E7F])', '', delta_content)
                                        
                                        yield delta_content
                                    if delta.tool_calls:
                                        tool_call  = delta.tool_calls[0]
                                        full_tool_calls = [
                                            {
                                                "id":tool_call.id,
                                                "type":"function",
                                                "function": {"name": tool_call.function.name, "arguments": tool_call.function.arguments}
                                            }
                                            ] 







                        except Exception as stream_err:
                            logger.error(f"Error during LLM stream ({model}): {stream_err}", exc_info=True)
                            yield f"\n[STREAM_ERROR: {stream_err}]\n"
                    return _async_stream_generator()
                else:
                    response = await client.chat.completions.create(
                        model=model, messages=messages,  stream=False
                    )
                    content = response.choices[0].message.content
                    return content.strip() if content else ""
            except (RateLimitError, APIError, Exception) as e:
                logger.warning(f"Error on attempt {attempt+1} for model {model}: {e}. Retrying...")
                attempt += 1
                if attempt > max_retries:
                    logger.error(f"Max retries exceeded for LLM call ({model}).")
                    if stream:
                        async def _error_gen(): yield f"\n[STREAM_ERROR: Max retries exceeded]\n"
                        return _error_gen()
                    return None
                await asyncio.sleep(3 * attempt)
        return None

    @observe()
    async def classify_rag_requirement(self, conversation: ConversationHistory) -> Optional[str]:
        """Classifies if the latest query requires RAG ('yes' or 'no') using full context."""
        if not conversation:
            return 'no'
        print(conversation)
        system_prompt = get_rag_classification_prompt()
        messages = [{"role": "user", "content": system_prompt+"/n"+conversation[0].get("content")}] 
        result = await self._call_llm(model=CLASSIFICATION_MODEL, messages=messages, temperature=0, max_tokens=10, stream=False)
        print(result)
        if isinstance(result, str):
            result_lower = result.lower().strip().rstrip('.')
            if 'yes' in result_lower: return 'yes'
            if 'no' in result_lower: return 'no'
            logger.error(f"RAG classification result '{result}' invalid. Defaulting to 'no'.")
        else:
            logger.error("RAG classification LLM call failed.")
            return 'yes'
    @observe()
    async def classify_relevance(self, query: str, document_content: str) -> bool:
        """
        Classifies if a document is relevant to a given query using an LLM.
        Returns True for 'yes', False otherwise.
        """
        # truncated_content = document_content # Truncate to manage token count

        prompt = (
            "You are an expert relevance classifier. Your task is to determine if the provided "
            "DOCUMENT is use to answer USER QUERY. Be strictly"
            # "Focus on direct relevance. If the document is only vaguely related or just mentions similar topics, it is not relevant. "
            "Respond with only the word 'yes' or 'no'."
        )
        messages = [
            {"role": "system", "content": prompt},
            {"role": "user", "content": f"USER QUERY:\n---\n{query}\n---\n\nDOCUMENT:\n---\n{document_content}\n---"}
        ]

        # Use a fast and cheap model for this simple classification task
        result = await self._call_llm(
            model=RERANKER_MODEL,
            messages=messages,
            temperature=0,
            stream=False
        )
        

        if isinstance(result, str) and 'no' in result.lower():
            logger.debug(f"Relevance classification for query '{query[:30]}...': NO")
            return False

        logger.debug(f"Relevance classification for query '{query[:30]}...': Yes (Result: '{result}')")
        return True
    
    @observe()
    async def select_relevant_documents(self, query: str, documents: str) -> bool:
        
        import ast
        messages = [
            
            {"role": "user", "content": f"""{documents}\n from the context, select a single or group(up to 4, if it's more than 4, rank from the most relavant) of documents that are relevant to the query: {query}. Here is the common knowledge:
1. The Rabbit Rewards program in Thailand: This program allows users to earn and redeem points for BTS Skytrain travel and at partner merchants.
2. Rabbit reward application and registration
3. Xtreme Saving: เเพ็กเกจเดินทางสำหรับรถไฟฟ้าสายสีเขียว สีชมพู(น้องนมเย็น) เเละสีเหลืองซึ่งเเตกตามกันในเเต่ละสาย
4. โครงการ 20 บาทตลอดสาย: เป็นนโยบายของรัฐบาลที่ต้องการลดภาระค่าใช้จ่ายในการเดินทางของประชาชน โดยมีเป้าหมายให้ผู้โดยสารรถไฟฟ้าทุกสายในกรุงเทพมหานครและปริมณฑล จ่ายค่าโดยสารสูงสุดไม่เกิน 20 บาทต่อเที่ยว.

Do not describe, answer as a list of number of the documents. example [0,2,4] \n\n"""}
        ]

        # Use a fast and cheap model for this simple classification task
        result = await self._call_llm(
            model=RERANKER_MODEL,
            messages=messages,
            temperature=0,
            max_tokens=5, # 'yes' or 'no' is very short
            stream=False
        )
        try : 
            result = ast.literal_eval(result)
            return result
        except Exception as e:
            logger.error(f"Error parsing result from select_relevant_documents: {e}")
            return None

        


    @observe()
    async def generate_subquery(self, conversation: ConversationHistory) -> Optional[str]:
        """Generates structured database query components based on the conversation without tool use."""
        if not conversation:
            logger.warning("generate_subquery called with empty conversation")
            return None

        client = self._get_client_for_model(SUBQUERY_MODEL)
        if not client:
            logger.error(f"Client for subquery model '{SUBQUERY_MODEL}' not available")
            return None

        system_prompt_content = get_subquery_prompt()
        messages = [{"role": "system", "content": system_prompt_content}] + conversation

        try:
            response = await client.chat.completions.create(
                model=SUBQUERY_MODEL,
                messages=messages,
                temperature=0,
            )
            final_content = response.choices[0].message.content
        except Exception as e:
            logger.error(f"API call error in generate_subquery: {e}", exc_info=True)
            return None

        if not final_content:
            logger.error("No content received from subquery model")
            return None

        return final_content

    @observe()
    async def generate_normal_response(self, conversation: ConversationHistory) -> AsyncGenerator[str, None]:
        """Generate a RAG response, yielding text chunks."""
        try:
            
            
            system_prompt = get_normal_prompt()
            messages = [{"role": "system", "content": system_prompt}] + conversation

            result_generator = await self._call_llm(
                model=NORMAL_RAG_MODEL, messages=messages, temperature=0.2, stream=True, tools = TOOL_DEFINITIONS
            )
            
            if isinstance(result_generator, AsyncGenerator):
                async for chunk in result_generator:
                    yield chunk
            else:
                yield "[ERROR: Failed to initiate normal RAG stream.]"
        except Exception as e:
            logger.error(f"Error in generate_normal_response setup: {e}", exc_info=True)
            yield f"[ERROR: {e}]"

    @observe()
    async def generate_non_rag_response(self, conversation: ConversationHistory) -> Optional[str]:
        """Generate response for non-RAG questions."""
        messages = [{"role": "system", "content": get_non_rag_prompt()}] + conversation
        result = await self._call_llm(model=NON_RAG_MODEL, messages=messages, temperature=0, stream=False)
        
        if isinstance(result, str):
            return result.replace("!","")
        
        logger.error("generate_non_rag_response call failed or returned non-string.")
        return None