mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 16:13:23 +00:00
Alex
37d93cbd86
Replace the sandbox Docling extractor with read_document, backed by the in-process backend parser (the same one ingestion uses) and offloaded to a dedicated 'parsing' Celery queue so it can run on GPU-capable workers with predictable RAM. The tool resolves the input ref under the run-scoped gate, enqueues the parse, and awaits it with a timeout (degrading to an error rather than hanging); the worker independently re-resolves the artifact through the same gate and never trusts a raw path. Untrusted files get the upload path's safeguards (extension whitelist, size cap, sanitized temp file, cleanup). Options: output (markdown/text/structured/chunks), ocr, pages, engine, max_chars, include_tables, persist, json_schema. The workflow native-file 'extract' fallback now uses the same worker path, so document parsing no longer needs the sandbox and works on every backend. Also fixes the branch's periodic-task test (the sandbox reaper made it 12) and points the dev and e2e Celery workers at the parsing queue.
78 lines
2.7 KiB
Python
78 lines
2.7 KiB
Python
import importlib
|
|
import inspect
|
|
import os
|
|
import pkgutil
|
|
|
|
from application.agents.tools.base import Tool
|
|
|
|
|
|
class ToolManager:
|
|
def __init__(self, config):
|
|
self.config = config
|
|
self.tools = {}
|
|
self.load_tools()
|
|
|
|
def load_tools(self):
|
|
tools_dir = os.path.join(os.path.dirname(__file__))
|
|
for finder, name, ispkg in pkgutil.iter_modules([tools_dir]):
|
|
if name == "base" or name.startswith("__"):
|
|
continue
|
|
module = importlib.import_module(f"application.agents.tools.{name}")
|
|
for member_name, obj in inspect.getmembers(module, inspect.isclass):
|
|
if issubclass(obj, Tool) and obj is not Tool and not obj.internal:
|
|
tool_config = self.config.get(name, {})
|
|
self.tools[name] = obj(tool_config)
|
|
|
|
def load_tool(self, tool_name, tool_config, user_id=None):
|
|
self.config[tool_name] = tool_config
|
|
module = importlib.import_module(f"application.agents.tools.{tool_name}")
|
|
for member_name, obj in inspect.getmembers(module, inspect.isclass):
|
|
if issubclass(obj, Tool) and obj is not Tool:
|
|
if (
|
|
tool_name
|
|
in {
|
|
"mcp_tool",
|
|
"notes",
|
|
"memory",
|
|
"todo_list",
|
|
"scheduler",
|
|
"remote_device",
|
|
"code_executor",
|
|
"artifact_generator",
|
|
"read_document",
|
|
}
|
|
and user_id
|
|
):
|
|
return obj(tool_config, user_id)
|
|
else:
|
|
return obj(tool_config)
|
|
|
|
def execute_action(self, tool_name, action_name, user_id=None, **kwargs):
|
|
if tool_name not in self.tools:
|
|
raise ValueError(f"Tool '{tool_name}' not loaded")
|
|
if (
|
|
tool_name
|
|
in {
|
|
"mcp_tool",
|
|
"memory",
|
|
"todo_list",
|
|
"notes",
|
|
"scheduler",
|
|
"remote_device",
|
|
"code_executor",
|
|
"artifact_generator",
|
|
"read_document",
|
|
}
|
|
and user_id
|
|
):
|
|
tool_config = self.config.get(tool_name, {})
|
|
tool = self.load_tool(tool_name, tool_config, user_id)
|
|
return tool.execute_action(action_name, **kwargs)
|
|
return self.tools[tool_name].execute_action(action_name, **kwargs)
|
|
|
|
def get_all_actions_metadata(self):
|
|
metadata = []
|
|
for tool in self.tools.values():
|
|
metadata.extend(tool.get_actions_metadata())
|
|
return metadata
|