eabdadfb62
Features: - FastAPI server with OpenAI-compatible endpoints (/v1/chat/completions, /v1/models) - RAG system with document processing and vector storage - Support for multiple document formats (PDF, DOCX, HTML, text, code) - Streaming response support - Tool integration with website_downloader - Document management API endpoints - GLM-4.7-Flash integration via z-ai-web-dev-sdk - Works transparently with Open WebUI and other OpenAI clients Components: - main.py: FastAPI application with OpenAI-compatible API - rag/: RAG system (document processor, vector store, retriever) - tools/: Tool manager with website_downloader integration - .env.example: Configuration template
157 lines
3.9 KiB
Python
157 lines
3.9 KiB
Python
"""
|
|
Tools Module - Tool management for the RAG system
|
|
|
|
Provides a unified interface for tool registration and execution.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from typing import Any, Callable, Optional
|
|
|
|
# Import the website downloader tool
|
|
from website_downloader_tool import (
|
|
website_downloader,
|
|
get_tool_schema as get_website_downloader_schema,
|
|
)
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class ToolManager:
|
|
"""
|
|
Manages tool registration and execution.
|
|
|
|
Provides:
|
|
- Tool registration
|
|
- Tool schema generation
|
|
- Tool execution with error handling
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._tools: dict[str, Callable] = {}
|
|
self._schemas: dict[str, dict] = {}
|
|
|
|
# Register built-in tools
|
|
self._register_builtin_tools()
|
|
|
|
def _register_builtin_tools(self) -> None:
|
|
"""Register built-in tools."""
|
|
# Register website downloader
|
|
self.register_tool(
|
|
name="website_downloader",
|
|
function=website_downloader,
|
|
schema=get_website_downloader_schema(),
|
|
)
|
|
|
|
log.info(f"Registered {len(self._tools)} built-in tools")
|
|
|
|
def register_tool(
|
|
self,
|
|
name: str,
|
|
function: Callable,
|
|
schema: dict,
|
|
) -> None:
|
|
"""
|
|
Register a new tool.
|
|
|
|
Args:
|
|
name: Tool name
|
|
function: Tool function
|
|
schema: OpenAI function schema
|
|
"""
|
|
self._tools[name] = function
|
|
self._schemas[name] = schema
|
|
log.info(f"Registered tool: {name}")
|
|
|
|
def get_tool_schema(self, name: str) -> Optional[dict]:
|
|
"""Get the schema for a tool."""
|
|
return self._schemas.get(name)
|
|
|
|
def get_all_schemas(self) -> list[dict]:
|
|
"""Get schemas for all registered tools."""
|
|
return list(self._schemas.values())
|
|
|
|
def list_tools(self) -> list[str]:
|
|
"""List all registered tool names."""
|
|
return list(self._tools.keys())
|
|
|
|
def execute_tool(
|
|
self,
|
|
name: str,
|
|
arguments: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Execute a tool with the given arguments.
|
|
|
|
Args:
|
|
name: Tool name
|
|
arguments: Tool arguments
|
|
|
|
Returns:
|
|
Tool result
|
|
"""
|
|
if name not in self._tools:
|
|
return {
|
|
"success": False,
|
|
"error": f"Unknown tool: {name}",
|
|
}
|
|
|
|
try:
|
|
log.info(f"Executing tool: {name} with args: {arguments}")
|
|
result = self._tools[name](**arguments)
|
|
return result
|
|
|
|
except TypeError as e:
|
|
log.error(f"Invalid arguments for tool {name}: {e}")
|
|
return {
|
|
"success": False,
|
|
"error": f"Invalid arguments: {str(e)}",
|
|
}
|
|
|
|
except Exception as e:
|
|
log.exception(f"Tool execution failed: {name}")
|
|
return {
|
|
"success": False,
|
|
"error": str(e),
|
|
}
|
|
|
|
def execute_tool_from_json(
|
|
self,
|
|
name: str,
|
|
arguments_json: str,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Execute a tool with JSON arguments.
|
|
|
|
Args:
|
|
name: Tool name
|
|
arguments_json: JSON string of arguments
|
|
|
|
Returns:
|
|
Tool result
|
|
"""
|
|
try:
|
|
arguments = json.loads(arguments_json)
|
|
return self.execute_tool(name, arguments)
|
|
except json.JSONDecodeError as e:
|
|
return {
|
|
"success": False,
|
|
"error": f"Invalid JSON arguments: {str(e)}",
|
|
}
|
|
|
|
|
|
# Global tool manager instance
|
|
_tool_manager: Optional[ToolManager] = None
|
|
|
|
|
|
def get_tool_manager() -> ToolManager:
|
|
"""Get or create the global tool manager instance."""
|
|
global _tool_manager
|
|
|
|
if _tool_manager is None:
|
|
_tool_manager = ToolManager()
|
|
|
|
return _tool_manager
|