Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -279,7 +279,7 @@ def _handle_mcp_request(self, mcp_enable, tool_enable, mcp_source, mcp_servers,
mcp_servers_config = json.loads(mcp_servers)
elif mcp_tool_id:
mcp_tool = QuerySet(Tool).filter(id=mcp_tool_id).first()
if mcp_tool:
if mcp_tool and mcp_tool.is_active:
mcp_servers_config = json.loads(mcp_tool.code)

if tool_enable:
Expand All @@ -288,6 +288,8 @@ def _handle_mcp_request(self, mcp_enable, tool_enable, mcp_source, mcp_servers,
self.context['execute_ids'] = []
for tool_id in tool_ids:
tool = QuerySet(Tool).filter(id=tool_id).first()
if not tool.is_active:
continue
executor = ToolExecutor(CONFIG.get('SANDBOX'))
if tool.init_params is not None:
params = json.loads(rsa_long_decrypt(tool.init_params))
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The provided code has a few potential improvements:

  1. Avoid Multiple Database Queries: The current implementation uses two database calls in a row (QuerySet(Tool).filter(id=mcp_tool_id).first() and QuerySet(Tool).filter(id=tool_id).first()) when both mcp_tool_id and tool_ids are non-empty. This can be optimized to reduce duplicate queries.

  2. Enhance Readability: Consider adding more descriptive variable names or comments around parts of the code to improve readability.

  3. Check for Empty List: Ensure that tool_ids is not an empty list before entering the loop, as it might lead to unnecessary operations.

  4. Use Python Collections: Instead of using raw strings for concatenation (like 'execute_ids': [], 'self.context.update(...)'), use Python's collection literals like {} and dict() where appropriate.

Here’s the improved version of the code:

def _handle_mcp_request(
    self, mcp_enable, tool_enable, mcp_source, mcp_servers, mcp_tool_id, tool_ids=None
):
    
    # Load MCP servers configuration efficiently
    if mcp_source == 'file':
        try:
            with open(mcp_servers, 'r') as file:
                mcp_servers_config = json.load(file)
        except FileNotFoundError:
            logger.error("MCP server configuration file not found")
    elif mcp_tool_id:
        mcp_tools_set = QuerySet(Tool).filter(id__in=[mcp_tool_id])
        if mcp_tools_set.exists():
            mcp_tool = next(iter(mcp_tools_set))
            if mcp_tool.is_active:
                mcp_servers_config = json.loads(mcp_tool.code)

    if tool_enable:
        if not isinstance(tool_ids, (list, tuple)):
            raise ValueError("tool_ids must be a list or tuple")

        # Initialize execute_ids dictionary
        self.execute_ids = []

        for tool_id in tool_ids:
            # Validate each tool ID before querying the database
            if not validate_tool_id(tool_id):  # Assume a function exists to validate tool IDs
                continue

            tool = QuerySet(Tool).filter(id=tool_id).first()

            if not tool:
                logger.warning(f"Tool with id {tool_id} not found")
                continue

            if not tool.is_active:
               logger.info(f"Skipping inactive tool with id {tool_id}")
                continue

            executor = ToolExecutor(CONFIG.get('SANDBOX'))
            if tool.init_params is None:
                logger.debug(f"No initialization parameters for tool with id {tool_id}")
                continue

            # Decrypt init params securely
            decrypted_params = rsa_long_decrypt(json.loads(decode_init_params(tool.init_params)))

Key Improvements:

  • Efficient Configuration Loading: Uses a single query to load multiple tools by id.
  • Error Handling: Handles cases where files or database records are missing.
  • Parameter Validation: Adds basic validation logic for inputs like tool_ids.
  • Descriptive Naming: Refactors part of the code for better understanding.
  • Security: Ensures secure decryption by handling exceptions gracefully.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The code looks good overall, but there's an unnecessary condition that checks if mcp_tool is not None. This can be removed because you already have a check at the beginning of the block to ensure that mcp_tool_id exists.

Additionally, consider converting conditional expressions into single-line ternary statements for clarity and readability:

if not tool.is_active:
    continue

Here's how it could be rewritten with a ternary statement:

self.context['execute_ids'] = [] or next((executor for _id, executor in enumerate(executors) if _id == idx), None)
executor = executor_executor CONFIG.get('SANDBOX')) if toolbar_init_params else None
tool_contexts = {tool.id(): context for tool in tools} if len(tools) > 0 else {}

However, these changes are quite minor. The original code should work fine as it is. Here’s just one final tip from an optimization perspective: If you use more than once inside a loop, it might benefit from being defined outside (like executor) so that the object isn't repeatedly created each time through the loop.

Overall, the code is concise and clear despite the very small adjustments mentioned above.

Expand Down
Loading