Unverified Commit 059c089f authored by takatost's avatar takatost Committed by GitHub

fix: external data tool batch retrieve bug (#1472)

parent c1e7193c
...@@ -205,21 +205,20 @@ class Completion: ...@@ -205,21 +205,20 @@ class Completion:
results = {} results = {}
with ThreadPoolExecutor() as executor: with ThreadPoolExecutor() as executor:
futures = {} futures = {}
for tools in grouped_tools.values(): for tool in external_data_tools:
# Only query the first tool in each group if not tool.get("enabled"):
first_tool = tools[0] continue
future = executor.submit( future = executor.submit(
cls.query_external_data_tool, current_app._get_current_object(), tenant_id, app_id, first_tool, cls.query_external_data_tool, current_app._get_current_object(), tenant_id, app_id, tool,
inputs, query inputs, query
) )
for tool in tools:
futures[future] = tool futures[future] = tool
for future in concurrent.futures.as_completed(futures): for future in concurrent.futures.as_completed(futures):
tool_key, result = future.result() tool_variable, result = future.result()
if tool_key in grouped_tools: results[tool_variable] = result
for tool in grouped_tools[tool_key]:
results[tool['variable']] = result
inputs.update(results) inputs.update(results)
return inputs return inputs
...@@ -246,9 +245,7 @@ class Completion: ...@@ -246,9 +245,7 @@ class Completion:
query=query query=query
) )
tool_key = (external_data_tool.get("type"), json.dumps(external_data_tool.get("config"), sort_keys=True)) return tool_variable, result
return tool_key, result
@classmethod @classmethod
def get_query_for_agent(cls, app: App, app_model_config: AppModelConfig, query: str, inputs: dict) -> str: def get_query_for_agent(cls, app: App, app_model_config: AppModelConfig, query: str, inputs: dict) -> str:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment