diff --git a/nodes/agent/BaseAgent.py b/nodes/agent/BaseAgent.py index c64a341..63dd938 100644 --- a/nodes/agent/BaseAgent.py +++ b/nodes/agent/BaseAgent.py @@ -1,3 +1,6 @@ +import logging + +from comfy_execution.graph import ExecutionBlocker from griptape.drivers import DummyVectorStoreDriver from griptape.tools import QueryTool, RagTool, VectorStoreTool from openai import OpenAIError @@ -179,8 +182,16 @@ class BaseAgent: # Warn for models model, simple_model = self.agent.model_check() - + if model == "": + error_msg = "You have provided a blank model for the Agent Configuration.\n\nPlease specify a model configuration, or disconnect it from the agent." + logging.error(error_msg) + raise Exception(error_msg) + return ( + ExecutionBlocker(error_msg), + self.agent, + ) if simple_model: + print(f"This is a simple model: {model}") return (self.agent.model_response(model), self.agent) # Check for inputs. If none, then just create the agent @@ -193,10 +204,6 @@ class BaseAgent: else: prompt_text = STRING + "\n\n" + input_string - # if len(tools) > 0: - # self.agent.add_task(ToolkitTask(prompt_text, tools=tools)) - # else: - # self.agent.add_task(PromptTask(prompt_text)) result = self.agent.run(prompt_text) output_string = result.output_task.output.value return ( diff --git a/nodes/agent/gtComfyAgent.py b/nodes/agent/gtComfyAgent.py index 41d8122..6084294 100644 --- a/nodes/agent/gtComfyAgent.py +++ b/nodes/agent/gtComfyAgent.py @@ -54,7 +54,7 @@ class gtComfyAgent(Agent): simple_models = ["llama3", "llama3:latest", "mistral", "LLama-3"] drivers = ["OllamaPromptDriver", "LMStudioPromptDriver"] agent_prompt_driver_name = self.prompt_driver.__class__.__name__ - + print(f"Agent Prompt Driver Name: {agent_prompt_driver_name=}") model = self.prompt_driver.model if agent_prompt_driver_name in drivers: if model == "": diff --git a/nodes/config/gtUILMStudioDriversConfig.py b/nodes/config/gtUILMStudioDriversConfig.py index 5f1c23f..3b1a09d 100644 --- a/nodes/config/gtUILMStudioDriversConfig.py +++ b/nodes/config/gtUILMStudioDriversConfig.py @@ -1,3 +1,6 @@ +import logging + +from comfy_execution.graph import ExecutionBlocker from griptape.configs import Defaults from griptape.configs.drivers import ( DriversConfig, @@ -48,33 +51,37 @@ class gtUILMStudioDriversConfig(gtUIBaseDriversConfig): def create(self, **kwargs): drivers_config_params = {} - # Create instances of the driver classes - prompt_driver_builder = gtUILMStudioChatPromptDriver() - embedding_driver_builder = gtUILMStudioEmbeddingDriver() - - # Build parameters for drivers - prompt_driver_params = prompt_driver_builder.build_params(**kwargs) - embedding_driver_params = embedding_driver_builder.build_params(**kwargs) - - # Create Driver Configs - drivers_config_params["prompt_driver"] = OpenAiChatPromptDriver( - **prompt_driver_params - ) - if "model" not in embedding_driver_params: - drivers_config_params["embedding_driver"] = DummyEmbeddingDriver() - else: - drivers_config_params["embedding_driver"] = OpenAiEmbeddingDriver( - **embedding_driver_params - ) - - drivers_config_params["vector_store_driver"] = LocalVectorStoreDriver( - embedding_driver=drivers_config_params["embedding_driver"] - ) - try: + # Create instances of the driver classes + prompt_driver_builder = gtUILMStudioChatPromptDriver() + embedding_driver_builder = gtUILMStudioEmbeddingDriver() + + # Build parameters for drivers + prompt_driver_params = prompt_driver_builder.build_params(**kwargs) + embedding_driver_params = embedding_driver_builder.build_params(**kwargs) + if "model" not in prompt_driver_params: + custom_config = ExecutionBlocker( + "Please provide a model for the prompt driver." + ) + return (custom_config,) + # Create Driver Configs + drivers_config_params["prompt_driver"] = OpenAiChatPromptDriver( + **prompt_driver_params + ) + if "model" not in embedding_driver_params: + drivers_config_params["embedding_driver"] = DummyEmbeddingDriver() + else: + drivers_config_params["embedding_driver"] = OpenAiEmbeddingDriver( + **embedding_driver_params + ) + + drivers_config_params["vector_store_driver"] = LocalVectorStoreDriver( + embedding_driver=drivers_config_params["embedding_driver"] + ) Defaults.drivers_config = DriversConfig(**drivers_config_params) custom_config = Defaults.drivers_config except Exception as e: - print(e) + custom_config = ExecutionBlocker(f"{e}") + logging.error(e) return (custom_config,) diff --git a/nodes/drivers/gtUILMStudioChatPromptDriver.py b/nodes/drivers/gtUILMStudioChatPromptDriver.py index c355a47..9eafe0e 100644 --- a/nodes/drivers/gtUILMStudioChatPromptDriver.py +++ b/nodes/drivers/gtUILMStudioChatPromptDriver.py @@ -1,3 +1,4 @@ +from comfy_execution.graph import ExecutionBlocker from griptape.drivers import OpenAiChatPromptDriver from .gtUIOpenAiCompatibleChatPromptDriver import gtUIOpenAiCompatibleChatPromptDriver @@ -78,9 +79,12 @@ class gtUILMStudioChatPromptDriver(gtUIOpenAiCompatibleChatPromptDriver): def create(self, **kwargs): params = self.build_params(**kwargs) + if not params.get("model"): + driver = ExecutionBlocker("Model is required.") + return (driver,) try: driver = OpenAiChatPromptDriver(**params) return (driver,) except Exception as e: - print(f"Error creating driver: {e}") - return (None, str(e)) + driver = ExecutionBlocker(f"Error creating driver: {e}") + return (driver,)