Explorar el Código

fix: retriever_resource missing (#1317)

takatost hace 1 año
padre
commit
9cba1c8bf4

+ 2 - 1
api/core/orchestrator_rule_parser.py

@@ -37,12 +37,13 @@ class OrchestratorRuleParser:
 
     def to_agent_executor(self, conversation_message_task: ConversationMessageTask, memory: Optional[BaseChatMemory],
                           rest_tokens: int, chain_callback: MainChainGatherCallbackHandler,
-                          return_resource: bool = False, retriever_from: str = 'dev') -> Optional[AgentExecutor]:
+                          retriever_from: str = 'dev') -> Optional[AgentExecutor]:
         if not self.app_model_config.agent_mode_dict:
             return None
 
         agent_mode_config = self.app_model_config.agent_mode_dict
         model_dict = self.app_model_config.model_dict
+        return_resource = self.app_model_config.retriever_resource_dict.get('enabled', False)
 
         chain = None
         if agent_mode_config and agent_mode_config.get('enabled'):

+ 1 - 1
api/core/tool/dataset_retriever_tool.py

@@ -30,7 +30,7 @@ class DatasetRetrieverTool(BaseTool):
     dataset_id: str
     k: int = 3
     conversation_message_task: ConversationMessageTask
-    return_resource: str
+    return_resource: bool
     retriever_from: str
 
     @classmethod

+ 2 - 1
api/models/model.py

@@ -196,7 +196,8 @@ class AppModelConfig(db.Model):
             user_input_form=self.user_input_form,
             dataset_query_variable=self.dataset_query_variable,
             pre_prompt=self.pre_prompt,
-            agent_mode=self.agent_mode
+            agent_mode=self.agent_mode,
+            retriever_resource=self.retriever_resource
         )
 
         return new_app_model_config