Rag_basic/src/chatbot/function_tools/router.py

88 lines
3.4 KiB
Python

# -*- coding: utf-8 -*-
"""
FunctionRouter — Bộ điều phối công cụ (function tools)
------------------------------------------------------
- Quản lý và điều hướng các công cụ mà LLM có thể gọi.
- Khi LLM trả JSON {"action": "<tên_tool>", "params": {...}},
router sẽ xác định đúng tool và gọi thực thi tương ứng.
"""
import logging
from typing import Dict
from .base_tool import BaseTool
from .txt_tool import TxtTool
from .google_tool import GoogleSearchTool # ✅ Thêm import Google search
logger = logging.getLogger("function_router")
class FunctionRouter:
"""
Router trung tâm, chịu trách nhiệm:
- Đăng ký tool vào registry (theo tên)
- Điều phối thực thi tool
"""
def __init__(self):
# Danh mục các tool khả dụng (tên -> instance)
self.tools: Dict[str, BaseTool] = {}
# 🔧 Đăng ký các tool mặc định
self.register_tool(TxtTool())
self.register_tool(GoogleSearchTool())
logger.info(f"✅ FunctionRouter khởi tạo, có {len(self.tools)} công cụ sẵn sàng.")
# ------------------------------------------------------------
# Đăng ký tool mới
# ------------------------------------------------------------
def register_tool(self, tool: BaseTool):
"""Đăng ký 1 tool mới vào router"""
if not isinstance(tool, BaseTool):
raise TypeError("Tool phải kế thừa BaseTool.")
if not hasattr(tool, "name"):
raise ValueError("Tool phải có thuộc tính name.")
if not callable(getattr(tool, "execute", None)):
raise ValueError("Tool phải có hàm execute().")
# Ghi đè nếu trùng tên
if tool.name in self.tools:
logger.warning(f"⚠️ Tool '{tool.name}' đã tồn tại, ghi đè instance mới.")
self.tools[tool.name] = tool
logger.info(f"🔧 Đã đăng ký tool: {tool.name}{tool.description}")
# ------------------------------------------------------------
# Lấy danh sách tool (cho debug / hiển thị)
# ------------------------------------------------------------
def list_tools(self) -> Dict[str, str]:
"""Trả về danh sách tên và mô tả của tất cả công cụ"""
return {name: t.description for name, t in self.tools.items()}
# ------------------------------------------------------------
# Thực thi tool dựa trên action
# ------------------------------------------------------------
def execute_tool(self, action: str, **params) -> str:
"""
Gọi tool theo tên hành động (action).
Ví dụ:
router.execute_tool("create_txt", text="Xin chào!")
router.execute_tool("search_google", query="tin tức hôm nay")
Trả về: Chuỗi kết quả hoặc message lỗi.
"""
tool = self.tools.get(action)
if not tool:
msg = f"❌ Không tìm thấy tool '{action}' trong router."
logger.error(msg)
return msg
try:
logger.info(f"⚙️ Thực thi tool '{action}' với params={params}")
result = tool.execute(**params)
logger.info(f"✅ Tool '{action}' thực thi thành công.")
return result
except Exception as e:
msg = f"❌ Lỗi khi thực thi tool '{action}': {e}"
logger.exception(msg)
return msg