refractor: store history in mongodb

.gitignore

@@ -10,4 +10,4 @@ __pycache__/
 *.mp3
 *.pem
 test.py
-test.ipynb
+test.ipynb

app.py

@@ -6,14 +6,13 @@ from dotenv import load_dotenv
 import time
 import numpy as np
 import sys
+import uuid
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 from backend.tts import synthesize_text
 from backend.asr import transcribe_audio, transcribe_typhoon
 from backend.utils import preprocess_audio, is_valid_turn, preprocess_audio_simplified
 from backend.main import stream_chat_response
-import json
 from pydub import AudioSegment
-import ast
 from backend.utils import get_device
 if get_device() == "cpu":
     load_dotenv(override=True)
@@ -25,7 +24,11 @@ sound_samples = np.array(phone_waiting_sound.get_array_of_samples(), dtype=np.in
 if phone_waiting_sound.channels > 1:
     sound_samples = sound_samples.reshape((-1, phone_waiting_sound.channels)).mean(axis=1)
 sound_samples = sound_samples.astype(np.float32) / 32768.0  # Normalize to [-1, 
-def startup(_):
+def startup(*arg):
+    print(arg[0])
+    print("_______")
+    print(arg[1])
+
     yield (phone_waiting_sound.frame_rate, sound_samples)
     STARTUP_MESSAGE = "สวัสดีค่ะ พลอย 1577Homeshopping ยินดีให้บริการค่ะ"
     yield from synthesize_text(STARTUP_MESSAGE)
@@ -65,145 +68,16 @@ h1 {
     box-shadow: 0 0 15px rgba(0, 0, 0, 0.2);
 }
 """
-def format_history_for_chatbot(history):
-    """
-    Formats the conversation history for display in the Gradio Chatbot.
-    It creates readable strings for tool calls and tool results.
-    """
-    formatted_history = []
-    if not history:
-        return []
-    for turn in history:
-        role = turn.get("role")
-        content = turn.get("content")
-        tool_calls = turn.get("tool_calls")
-        
-        if role == "user":
-            formatted_history.append({"role": "user", "content": content})
-        elif role == "assistant":
-            if tool_calls:
-                # Display a user-friendly message for the tool call
-                id = tool_calls[0]['id']
-                func_name = tool_calls[0]['function']['name']
-                func_args = tool_calls[0]['function']['arguments']
-                display_content = f"<id>{id}</id><func_name>{func_name}</func_name><func_args>{func_args}</func_args>"
-                # display_content = (
-                #     f"**Calling Tool:**\n"
-                #     f"```json\n"
-                #     f"{{\n"
-                #     f'  "name": "{func_name}",\n'
-                #     f'  "arguments": {func_args}\n'
-                #     f"}}\n"
-                #     f"```"
-                # )
-                formatted_history.append({"role": "assistant", "content": display_content})
-            else:
-                # Regular assistant message
-                formatted_history.append({"role": "assistant", "content": content})
-        elif role == "tool":
-            # Display a user-friendly message for the tool result
-            id = turn.get("tool_call_id")
-            result_content = json.dumps(json.loads(content), indent=2, ensure_ascii=False)
-            display_content = f"<id>{id}</id><content>{content}</content>"
-            # display_content = (
-            #     f"**Tool Result:**\n"
-            #     f"```json\n"
-            #     f"{result_content}\n"
-            #     f"```"
-            # )
-            # Represent tool results as if the "assistant" is providing them
-            formatted_history.append({"role": "assistant", "content": display_content})
-    
-    return formatted_history
-import re
-def revert_to_openai_format(formatted_history):
-    """
-    Converts a history list formatted for the Gradio Chatbot UI back into
-    the standard OpenAI API format. It parses custom string formats for
-    tool calls and tool results.
+def snapshot_history(history):
+    """Return a shallow copy of the current chatbot history."""
+    return [dict(turn) for turn in history] if history else []
 
-    Args:
-        formatted_history (list): A list of message dictionaries as they appear
-                                  in the Gradio Chatbot component.
 
-    Returns:
-        list: A list of message dictionaries compliant with the OpenAI API format.
-    """
-    openai_history = []
-    
-    # Pre-compile regex patterns for efficiency
-    # Pattern to find a tool call message
-    tool_call_pattern = re.compile(
-        r"<id>(.*?)</id><func_name>(.*?)</func_name><func_args>(.*?)</func_args>",
-        re.DOTALL  # Use DOTALL in case arguments contain newlines
-    )
-    # Pattern to find a tool result message
-    tool_result_pattern = re.compile(
-        r"<id>(.*?)</id><content>(.*?)</content>",
-        re.DOTALL
-    )
-
-    if not formatted_history:
-        return []
-
-    for turn in formatted_history:
-        role = turn.get("role")
-        content = turn.get("content") 
-        # If content is None, treat it as an empty string for the regex search.
-        if content is None:
-            content = ""
-
-        if role == "user":
-            openai_history.append(turn)
-            continue
-
-        if role == "assistant":
-            # Check if this is a formatted tool call
-            tool_call_match = tool_call_pattern.search(content)
-            if tool_call_match:
-                call_id, func_name, func_args_str = tool_call_match.groups()
-                
-                # Reconstruct the original tool_calls structure
-                reverted_turn = {
-                    "role": "assistant",
-                    "content": None,
-                    "tool_calls": [
-                        {
-                            "id": call_id,
-                            "type": "function",
-                            "function": {
-                                "name": func_name,
-                                "arguments": func_args_str
-                            },
-                        }
-                    ],
-                }
-                openai_history.append(reverted_turn)
-                continue
-
-            # Check if this is a formatted tool result (as per your formatter's logic)
-            tool_result_match = tool_result_pattern.search(content)
-            if tool_result_match:
-                tool_call_id, tool_content = tool_result_match.groups()
-                
-                # Reconstruct the original tool message
-                # NOTE: The role must be 'tool' for the API
-                reverted_turn = {
-                    "role": "tool",
-                    "tool_call_id": tool_call_id,
-                    "content": tool_content.strip() # Remove trailing space
-                }
-                openai_history.append(reverted_turn)
-                continue
-            
-            # If no patterns match, it's a regular assistant message
-            if turn.get("content") is not None:
-                openai_history.append(turn)
-
-    return openai_history
-
-
-def response(audio: tuple[int, np.ndarray] | None, conversation_history):
+def response(
+    audio: tuple[int, np.ndarray] | None,
+    conversation_history,
+    session_id: str | None,
+):
     """
     Handles user audio input, transcribes it, streams LLM text via backend.main,
     and synthesizes chunks to audio while updating the conversation history.
@@ -216,13 +90,12 @@ def response(audio: tuple[int, np.ndarray] | None, conversation_history):
     # print(f"Initial conver:{conversation_history}")
     # print('-----------------------------')
 
-    conversation_history = revert_to_openai_format(conversation_history)
-    # print(f"After convert:{conversation_history}")
+    conversation_history = conversation_history or []
     start_time = time.time()
-    if conversation_history is None:
-        conversation_history = []
-
-    previous_history = list(conversation_history)
+    session_identifier = session_id or ""
+    if not session_identifier:
+        session_identifier = str(uuid.uuid4())
+        print(f"[WARN] Missing session_id; generated temporary session {session_identifier}")
 
     if not audio or audio[1] is None or not np.any(audio[1]):
         print("No audio input detected; skipping response generation.")
@@ -272,17 +145,13 @@ def response(audio: tuple[int, np.ndarray] | None, conversation_history):
     print(f"User: {transcription}")
     if is_valid_turn(user_turn):
         conversation_history.append(user_turn)
-        yield AdditionalOutputs(format_history_for_chatbot(conversation_history))
+        yield AdditionalOutputs(snapshot_history(conversation_history))
 
     # print("Conversation history:", conversation_history)
 
     assistant_turn = {"role": "assistant", "content": ""}
     conversation_history.append(assistant_turn)
     
-    # print(previous_history)
-    history_for_stream = [dict(turn) for turn in previous_history if is_valid_turn(turn)]
-    # print(f"history_for_stream{history_for_stream}")
-
     text_buffer = ""
     full_response = ""
     delimiter_count = 0
@@ -294,7 +163,7 @@ def response(audio: tuple[int, np.ndarray] | None, conversation_history):
     start_llm_stream = time.time()
 
     try:
-        for chunk in stream_chat_response(history_for_stream, transcription):
+        for chunk in stream_chat_response(session_identifier, transcription):
             # print(f"LLM chunk: {text_chunk}")
             if  isinstance(chunk, str):
                 text_chunk = chunk
@@ -356,20 +225,9 @@ def response(audio: tuple[int, np.ndarray] | None, conversation_history):
                             first_chunk_sent = True
                             text_buffer = ""
                             delimiter_count = 0
-                            yield AdditionalOutputs(format_history_for_chatbot(conversation_history))
+                            yield AdditionalOutputs(snapshot_history(conversation_history))
 
                     i += 1
-            elif isinstance(chunk, dict) and "role" in chunk:
-                # print(f"Received tool message for history: {chunk}")
-                
-                if chunk.get("content") is None:
-                    chunk["content"] = ""
-                conversation_history.insert(-1, chunk)
-
-                # Update the chatbot UI to reflect the new history structure
-                
-                yield AdditionalOutputs(format_history_for_chatbot(conversation_history))
-        
         if text_buffer.strip():
             buffer_to_send = text_buffer.strip()
             try:
@@ -391,7 +249,7 @@ def response(audio: tuple[int, np.ndarray] | None, conversation_history):
                 text_buffer = ""
                 delimiter_count = 0
                 
-                yield AdditionalOutputs(format_history_for_chatbot(conversation_history))
+                yield AdditionalOutputs(snapshot_history(conversation_history))
 
     except Exception as e:
         print(f"An error occurred during response generation or synthesis: {e}")
@@ -401,7 +259,7 @@ def response(audio: tuple[int, np.ndarray] | None, conversation_history):
         except Exception as synth_error:
             print(f"Could not synthesize error message: {synth_error}")
         assistant_turn["content"] = (assistant_turn.get("content", "") + f" [Error: {e}]").strip()
-        yield AdditionalOutputs(format_history_for_chatbot(conversation_history))
+        yield AdditionalOutputs(snapshot_history(conversation_history))
 
     total_latency = time.time() - start_time
     print(f"Total:      {total_latency:.4f}s")
@@ -411,9 +269,15 @@ def response(audio: tuple[int, np.ndarray] | None, conversation_history):
 
 async def get_credentials():
     return await get_cloudflare_turn_credentials_async(hf_token=os.getenv('HF_TOKEN'))
+
+
+def initialize_session_id():
+    """Create a new session identifier for syncing backend history."""
+    return str(uuid.uuid4())
    
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="orange", secondary_hue="orange")) as demo:
     gr.HTML("""<h1 style='text-align: center'>1577 Voicebot Demo</h1>""")
+    session_state = gr.State(value=None)
     with gr.Row():
         with gr.Column(scale=1, elem_classes=["phone-column"]):
             audio = WebRTC(
@@ -447,6 +311,13 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="orange", second
             )
     gr.DeepLinkButton()
 
+    demo.load(
+        fn=initialize_session_id,
+        inputs=None,
+        outputs=[session_state],
+        queue=False,
+    )
+
     audio.stream(
         fn=ReplyOnPause(
             response,
@@ -460,11 +331,12 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="orange", second
                 min_speech_duration_ms=200,
                 max_speech_duration_s=float("inf"),
                 min_silence_duration_ms=1200,
+                speech_pad_ms=300  
             ),
             can_interrupt=False,
             startup_fn=startup,
         ),
-        inputs=[audio, conversation_history], 
+        inputs=[audio, conversation_history, session_state], 
         outputs=[audio],
         concurrency_limit=1000,
         time_limit=8192
@@ -484,4 +356,4 @@ demo.launch(
     share=False,
     server_name="0.0.0.0",
     server_port=int(os.getenv("PORT", 7860)),
-)
+)

backend/conversation_store.py

@@ -0,0 +1,116 @@
+"""Async helpers for persisting per-session conversation history."""
+
+from __future__ import annotations
+
+import logging
+import os
+from datetime import datetime, timezone
+from typing import Any, Dict, List
+
+from motor.motor_asyncio import AsyncIOMotorClient
+from pymongo.errors import PyMongoError
+
+logger = logging.getLogger(__name__)
+
+
+class ConversationStore:
+    """Wrapper around MongoDB for storing chat histories."""
+
+    def __init__(self) -> None:
+        mongo_uri = os.getenv("MONGO_URL")
+        if not mongo_uri:
+            raise RuntimeError("MONGO_URL is not configured for ConversationStore")
+
+        db_name = os.getenv("CONVERSATION_DB", "homeshopping")
+        collection_name = os.getenv("CONVERSATION_COLLECTION", "conversation_history")
+        self._client = AsyncIOMotorClient(mongo_uri)
+        self._collection = self._client[db_name][collection_name]
+        logger.info(
+            "ConversationStore connected to %s.%s",
+            db_name,
+            collection_name,
+        )
+
+    async def get_history(self, session_id: str) -> List[Dict[str, Any]]:
+        """Return a shallow copy of the stored history for the given session."""
+        if not session_id:
+            return []
+
+        try:
+            doc = await self._collection.find_one(
+                {"session_id": session_id},
+                {"_id": 0, "history": 1},
+            )
+        except PyMongoError as exc:
+            logger.error("Failed to load conversation history: %s", exc, exc_info=True)
+            return []
+
+        history = doc.get("history", []) if doc else []
+        return [dict(message) for message in history]
+
+    async def append_messages(self, session_id: str, messages: List[Dict[str, Any]]) -> None:
+        """Append one or more messages to the stored history."""
+        if not session_id or not messages:
+            return
+
+        safe_messages = [self._sanitize_message(message) for message in messages if message]
+        if not safe_messages:
+            return
+
+        now = datetime.now(timezone.utc)
+        try:
+            await self._collection.update_one(
+                {"session_id": session_id},
+                {
+                    "$setOnInsert": {"session_id": session_id, "created_at": now},
+                    "$set": {"updated_at": now},
+                    "$push": {"history": {"$each": safe_messages}},
+                },
+                upsert=True,
+            )
+        except PyMongoError as exc:
+            logger.error("Failed to append conversation messages: %s", exc, exc_info=True)
+
+    async def upsert_session_metadata(
+        self,
+        session_id: str,
+        persona: Dict[str, Any] | None = None,
+        user_info: Any | None = None,
+    ) -> None:
+        """Persist persona/user info for the session without overwriting history."""
+        if not session_id:
+            return
+
+        updates: Dict[str, Any] = {}
+        if persona is not None:
+            updates["persona"] = persona
+        if user_info is not None:
+            updates["user_info"] = user_info
+
+        if not updates:
+            return
+
+        now = datetime.now(timezone.utc)
+        updates["updated_at"] = now
+
+        try:
+            await self._collection.update_one(
+                {"session_id": session_id},
+                {
+                    "$setOnInsert": {"session_id": session_id, "created_at": now},
+                    "$set": updates,
+                },
+                upsert=True,
+            )
+        except PyMongoError as exc:
+            logger.error("Failed to update session metadata: %s", exc, exc_info=True)
+
+    @staticmethod
+    def _sanitize_message(message: Dict[str, Any]) -> Dict[str, Any]:
+        allowed_keys = {"role", "content", "tool_call_id", "tool_calls"}
+        return {key: message.get(key) for key in allowed_keys if key in message}
+
+
+conversation_store = ConversationStore()
+
+__all__ = ["ConversationStore", "conversation_store"]

backend/main.py

@@ -7,7 +7,7 @@ import logging
 import os
 from queue import Queue
 from threading import Lock, Thread
-from typing import AsyncGenerator, Dict, Iterator, List, Optional
+from typing import Any, AsyncGenerator, Dict, Iterator, List, Optional
 
 from dotenv import load_dotenv
 from langfuse import Langfuse
@@ -16,6 +16,7 @@ import sys
 sys.path.append(os.path.abspath('./backend'))
 from models import LLMFinanceAnalyzer
 from functions import MongoHybridSearch
+from conversation_store import conversation_store
 from utils import get_device
 if get_device() == "cpu":
     load_dotenv(override=True)
@@ -84,8 +85,36 @@ def _generate_pseudo_conversation(conversation: List[Dict[str, str]]) -> List[Di
 
 
 @observe()
-async def _stream_chat_async(history: List[Dict[str, str]], message: str) -> AsyncGenerator[str, None]:
-    full_conversation = [msg.copy() for msg in history] + [{"role": "user", "content": message}]
+async def _stream_chat_async(
+    session_id: str,
+    message: str,
+    persona: Optional[str] = None,
+    persona_state: Optional[Dict[str, Any]] = None,
+    user_info: Optional[Any] = None,
+) -> AsyncGenerator[str, None]:
+    try:
+        stored_history = await conversation_store.get_history(session_id)
+    except Exception as exc:  # noqa: BLE001
+        logger.error("Failed to load conversation history for %s: %s", session_id, exc, exc_info=True)
+        stored_history = []
+
+    if session_id:
+        try:
+            await conversation_store.upsert_session_metadata(
+                session_id,
+                persona=persona_state,
+                user_info=user_info,
+            )
+        except Exception as exc:  # noqa: BLE001
+            logger.error("Failed to persist session metadata for %s: %s", session_id, exc, exc_info=True)
+
+    full_conversation = [msg.copy() for msg in stored_history] + [{"role": "user", "content": message}]
+    if message and session_id:
+        try:
+            await conversation_store.append_messages(session_id, [{"role": "user", "content": message}])
+        except Exception as exc:  # noqa: BLE001
+            logger.error("Failed to persist user message for %s: %s", session_id, exc, exc_info=True)
+
     truncated_history = _create_truncated_history(full_conversation, 300)
     pseudo_conversation = _generate_pseudo_conversation(truncated_history)
 
@@ -116,7 +145,10 @@ async def _stream_chat_async(history: List[Dict[str, str]], message: str) -> Asy
         # else:
         #     limited_conversation = full_conversation
         limited_conversation = full_conversation
-        response_generator = llm_analyzer.generate_normal_response(limited_conversation)
+        response_generator = llm_analyzer.generate_normal_response(
+            limited_conversation,
+            session_id=session_id,
+        )
 
         async for chunk in response_generator:
             if chunk:
@@ -124,14 +156,23 @@ async def _stream_chat_async(history: List[Dict[str, str]], message: str) -> Asy
                 await asyncio.sleep(0.05)
     else:
         limited_conversation = full_conversation[-9:] if len(full_conversation) > 9 else full_conversation
-        final_response = await llm_analyzer.generate_non_rag_response(limited_conversation)
+        final_response = await llm_analyzer.generate_non_rag_response(
+            limited_conversation,
+            session_id=session_id,
+        )
         if final_response:
             yield final_response
         else:
             yield "ขออภัยค่ะ เกิดข้อผิดพลาดในการประมวลผลคำถามของคุณ"
 
 
-def stream_chat_response(history: List[Dict[str, str]], message: str) -> Iterator[str]:
+def stream_chat_response(
+    session_id: str,
+    message: str,
+    persona: Optional[str] = None,
+    persona_state: Optional[Dict[str, Any]] = None,
+    user_info: Optional[Any] = None,
+) -> Iterator[str]:
     """Synchronously iterate over streaming LLM chunks."""
 
     loop = _ensure_stream_loop()
@@ -139,7 +180,7 @@ def stream_chat_response(history: List[Dict[str, str]], message: str) -> Iterato
 
     async def runner() -> None:
         try:
-            async for chunk in _stream_chat_async(history, message):
+            async for chunk in _stream_chat_async(session_id, message, persona, persona_state, user_info):
                 output_queue.put_nowait(chunk) 
         except Exception as exc:  # noqa: BLE001
             logger.error("Unhandled error in async chat stream: %s", exc, exc_info=True)

backend/models.py

@@ -12,6 +12,7 @@ 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 conversation_store import conversation_store
 
 from systemprompt import (
     get_rag_classification_prompt,
@@ -110,6 +111,7 @@ class LLMFinanceAnalyzer:
         max_retries: int = 2,
         stream: bool = False,
         tools: Optional[List[Dict[str, Any]]] = None,
+        session_id: Optional[str] = None,
     ) -> Union[Optional[str], AsyncGenerator[str, None]]:
         """Internal helper to call the appropriate LLM client with retries."""
         client = self._get_client_for_model(model)
@@ -134,6 +136,18 @@ class LLMFinanceAnalyzer:
                         token_input = 0
                         token_output = 0
                         tokenin = 0
+                        assistant_chunks: List[str] = []
+
+                        async def flush_assistant() -> None:
+                            nonlocal assistant_chunks
+                            if session_id and assistant_chunks:
+                                text = "".join(assistant_chunks).strip()
+                                if text:
+                                    await conversation_store.append_messages(
+                                        session_id,
+                                        [{"role": "assistant", "content": text}],
+                                    )
+                            assistant_chunks = []
 
                         try:
                             async for  chunk in response_stream:
@@ -155,10 +169,11 @@ class LLMFinanceAnalyzer:
                                         
                                         delta_content = content.replace("•", "\n•").replace("!","")
                                         delta_content = re.sub(r'(?<=[\u0E00-\u0E7F]) +(?=[\u0E00-\u0E7F])', '', delta_content)
-                                        
+                                        assistant_chunks.append(delta_content)
                                         yield delta_content
                                     
                                     if delta.tool_calls:
+                                        await flush_assistant()
                                         tool_call  = delta.tool_calls[0]
                                         full_tool_calls = [
                                             {
@@ -174,8 +189,12 @@ class LLMFinanceAnalyzer:
                                     "content": None,
                                     "tool_calls": full_tool_calls
                                 }
-                                # Yield this message to be added to the main history
                                 yield assistant_tool_call_msg
+                                if session_id:
+                                    await conversation_store.append_messages(
+                                        session_id,
+                                        [assistant_tool_call_msg],
+                                    )
 
                                 messages_for_next_call = messages + [assistant_tool_call_msg]
                                 
@@ -203,6 +222,11 @@ class LLMFinanceAnalyzer:
                                 }
                                     # Yield this message for the history as well
                                 yield tool_result_msg
+                                if session_id:
+                                    await conversation_store.append_messages(
+                                        session_id,
+                                        [tool_result_msg],
+                                    )
                                 messages_for_next_call.append(tool_result_msg)
 
                                 i += 1
@@ -231,9 +255,10 @@ class LLMFinanceAnalyzer:
                                         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)
-                                        
+                                        assistant_chunks.append(delta_content)
                                         yield delta_content
                                     if delta.tool_calls:
+                                        await flush_assistant()
                                         tool_call  = delta.tool_calls[0]
                                         full_tool_calls = [
                                             {
@@ -254,7 +279,9 @@ class LLMFinanceAnalyzer:
                         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"
-                        print(f"Total tokens used - Input: {token_input}, Output: {token_output}")  
+                        finally:
+                            await flush_assistant()
+                            print(f"Total tokens used - Input: {token_input}, Output: {token_output}")  
                         # response = requests.post("https://1577shop-api.jts.co.th/count_tokens", json={
                         #                                                                             "input_token": token_input,
                         #                                                                             "output_token": token_output
@@ -266,7 +293,13 @@ class LLMFinanceAnalyzer:
                         model=model, messages=messages,  stream=False
                     )
                     content = response.choices[0].message.content
-                    return content.strip() if content else ""
+                    text = content.strip() if content else ""
+                    if session_id and text:
+                        await conversation_store.append_messages(
+                            session_id,
+                            [{"role": "assistant", "content": text}],
+                        )
+                    return text
             except (RateLimitError, APIError, Exception) as e:
                 logger.warning(f"Error on attempt {attempt+1} for model {model}: {e}. Retrying...")
                 attempt += 1
@@ -398,7 +431,11 @@ Do not describe, answer as a list of number of the documents. example [0,2,4] \n
         return final_content
 
     @observe()
-    async def generate_normal_response(self, conversation: ConversationHistory) -> AsyncGenerator[str, None]:
+    async def generate_normal_response(
+        self,
+        conversation: ConversationHistory,
+        session_id: Optional[str] = None,
+    ) -> AsyncGenerator[str, None]:
         """Generate a RAG response, yielding text chunks."""
         try:
             
@@ -407,7 +444,12 @@ Do not describe, answer as a list of number of the documents. example [0,2,4] \n
             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
+                model=NORMAL_RAG_MODEL,
+                messages=messages,
+                temperature=0.2,
+                stream=True,
+                tools=TOOL_DEFINITIONS,
+                session_id=session_id,
             )
             
             if isinstance(result_generator, AsyncGenerator):
@@ -420,13 +462,23 @@ Do not describe, answer as a list of number of the documents. example [0,2,4] \n
             yield f"[ERROR: {e}]"
 
     @observe()
-    async def generate_non_rag_response(self, conversation: ConversationHistory) -> Optional[str]:
+    async def generate_non_rag_response(
+        self,
+        conversation: ConversationHistory,
+        session_id: Optional[str] = None,
+    ) -> 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)
+        result = await self._call_llm(
+            model=NON_RAG_MODEL,
+            messages=messages,
+            temperature=0,
+            stream=False,
+            session_id=session_id,
+        )
         
         if isinstance(result, str):
             return result.replace("!","")
         
         logger.error("generate_non_rag_response call failed or returned non-string.")
-        return None
+        return None

backend/tools.py

@@ -97,7 +97,7 @@ TOOL_DEFINITIONS =[
                 "properties": {
                     "cause": {
                         "type": "string",
-                        "description": "Short description of the problem.",
+                        "description": "Short description of the problem in Thai language.",
                     }
                 },
                 "required": ["cause"],