Upload 72 files

src/agents/custom_chatbot/func.py

@@ -1,4 +1,4 @@
-from typing import TypedDict,Optional
+from typing import TypedDict, Optional
 from langchain_core.messages import AnyMessage, ToolMessage
 from langgraph.graph.message import add_messages
 from typing import Sequence, Annotated
@@ -26,7 +26,7 @@ def get_info_collection(messages):
 
 
 async def collection_info_agent(state: State):
-    model_name = state.get("model_name", "gpt-4o")
+    model_name = state.get("model_name")
     _, collection_info_agent = get_custom_chatbot_chains(model_name)
     return await collection_info_agent.ainvoke(state)
 
@@ -35,7 +35,7 @@ async def create_prompt(state: State):
     messages = state.get("messages")
     name, info = get_info_collection(messages)
     logger.info(f"create_prompt {info}")
-    model_name = state.get("model_name", "gpt-4o")
+    model_name = state.get("model_name")
     create_system_chain, _ = get_custom_chatbot_chains(model_name)
     res = await create_system_chain.ainvoke({"info": info})
     return {"prompt": res.content, "name": name}

src/agents/rag_agent_template/func.py

@@ -1,27 +1,30 @@
 from typing import TypedDict, Optional, List
 from langchain_core.messages import AnyMessage, ToolMessage
 from langgraph.graph.message import add_messages
-from .prompt import get_rag_chains
 from typing import Sequence, Annotated
 from langchain_core.messages import RemoveMessage
 from langchain_core.documents import Document
-from .tools import retrieve_document
+from .tools import retrieve_document, python_repl, duckduckgo_search
 from src.utils.logger import logger
+from src.config.llm import get_llm
+from .prompt import template_prompt
 
-tools = [retrieve_document]
+tools = [retrieve_document, python_repl, duckduckgo_search]
 
 
 class State(TypedDict):
     messages: Annotated[Sequence[AnyMessage], add_messages]
     selected_ids: Optional[List[str]]
     selected_documents: Optional[List[Document]]
-    tools: list
+    tools: Optional[List[str]]
     prompt: str
     model_name: Optional[str]
 
 
 def trim_history(state: State):
     history = state.get("messages", [])
+    tool_names = state.get("tools", [])
+    
     if len(history) > 10:
         num_to_remove = len(history) - 10
         remove_messages = [
@@ -38,7 +41,10 @@ def trim_history(state: State):
 
 def execute_tool(state: State):
     tool_calls = state["messages"][-1].tool_calls
+    tool_names = state.get("tools", [])
     tool_name_to_func = {tool.name: tool for tool in tools}
+    tool_functions = [tool_name_to_func[name] for name in tool_names if name in tool_name_to_func]
+    
     selected_ids = []
     selected_documents = []
     tool_messages = []
@@ -64,7 +70,11 @@ def execute_tool(state: State):
                     )
                 continue
             tool_response = tool_func.invoke(tool_args)
-            tool_messages.append(tool_response)
+            print(f"tool_response: {tool_response}")
+            tool_messages.append(ToolMessage(
+                tool_call_id=tool_id,
+                content=tool_response,
+            ))
 
     return {
         "selected_ids": selected_ids,
@@ -75,22 +85,30 @@ def execute_tool(state: State):
 
 def generate_answer_rag(state: State):
     messages = state["messages"]
-    tools = state["tools"]
-    model_name = state.get("model_name", "gemini-2.0-flash")  # default to gemini-2.0-flash
-    rag_answering_chain_tool, rag_answering_chain = get_rag_chains(model_name)
-    logger.info(f"tools: {tools}")
-    if tools:
-        response = rag_answering_chain_tool.invoke(
-            {
-                "messages": messages,
-                "prompt": state["prompt"] + "Sử dụng tool `retrieve_document` để truy xuất tài liệu để bổ sung thông tin cho câu trả lời",
-            }
-        )
-    else:
-        response = rag_answering_chain.invoke(
-            {
-                "messages": messages,
-                "prompt": state["prompt"],
-            }
-        )
+    tool_names = state.get("tools", [])
+    prompt = state["prompt"]
+    model_name = state.get("model_name", "gemini-2.0-flash")
+    
+    tool_name_to_func = {tool.name: tool for tool in tools}
+    tool_functions = [tool_name_to_func[name] for name in tool_names if name in tool_name_to_func]
+    
+    print(f"tools: {tool_functions}")
+    llm_call = template_prompt | get_llm(model_name).bind_tools(tool_functions)
+
+    if tool_functions:
+        for tool in tool_functions:
+            if tool.name == "retrieve_document":
+                prompt += "Sử dụng tool `retrieve_document` để truy xuất tài liệu để bổ sung thông tin cho câu trả lời"
+            if tool.name == "python_repl":
+                prompt += "Sử dụng tool `python_repl` để thực hiện các tác vụ liên quan đến tính toán phức tạp"
+            if tool.name == "duckduckgo_search":
+                prompt += "Sử dụng tool `duckduckgo_search` để tìm kiếm thông tin trên internet"
+
+    response = llm_call.invoke(
+        {
+            "messages": messages,
+            "prompt": prompt,
+        }
+    )
+
     return {"messages": response}

src/agents/rag_agent_template/prompt.py

@@ -2,17 +2,9 @@ from langchain_core.prompts import ChatPromptTemplate
 from src.config.llm import get_llm
 from .tools import retrieve_document
 
-rag_prompt = ChatPromptTemplate.from_messages(
+template_prompt = ChatPromptTemplate.from_messages(
     [
         ("system", "{prompt}"),
         ("placeholder", "{messages}"),
     ]
 )
-
-
-def get_rag_chains(model_name: str):
-    llm = get_llm(model_name)
-    llm_rag = llm.bind_tools([retrieve_document])
-    rag_answering_chain_tool = rag_prompt | llm_rag
-    rag_answering_chain = rag_prompt | llm
-    return rag_answering_chain_tool, rag_answering_chain

src/agents/rag_agent_template/tools.py

@@ -3,12 +3,18 @@ from src.config.vector_store import test_rag_vector_store
 from src.utils.helper import convert_list_context_source_to_str
 from src.utils.logger import logger
 from langchain_core.runnables import RunnableConfig
+from langchain_experimental.utilities import PythonREPL
+from langchain_community.tools import DuckDuckGoSearchRun
+
+
+duckduckgo_search = DuckDuckGoSearchRun()
+
+python_exec = PythonREPL()
 
 
 @tool
 def retrieve_document(query: str, config: RunnableConfig):
     """Ưu tiên truy xuất tài liệu từ vector store nếu câu hỏi liên quan đến vai trò của chatbot.
-    
 
     Args:
         query (str): Câu truy vấn của người dùng bằng tiếng Việt
@@ -34,3 +40,16 @@ def retrieve_document(query: str, config: RunnableConfig):
         "selected_documents": selected_documents,
         "selected_ids": selected_ids,
     }
+
+
+@tool
+def python_repl(code: str):
+    """
+    A Python shell. Use this to execute python commands. Input should be a valid python command. If you want to see the output of a value, you should print it out with `print(...)`.
+
+    Args:
+        code (str): Python code to execute
+    Returns:
+        str: Output of the Python code
+    """
+    return python_exec.run(code)

src/apis/create_app.py

@@ -4,13 +4,13 @@ from src.apis.routers.rag_agent_template import router as router_rag_agent_templ
 from src.apis.routers.file_processing_router import router as router_file_processing
 from src.apis.routers.custom_chatbot_router import router as custom_chatbot_processing
 from src.apis.routers.vector_store_router import router as vector_store_router
-
+from src.apis.routers.tts_router import router as tts_router
 api_router = APIRouter()
 api_router.include_router(router_rag_agent_template)
 api_router.include_router(router_file_processing)
 api_router.include_router(custom_chatbot_processing)
 api_router.include_router(vector_store_router)
-
+api_router.include_router(tts_router)
 def create_app():
     app = FastAPI(
         docs_url="/",

src/apis/routers/tts_router.py

@@ -0,0 +1,145 @@
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel
+from transformers import VitsModel, AutoTokenizer
+import torch
+import os
+import uuid
+from fastapi.responses import FileResponse
+import soundfile as sf
+import numpy as np
+from src.utils.logger import logger
+from google import genai
+from google.genai import types
+import wave
+from typing import Literal
+
+router = APIRouter()
+
+
+class TTSRequest(BaseModel):
+    text: str
+
+
+class GeminiTTSRequest(BaseModel):
+    text: str
+    voice_name: str = "Kore"
+
+
+# Initialize model and tokenizer globally
+try:
+    logger.info("Loading TTS model and tokenizer...")
+    model = VitsModel.from_pretrained("facebook/mms-tts-vie")
+    tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-vie")
+    logger.info("TTS model and tokenizer loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to load TTS model: {str(e)}")
+    raise
+
+# Initialize Google Gemini client
+try:
+    logger.info("Initializing Google Gemini client...")
+    gemini_client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))
+    logger.info("Google Gemini client initialized successfully")
+except Exception as e:
+    logger.error(f"Failed to initialize Google Gemini client: {str(e)}")
+    raise
+
+
+def save_wave_file(
+    filename: str,
+    pcm: bytes,
+    channels: int = 1,
+    rate: int = 24000,
+    sample_width: int = 2,
+):
+    """Save PCM data to a WAV file."""
+    with wave.open(filename, "wb") as wf:
+        wf.setnchannels(channels)
+        wf.setsampwidth(sample_width)
+        wf.setframerate(rate)
+        wf.writeframes(pcm)
+
+
+@router.post("/tts/huggingface")
+async def huggingface_tts(request: TTSRequest):
+    try:
+        logger.info(
+            f"Processing HuggingFace TTS request for text: {request.text[:50]}..."
+        )
+
+        # Tokenize input
+        inputs = tokenizer(request.text, return_tensors="pt")
+        logger.info("Text tokenized successfully")
+
+        # Generate audio
+        with torch.no_grad():
+            output = model(**inputs).waveform
+        logger.info("Audio generated successfully")
+
+        # Convert tensor to numpy array
+        audio_numpy = output.squeeze().cpu().numpy()
+
+        # Create audio directory if it doesn't exist
+        audio_dir = os.path.join(os.getcwd(), "audio_files")
+        os.makedirs(audio_dir, exist_ok=True)
+
+        # Generate unique filename
+        audio_filename = f"huggingface_{uuid.uuid4()}.wav"
+        audio_path = os.path.join(audio_dir, audio_filename)
+
+        # Save audio file using soundfile
+        sf.write(audio_path, audio_numpy, model.config.sampling_rate)
+        logger.info(f"Audio saved to {audio_path}")
+
+        # Return audio file
+        return FileResponse(audio_path, media_type="audio/wav", filename=audio_filename)
+
+    except Exception as e:
+        logger.error(f"Error in huggingface_tts: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Failed to generate speech: {str(e)}"
+        )
+
+
+@router.post("/tts/gemini")
+def gemini_tts(request: GeminiTTSRequest):
+    try:
+        logger.info(f"Processing Gemini TTS request for text: {request.text[:50]}...")
+
+        response = gemini_client.models.generate_content(
+            model="gemini-2.5-flash-preview-tts",
+            contents=request.text,
+            config=types.GenerateContentConfig(
+                response_modalities=["AUDIO"],
+                speech_config=types.SpeechConfig(
+                    voice_config=types.VoiceConfig(
+                        prebuilt_voice_config=types.PrebuiltVoiceConfig(
+                            voice_name=request.voice_name,
+                        )
+                    )
+                ),
+            ),
+        )
+
+        data = response.candidates[0].content.parts[0].inline_data.data
+
+        # Create audio directory if it doesn't exist
+        audio_dir = os.path.join(os.getcwd(), "audio_files")
+        os.makedirs(audio_dir, exist_ok=True)
+
+        # Generate unique filename
+        audio_filename = f"gemini_{uuid.uuid4()}.wav"
+        audio_path = os.path.join(audio_dir, audio_filename)
+
+        # Save audio file
+        save_wave_file(audio_path, data)
+        logger.info(f"Audio saved to {audio_path}")
+
+        # Return audio file
+        return FileResponse(audio_path, media_type="audio/wav", filename=audio_filename)
+
+    except Exception as e:
+        logger.error(f"Error in gemini_tts: {str(e)}")
+        raise HTTPException(
+            status_code=500, detail=f"Failed to generate speech: {str(e)}"
+        )