Hoàn thiện luồng xử lý pdf và ảnh

core.py

@@ -2,14 +2,28 @@
 Core agent orchestration — entry point dùng chung cho API và UI.
 """
 
+import base64
+import mimetypes
 import time
 from datetime import datetime
+from typing import Optional
+
+from langchain_core.messages import HumanMessage, ToolMessage
 
 from src.graph import run
+from src.nodes import final_response_node, image_response_node
+from src.pdf_processing import format_chat_history, pdf_to_markdown
+from src.redis_client import redis_client
 from src.state import MAX_ITERS, AgentState
 
 
-def final_answer(conversation_id: str, sender_id: str, query: str) -> tuple[str, str]:
+def final_answer(
+    conversation_id: str,
+    sender_id: str,
+    query: str,
+    pdf_path: Optional[str] = None,
+    image_path: Optional[str] = None,
+) -> tuple[str, str]:
     """
     Khởi tạo AgentState, chạy graph, trả về (câu trả lời, thời gian xử lý).
 
@@ -20,8 +34,8 @@ def final_answer(conversation_id: str, sender_id: str, query: str) -> tuple[str,
         ValueError: nếu bất kỳ tham số bắt buộc nào rỗng.
     """
     conversation_id = conversation_id.strip()
-    sender_id = sender_id.strip()
-    query = query.strip()
+    sender_id       = sender_id.strip()
+    query           = query.strip()
 
     if not conversation_id:
         raise ValueError("conversation_id không được để trống.")
@@ -30,6 +44,72 @@ def final_answer(conversation_id: str, sender_id: str, query: str) -> tuple[str,
     if not query:
         raise ValueError("query không được để trống.")
 
+    if pdf_path is not None:
+        pdf_content = pdf_to_markdown(pdf_path)
+        chat_history = redis_client.get_chat_history(conversation_id)
+        chat_text    = format_chat_history(chat_history)
+
+        tool_content = (
+            f"[Nội dung PDF]\n{pdf_content}"
+            f"\n\n[Lịch sử trò chuyện]\n{chat_text}"
+        )
+
+        state: AgentState = {
+            "conversation_id": conversation_id,
+            "sender_id":       sender_id,
+            "time":            datetime.now().isoformat(),
+            "raw_query":       query,
+            "query_type":      None,
+            "messages":        [
+                HumanMessage(content=query),
+                ToolMessage(content=tool_content, tool_call_id="pdf_reader", name="pdf_reader"),
+            ],
+            "iters":           0,
+            "max_iters":       MAX_ITERS,
+            "final_answer":    None,
+        }
+
+        t0 = time.perf_counter()
+        result  = final_response_node(state)
+        elapsed = f"{time.perf_counter() - t0:.2f}s"
+        answer  = result.get("final_answer") or "(Không có kết quả)"
+        return answer, elapsed
+
+    if image_path is not None:
+        mime_type, _ = mimetypes.guess_type(image_path)
+        mime_type = mime_type or "image/jpeg"
+
+        with open(image_path, "rb") as f:
+            image_b64 = base64.b64encode(f.read()).decode()
+
+        chat_history = redis_client.get_chat_history(conversation_id)
+        chat_text    = format_chat_history(chat_history)
+
+        text_content = f"{query}\n\n[Lịch sử trò chuyện]\n{chat_text}"
+
+        state: AgentState = {
+            "conversation_id": conversation_id,
+            "sender_id":       sender_id,
+            "time":            datetime.now().isoformat(),
+            "raw_query":       query,
+            "query_type":      None,
+            "messages":        [
+                HumanMessage(content=[
+                    {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{image_b64}"}},
+                    {"type": "text", "text": text_content},
+                ]),
+            ],
+            "iters":           0,
+            "max_iters":       MAX_ITERS,
+            "final_answer":    None,
+        }
+
+        t0 = time.perf_counter()
+        result  = image_response_node(state)
+        elapsed = f"{time.perf_counter() - t0:.2f}s"
+        answer  = result.get("final_answer") or "(Không có kết quả)"
+        return answer, elapsed
+
     initial_state: AgentState = {
         "conversation_id": conversation_id,
         "sender_id":       sender_id,
@@ -43,8 +123,19 @@ def final_answer(conversation_id: str, sender_id: str, query: str) -> tuple[str,
     }
 
     t0 = time.perf_counter()
-    result = run(initial_state)
+    result  = run(initial_state)
     elapsed = f"{time.perf_counter() - t0:.2f}s"
 
     answer = result.get("final_answer") or "(Không có kết quả)"
     return answer, elapsed
+
+
+if __name__ == "__main__":
+    answer, elapsed = final_answer(
+        conversation_id="04ba40fe-61c7-4906-9f51-5ada0a392dac",
+        sender_id="@slavakpa",
+        query="tóm tắt nội dung tài liệu này",
+        pdf_path="temp/test_doc.pdf",
+    )
+    print(answer)
+    print(f"\n({elapsed})")

llm.py

@@ -1,17 +1,20 @@
 from langchain_google_genai import ChatGoogleGenerativeAI
 from src.config import GEMINI_API_KEY, DEFAULT_MODEL
 
-llm = ChatGoogleGenerativeAI(
-    model=DEFAULT_MODEL,
+_base_kwargs = dict(
     temperature=0,
     top_p=1,
     top_k=1,
     max_tokens=None,
     timeout=None,
     max_retries=2,
-    google_api_key=GEMINI_API_KEY
+    google_api_key=GEMINI_API_KEY,
 )
 
+llm = ChatGoogleGenerativeAI(model=DEFAULT_MODEL, **_base_kwargs)
+
+multimodal_llm = ChatGoogleGenerativeAI(model=DEFAULT_MODEL, **_base_kwargs)
+
 if __name__ == "__main__":
     response = llm.invoke("Hello World là gì?").content
     print(response)

nodes.py

@@ -6,7 +6,7 @@ from typing import Any, Literal
 
 from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 
-from src.llm import llm
+from src.llm import llm, multimodal_llm
 from src.prompts import (
     final_response_prompt,
     orchestrator_prompt,
@@ -162,6 +162,23 @@ def _extract_tool_results(state: AgentState) -> str:
     return "\n\n".join(parts) if parts else "(Không có kết quả từ tool)"
 
 
+# ════════════════════════════════════════════════════════════════════
+# NODE 6 — ImageResponseNode
+# ════════════════════════════════════════════════════════════════════
+def image_response_node(state: AgentState) -> AgentState:
+    """Nhận HumanMessage chứa ảnh + text, gọi multimodal LLM sinh câu trả lời."""
+    logger.info("[ImageResponseNode] Xử lý ảnh cho %s", state["sender_id"])
+
+    response = multimodal_llm.invoke(state["messages"])
+    answer   = response.content
+
+    logger.info("[ImageResponseNode] Hoàn thành (%d ký tự)", len(answer))
+    return {**state, "messages": [AIMessage(content=answer)], "final_answer": answer}
+
+
+# ════════════════════════════════════════════════════════════════════
+# NODE 5 — FinalResponseNode
+# ════════════════════════════════════════════════════════════════════
 def final_response_node(state: AgentState) -> AgentState:
     """Tổng hợp ToolMessage(s) và sinh câu trả lời cuối cùng."""
     logger.info("[FinalResponseNode] Tổng hợp câu trả lời...")