fix: enforce model_instance requirement in agent runners for improved…

… reliability Signed-off-by: -LAN- <[email protected]>
langgenius · Dec 24, 2024 · b7c91c4 · b7c91c4
1 parent 41fac0b
commit b7c91c4
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 6 deletions.
diff --git a/api/core/agent/base_agent_runner.py b/api/core/agent/base_agent_runner.py
@@ -52,6 +52,7 @@
 class BaseAgentRunner(AppRunner):
     def __init__(
         self,
+        *,
         tenant_id: str,
         application_generate_entity: AgentChatAppGenerateEntity,
         conversation: Conversation,
@@ -65,7 +66,7 @@ def __init__(
         prompt_messages: Optional[list[PromptMessage]] = None,
         variables_pool: Optional[ToolRuntimeVariablePool] = None,
         db_variables: Optional[ToolConversationVariables] = None,
-        model_instance: ModelInstance | None = None,
+        model_instance: ModelInstance,
     ) -> None:
         self.tenant_id = tenant_id
         self.application_generate_entity = application_generate_entity
@@ -111,8 +112,6 @@ def __init__(
         db.session.close()
 
         # check if model supports stream tool call
-        # FIXME confirm here, model_instance is not None
-        assert model_instance is not None
         llm_model = cast(LargeLanguageModel, model_instance.model_type_instance)
         model_schema = llm_model.get_model_schema(model_instance.model, model_instance.credentials)
         features = model_schema.features if model_schema and model_schema.features else []

diff --git a/api/core/agent/cot_agent_runner.py b/api/core/agent/cot_agent_runner.py
@@ -82,8 +82,6 @@ def increase_usage(final_llm_usage_dict: dict[str, Optional[LLMUsage]], usage: L
                 llm_usage.total_price += usage.total_price
 
         model_instance = self.model_instance
-        if not model_instance:
-            raise ValueError("failed to get model instance")
 
         while function_call_state and iteration_step <= max_iteration_steps:
             # continue to run until there is not any tool call

diff --git a/api/core/agent/fc_agent_runner.py b/api/core/agent/fc_agent_runner.py
@@ -69,7 +69,6 @@ def increase_usage(final_llm_usage_dict: dict[str, LLMUsage], usage: LLMUsage):
                 llm_usage.total_price += usage.total_price
 
         model_instance = self.model_instance
-        assert model_instance is not None
 
         while function_call_state and iteration_step <= max_iteration_steps:
             function_call_state = False