alita-sdk 0.3.462__py3-none-any.whl → 0.3.627__py3-none-any.whl

alita_sdk/runtime/langchain/assistant.py

@@ -13,7 +13,10 @@ from langchain_core.messages import (
     BaseMessage, SystemMessage, HumanMessage
 )
 from langchain_core.prompts import MessagesPlaceholder
-from .constants import REACT_ADDON, REACT_VARS, XML_ADDON
+from .constants import (REACT_ADDON, REACT_VARS, XML_ADDON, USER_ADDON,
+                        QA_ASSISTANT, NERDY_ASSISTANT, QUIRKY_ASSISTANT, CYNICAL_ASSISTANT,
+                        DEFAULT_ASSISTANT, PLAN_ADDON, PYODITE_ADDON, DATA_ANALYSIS_ADDON,
+                        SEARCH_INDEX_ADDON, FILE_HANDLING_INSTRUCTIONS)
 from .chat_message_template import Jinja2TemplatedChatMessagesTemplate
 from ..tools.echo import EchoTool
 from langchain_core.tools import BaseTool, ToolException
@@ -32,12 +35,18 @@ class Assistant:
                  memory: Optional[Any] = None,
                  store: Optional[BaseStore] = None,
                  debug_mode: Optional[bool] = False,
-                 mcp_tokens: Optional[dict] = None):
+                 mcp_tokens: Optional[dict] = None,
+                 conversation_id: Optional[str] = None,
+                 ignored_mcp_servers: Optional[list] = None,
+                 persona: Optional[str] = "generic",
+                 is_subgraph: bool = False):
 
         self.app_type = app_type
         self.memory = memory
         self.store = store
+        self.persona = persona
         self.max_iterations = data.get('meta', {}).get('step_limit', 25)
+        self.is_subgraph = is_subgraph  # Store is_subgraph flag
 
         logger.debug("Data for agent creation: %s", data)
         logger.info("App type: %s", app_type)
@@ -87,8 +96,24 @@ class Assistant:
         # Handle internal tools
         meta = data.get('meta', {})
         if meta.get("internal_tools"):
+            # Find bucket from artifact toolkit marked with is_attachment flag
+            bucket_name = None
+            for tool in version_tools:
+                if tool.get('type') == 'artifact' and tool.get('is_attachment'):
+                    bucket_name = tool.get('settings', {}).get('bucket')
+                    break
+            # Fallback: use first artifact toolkit with a bucket
+            if not bucket_name:
+                for tool in version_tools:
+                    if tool.get('type') == 'artifact' and tool.get('settings', {}).get('bucket'):
+                        bucket_name = tool['settings']['bucket']
+                        break
+
             for internal_tool_name in meta.get("internal_tools"):
-                version_tools.append({"type": "internal_tool", "name": internal_tool_name})
+                tool_config = {"type": "internal_tool", "name": internal_tool_name, "settings": {}}
+                if bucket_name:
+                    tool_config["settings"]["bucket_name"] = bucket_name
+                version_tools.append(tool_config)
 
         self.tools = get_tools(
             version_tools,
@@ -96,10 +121,56 @@ class Assistant:
             llm=self.client,
             memory_store=self.store,
             debug_mode=debug_mode,
-            mcp_tokens=mcp_tokens
+            mcp_tokens=mcp_tokens,
+            conversation_id=conversation_id,
+            ignored_mcp_servers=ignored_mcp_servers
         )
         if tools:
             self.tools += tools
+        
+        # Create ToolRegistry to track tool metadata and handle name collisions
+        self.tool_registry = {}
+        tool_name_counts = {}  # Track how many times each base name appears
+        
+        for tool in self.tools:
+            if hasattr(tool, 'name'):
+                original_name = tool.name
+                base_name = original_name
+                
+                # Extract toolkit metadata from tool configuration
+                toolkit_name = ""
+                toolkit_type = ""
+                
+                # Find matching tool config to extract metadata
+                for tool_config in version_tools:
+                    # Try to match by toolkit_name or name field
+                    config_toolkit_name = tool_config.get('toolkit_name', tool_config.get('name', ''))
+                    # Simple heuristic: toolkit info should be accessible from tool config
+                    # For now, use toolkit_name and type from config
+                    toolkit_name = config_toolkit_name
+                    toolkit_type = tool_config.get('type', '')
+                    break  # Use first match for now; will refine with better matching
+                
+                # Handle duplicate tool names by appending numeric suffix
+                if base_name in tool_name_counts:
+                    tool_name_counts[base_name] += 1
+                    # Append suffix to make unique
+                    new_name = f"{base_name}_{tool_name_counts[base_name]}"
+                    tool.name = new_name
+                    logger.info(f"Tool name collision detected: '{base_name}' -> '{new_name}'")
+                else:
+                    tool_name_counts[base_name] = 0
+                    new_name = base_name
+                
+                # Store in registry
+                self.tool_registry[tool.name] = {
+                    'toolkit_name': toolkit_name,
+                    'toolkit_type': toolkit_type,
+                    'original_tool_name': base_name
+                }
+                
+        logger.info(f"ToolRegistry initialized with {len(self.tool_registry)} tools")
+        
         # Handle prompt setup
         if app_type in ["pipeline", "predict", "react"]:
             self.prompt = data['instructions']
@@ -226,36 +297,61 @@ class Assistant:
                 chat_history_messages.append(message)
         
         # Only use prompt_instructions if explicitly specified (for predict app_type)
-        if self.app_type == "predict" and isinstance(self.prompt, str):
+        if self.app_type in ["predict", "react"] and isinstance(self.prompt, str):
             prompt_instructions = self.prompt
-
-        # take the system message from the openai prompt as a prompt instructions
-        if self.app_type == "openai" and hasattr(self.prompt, 'messages'):
-            prompt_instructions = self.__take_prompt_from_openai_messages()
-        
-        # Create a unified YAML schema with conditional tool binding
-        # Build the base node configuration
-        node_config = {
-            'id': 'agent',
-            'type': 'llm',
-            'prompt': {
-                'template': prompt_instructions or "You are a helpful assistant."
-            },
-            'input': ['messages'],
-            'output': ['messages'],
-            'transition': 'END'
-        }
         
         # Add tool binding only if tools are present
+        tool_names = []
         if simple_tools:
             tool_names = [tool.name for tool in simple_tools]
-            tool_names_yaml = str(tool_names).replace("'", '"')  # Convert to YAML-compatible format
-            node_config['tool_names'] = tool_names_yaml
             logger.info("Binding tools: %s", tool_names)
         
+        # take the system message from the openai prompt as a prompt instructions
+        if self.app_type == "openai" and hasattr(self.prompt, 'messages'):
+            prompt_instructions = self.__take_prompt_from_openai_messages()
+        
+        user_addon = USER_ADDON.format(prompt=str(prompt_instructions)) if prompt_instructions else ""
+        plan_addon = PLAN_ADDON if 'update_plan' in tool_names else ""
+        data_analysis_addon = DATA_ANALYSIS_ADDON if 'pandas_analyze_data' in tool_names else ""
+        pyodite_addon = PYODITE_ADDON if 'pyodide_sandbox' in tool_names else ""
+        search_index_addon = SEARCH_INDEX_ADDON if 'stepback_summary_index' in tool_names else ""
+
+        # Select assistant template based on persona
+        persona_templates = {
+            "qa": QA_ASSISTANT,
+            "nerdy": NERDY_ASSISTANT,
+            "quirky": QUIRKY_ASSISTANT,
+            "cynical": CYNICAL_ASSISTANT,
+        }
+
+        # For predict agents with their own instructions, use those directly
+        # instead of wrapping in DEFAULT_ASSISTANT
+        if self.app_type == "openai" and prompt_instructions:
+            # Use agent's own instructions as the base system prompt
+            # Append addons only when their corresponding tools are present
+            addons = "\n\n---\n\n".join(filter(None, [
+                plan_addon,
+                search_index_addon,
+                FILE_HANDLING_INSTRUCTIONS if simple_tools else "",
+                pyodite_addon,
+                data_analysis_addon
+            ]))
+            escaped_prompt = f"{prompt_instructions}\n\n---\n\n{addons}" if addons else str(prompt_instructions)
+            logger.info("Using agent's own instructions directly (app_type=predict)")
+        else:
+            base_assistant = persona_templates.get(self.persona, DEFAULT_ASSISTANT)
+            escaped_prompt = base_assistant.format(
+                users_instructions=user_addon,
+                planning_instructions=plan_addon,
+                pyodite_addon=pyodite_addon,
+                data_analysis_addon=data_analysis_addon,
+                search_index_addon=search_index_addon,
+                file_handling_instructions=FILE_HANDLING_INSTRUCTIONS
+            )
+
         # Properly setup the prompt for YAML
         import yaml
-        escaped_prompt = prompt_instructions or "You are a helpful assistant."
+        
         
         # Create the schema as a dictionary first, then convert to YAML
         state_messages_config = {'type': 'list'}
@@ -341,7 +437,8 @@ class Assistant:
             client=self.client, tools=self.tools,
             yaml_schema=self.prompt, memory=memory,
             alita_client=self.alita_client,
-            steps_limit=self.max_iterations
+            steps_limit=self.max_iterations,
+            for_subgraph=self.is_subgraph  # Pass for_subgraph flag to filter PrinterNodes
         )
         #
         return agent