Skip to content

fix: ensure only active tools are processed in base_chat_node.py #3982

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
liuruibin merged 2 commits into from
Sep 2, 2025
Merged

fix: ensure only active tools are processed in base_chat_node.py #3982

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ef74377
fix: ensure only active tools are processed in base_chat_node.py
liuruibin Sep 2, 2025
afe78d1
fix: update chat update_time to use timezone.now for consistency
liuruibin Sep 2, 2025
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
4 changes: 3 additions & 1 deletion apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py
View file
Open in desktop
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

@shaohuzhang1 shaohuzhang1 Sep 2, 2025

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

@shaohuzhang1 shaohuzhang1 Sep 2, 2025

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
3 changes: 2 additions & 1 deletion apps/application/serializers/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from django.core.cache import cache
from django.db.models import QuerySet
from django.utils.translation import gettext_lazy as _
from django.utils import timezone

from application.chat_pipeline.step.chat_step.i_chat_step import PostResponseHandler
from application.models import Application, ChatRecord, Chat, ApplicationVersion, ChatUserType, ApplicationTypeChoices, \
Expand Down Expand Up @@ -213,7 +214,7 @@ def append_chat_record(self, chat_record: ChatRecord):
chat_user_id=self.chat_user_id, chat_user_type=self.chat_user_type,
asker=self.get_chat_user()).save()
else:
QuerySet(Chat).filter(id=self.chat_id).update(update_time=datetime.now())
QuerySet(Chat).filter(id=self.chat_id).update(update_time=timezone.now())
# 插入会话记录
chat_record.save()
ChatCountSerializer(data={'chat_id': self.chat_id}).update_chat()
Copy link
Copy Markdown
Contributor Author

@shaohuzhang1 shaohuzhang1 Sep 2, 2025

Choose a reason for hiding this comment

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

The code looks generally correct and well-structured according to the provided context. However, here's a few minor suggestions and optimizations:

  1. Use QuerySet.first(): It might be more efficient to retrieve a single instance of Chat instead of using QuerySet.

    queryset = QuerySet(Chat).filter(id=self.chat_id)
if queryset.exists():
    chat_instance = queryset.first()
    chat_instance.update_time = timezone.now()
    chat_instance.save()

  2. Consider Using get_object_or_404(): If you want to raise an exception when no object is found, this can help manage errors gracefully.

    from django.shortcuts import get_object_or_404

chat_instance = get_object_or_404(QuerySet(Chat), id=self.chat_id
chat_instance.update_time = timezone.now()
chat_instance.save()

These changes make the code slightly more robust and potentially more performant in certain scenarios depending on the load and environment it will run in.

Expand Down
Loading