griptape-nodes 0.41.0__py3-none-any.whl → 0.43.0__py3-none-any.whl

griptape_nodes/retained_mode/managers/agent_manager.py

@@ -1,7 +1,9 @@
 import json
 import logging
+import os
 import threading
 import uuid
+from typing import TYPE_CHECKING
 
 from attrs import define, field
 from griptape.artifacts import ErrorArtifact, ImageUrlArtifact, JsonArtifact
@@ -14,6 +16,7 @@ from griptape.memory.structure import ConversationMemory
 from griptape.rules import Rule, Ruleset
 from griptape.structures import Agent
 from griptape.tools import BaseImageGenerationTool
+from griptape.tools.mcp.tool import MCPTool
 from griptape.utils.decorators import activity
 from json_repair import repair_json
 from schema import Literal, Schema
@@ -42,10 +45,14 @@ from griptape_nodes.retained_mode.managers.static_files_manager import (
     StaticFilesManager,
 )
 
+if TYPE_CHECKING:
+    from griptape.tools.mcp.sessions import StreamableHttpConnection
+
 logger = logging.getLogger("griptape_nodes")
 
 API_KEY_ENV_VAR = "GT_CLOUD_API_KEY"
 SERVICE = "Griptape"
+GTN_MCP_SERVER_PORT = int(os.getenv("GTN_MCP_SERVER_PORT", "9927"))
 
 config_manager = ConfigManager()
 secrets_manager = SecretsManager(config_manager)
@@ -84,6 +91,7 @@ class AgentManager:
         self.conversation_memory = ConversationMemory()
         self.prompt_driver = None
         self.image_tool = None
+        self.mcp_tool = None
         self.static_files_manager = static_files_manager
 
         if event_manager is not None:
@@ -109,18 +117,57 @@ class AgentManager:
             msg = f"Secret '{API_KEY_ENV_VAR}' not found"
             raise ValueError(msg)
         return NodesPromptImageGenerationTool(
-            image_generation_driver=GriptapeCloudImageGenerationDriver(api_key=api_key, model="dall-e-3"),
+            image_generation_driver=GriptapeCloudImageGenerationDriver(api_key=api_key, model="gpt-image-1"),
             static_files_manager=self.static_files_manager,
         )
 
+    def _initialize_mcp_tool(self) -> MCPTool:
+        connection: StreamableHttpConnection = {  # type: ignore[reportAssignmentType]
+            "transport": "streamable_http",
+            "url": f"http://localhost:{GTN_MCP_SERVER_PORT}/mcp/",
+        }
+        return MCPTool(connection=connection)
+
     def on_handle_run_agent_request(self, request: RunAgentRequest) -> ResultPayload:
         if self.prompt_driver is None:
             self.prompt_driver = self._initialize_prompt_driver()
         if self.image_tool is None:
             self.image_tool = self._initialize_image_tool()
+        if self.mcp_tool is None:
+            self.mcp_tool = self._initialize_mcp_tool()
         threading.Thread(target=self._on_handle_run_agent_request, args=(request, EventBus.event_listeners)).start()
         return RunAgentResultStarted()
 
+    def _create_agent(self) -> Agent:
+        output_schema = Schema(
+            {
+                "generated_image_urls": [str],
+                "conversation_output": str,
+            }
+        )
+
+        tools = []
+        if self.image_tool is not None:
+            tools.append(self.image_tool)
+        if self.mcp_tool is not None:
+            tools.append(self.mcp_tool)
+
+        return Agent(
+            prompt_driver=self.prompt_driver,
+            conversation_memory=self.conversation_memory,
+            tools=tools,
+            output_schema=output_schema,
+            rulesets=[
+                Ruleset(
+                    name="generated_image_urls",
+                    rules=[
+                        Rule("Do not hallucinate generated_image_urls."),
+                        Rule("Only set generated_image_urls with images generated with your tools."),
+                    ],
+                ),
+            ],
+        )
+
     def _on_handle_run_agent_request(
         self, request: RunAgentRequest, event_listeners: list[EventListener]
     ) -> ResultPayload:
@@ -131,28 +178,7 @@ class AgentManager:
                 for url_artifact in request.url_artifacts
                 if url_artifact["type"] == "ImageUrlArtifact"
             ]
-
-            output_schema = Schema(
-                {
-                    "generated_image_urls": [str],
-                    "conversation_output": str,
-                }
-            )
-            agent = Agent(
-                prompt_driver=self.prompt_driver,
-                conversation_memory=self.conversation_memory,
-                tools=[self.image_tool] if self.image_tool else [],
-                output_schema=output_schema,
-                rulesets=[
-                    Ruleset(
-                        name="generated_image_urls",
-                        rules=[
-                            Rule("Do not hallucinate generated_image_urls."),
-                            Rule("Only set generated_image_urls with images generated with your tools."),
-                        ],
-                    ),
-                ],
-            )
+            agent = self._create_agent()
             *events, last_event = agent.run_stream([request.input, *artifacts])
             full_result = ""
             last_conversation_output = ""

griptape_nodes/retained_mode/managers/engine_identity_manager.py

@@ -0,0 +1,146 @@
+"""Manages engine identity state.
+
+Centralizes engine identity management, providing a consistent interface for
+engine ID and name operations.
+"""
+
+import logging
+from pathlib import Path
+
+from griptape_nodes.retained_mode.events.base_events import BaseEvent
+from griptape_nodes.retained_mode.managers.event_manager import EventManager
+from griptape_nodes.retained_mode.utils.engine_identity import EngineIdentity
+
+logger = logging.getLogger("griptape_nodes")
+
+
+class EngineIdentityManager:
+    """Manages engine identity and active engine state."""
+
+    _active_engine_id: str | None = None
+
+    def __init__(self, event_manager: EventManager | None = None) -> None:
+        """Initialize the EngineIdentityManager.
+
+        Args:
+            event_manager: The EventManager instance to use for event handling.
+        """
+        if event_manager is not None:
+            # Register event handlers here when engine events are defined
+            pass
+
+    @classmethod
+    def get_active_engine_id(cls) -> str | None:
+        """Get the active engine ID.
+
+        Returns:
+            str | None: The active engine ID or None if not set
+        """
+        return cls._active_engine_id
+
+    @classmethod
+    def set_active_engine_id(cls, engine_id: str) -> None:
+        """Set the active engine ID.
+
+        Args:
+            engine_id: The engine ID to set as active
+        """
+        cls._active_engine_id = engine_id
+        logger.debug("Set active engine ID to: %s", engine_id)
+
+    @classmethod
+    def initialize_engine_id(cls) -> str:
+        """Initialize the engine ID if not already set."""
+        if cls._active_engine_id is None:
+            engine_id = EngineIdentity.get_engine_id()
+            BaseEvent._engine_id = engine_id
+            cls._active_engine_id = engine_id
+            logger.debug("Initialized engine ID: %s", engine_id)
+
+        return cls._active_engine_id
+
+    @classmethod
+    def get_engine_data(cls) -> dict:
+        """Get the current engine data, creating default if it doesn't exist.
+
+        Returns:
+            dict: The current engine data
+        """
+        return EngineIdentity.get_engine_data()
+
+    @classmethod
+    def get_engine_name(cls) -> str:
+        """Get the engine name.
+
+        Returns:
+            str: The engine name
+        """
+        return EngineIdentity.get_engine_name()
+
+    @classmethod
+    def set_engine_name(cls, engine_name: str) -> None:
+        """Set and persist the current engine name.
+
+        Args:
+            engine_name: The new engine name to set
+        """
+        EngineIdentity.set_engine_name(engine_name)
+        logger.info("Updated engine name to: %s", engine_name)
+
+    @classmethod
+    def get_all_engines(cls) -> list[dict]:
+        """Get all registered engines.
+
+        Returns:
+            list[dict]: List of all engine data
+        """
+        return EngineIdentity.get_all_engines()
+
+    @classmethod
+    def get_default_engine_id(cls) -> str | None:
+        """Get the default engine ID.
+
+        Returns:
+            str | None: The default engine ID or None if not set
+        """
+        return EngineIdentity.get_default_engine_id()
+
+    @classmethod
+    def set_default_engine_id(cls, engine_id: str) -> None:
+        """Set the default engine ID.
+
+        Args:
+            engine_id: The engine ID to set as default
+
+        Raises:
+            ValueError: If engine_id is not found in registered engines
+        """
+        try:
+            EngineIdentity.set_default_engine_id(engine_id)
+            logger.info("Set default engine ID to: %s", engine_id)
+        except ValueError as e:
+            logger.error("Failed to set default engine ID: %s", e)
+            raise
+
+    @classmethod
+    def get_engine_data_file_path(cls) -> Path:
+        """Get the path where engine data is stored (for debugging/inspection).
+
+        Returns:
+            Path: The path to the engine data file
+        """
+        return EngineIdentity.get_engine_data_file_path()
+
+    @classmethod
+    def ensure_engine_initialized(cls) -> str:
+        """Ensure engine is initialized and return the engine ID.
+
+        Returns:
+            str: The initialized engine ID
+        """
+        cls.initialize_engine_id()
+        engine_id = cls.get_active_engine_id()
+        if engine_id is None:
+            msg = "Failed to initialize engine ID"
+            raise RuntimeError(msg)
+        return engine_id

griptape_nodes/retained_mode/managers/event_manager.py

@@ -66,7 +66,13 @@ class EventManager:
             del self._request_type_to_manager[request_type]
 
     def handle_request(
-        self, request: RP, operation_depth_mgr: "OperationDepthManager", workflow_mgr: "WorkflowManager"
+        self,
+        request: RP,
+        *,
+        operation_depth_mgr: "OperationDepthManager",
+        workflow_mgr: "WorkflowManager",
+        response_topic: str | None = None,
+        request_id: str | None = None,
     ) -> ResultPayload:
         """Publish an event to the manager assigned to its type.
 
@@ -74,6 +80,8 @@ class EventManager:
             request: The request to handle
             operation_depth_mgr: The operation depth manager to use
             workflow_mgr: The workflow manager to use
+            response_topic: The topic to send the response to (optional)
+            request_id: The ID of the request to correlate with the response (optional)
         """
         # Notify the manager of the event type
         with operation_depth_mgr as depth_manager:
@@ -91,7 +99,7 @@ class EventManager:
 
                 retained_mode_str = None
                 # If request_id exists, that means it's a direct request from the GUI (not internal), and should be echoed by retained mode.
-                if depth_manager.is_top_level() and request.request_id is not None:
+                if depth_manager.is_top_level() and request_id is not None:
                     retained_mode_str = depth_manager.request_retained_mode_translation(request)
 
                 # Some requests have fields marked as "omit_from_result" which should be removed from the request
@@ -101,8 +109,10 @@ class EventManager:
                 if result_payload.succeeded():
                     result_event = EventResultSuccess(
                         request=request,
+                        request_id=request_id,
                         result=result_payload,
                         retained_mode=retained_mode_str,
+                        response_topic=response_topic,
                     )
                     # If the result is a success, and the WorkflowAlteredMixin is present, that means the flow has been changed in some way.
                     # In that case, we need to flush the element changes, so we add one to the event queue.
@@ -114,8 +124,10 @@ class EventManager:
                 else:
                     result_event = EventResultFailure(
                         request=request,
+                        request_id=request_id,
                         result=result_payload,
                         retained_mode=retained_mode_str,
+                        response_topic=response_topic,
                     )
                 wrapped_event = GriptapeNodeEvent(wrapped_event=result_event)
                 EventBus.publish_event(wrapped_event)