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

alita_sdk/runtime/langchain/assistant.py

@@ -13,10 +13,14 @@ 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
+from jinja2 import Environment, DebugUndefined
 
 logger = logging.getLogger(__name__)
 
@@ -29,12 +33,20 @@ class Assistant:
                  app_type: str = "openai",
                  tools: Optional[list] = [],
                  memory: Optional[Any] = None,
-                 store: Optional[BaseStore] = None):
+                 store: Optional[BaseStore] = None,
+                 debug_mode: Optional[bool] = False,
+                 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)
@@ -84,12 +96,81 @@ 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, alita_client=alita, llm=self.client, memory_store=self.store)
+        self.tools = get_tools(
+            version_tools,
+            alita_client=alita,
+            llm=self.client,
+            memory_store=self.store,
+            debug_mode=debug_mode,
+            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']
@@ -184,17 +265,6 @@ class Assistant:
         """
         # Exclude compiled graph runnables from simple tool agents
         simple_tools = [t for t in self.tools if isinstance(t, (BaseTool, CompiledStateGraph))]
-        # Add image generation tool if model is configured
-        if self.alita_client.model_image_generation is not None:
-            try:
-                from ..tools.image_generation import (
-                    create_image_generation_tool
-                )
-                image_tool = create_image_generation_tool(self.alita_client)
-                simple_tools.append(image_tool)
-                logger.info("Added ImageGenerationTool to react agent")
-            except Exception as e:
-                logger.error(f"Error adding ImageGenerationTool: {e}")
         
         # Set up memory/checkpointer if available
         checkpointer = None
@@ -227,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'}
@@ -274,6 +369,9 @@ class Assistant:
         schema_dict = {
             'name': 'react_agent',
             'state': {
+                'input': {
+                    'type': 'str'
+                },
                 'messages': state_messages_config
             },
             'nodes': [{
@@ -282,6 +380,21 @@ class Assistant:
                 'prompt': {
                     'template': escaped_prompt
                 },
+                'input_mapping': {
+                    'system': {
+                        'type': 'fixed',
+                        'value': escaped_prompt
+                    },
+                    'task': {
+                        'type': 'variable',
+                        'value': 'input'
+                    },
+                    'chat_history': {
+                        'type': 'variable',
+                        'value': 'messages'
+                    }
+                },
+                'step_limit': self.max_iterations,
                 'input': ['messages'],
                 'output': ['messages'],
                 'transition': 'END'
@@ -307,7 +420,8 @@ class Assistant:
             store=self.store,
             debug=False,
             for_subgraph=False,
-            alita_client=self.alita_client
+            alita_client=self.alita_client,
+            steps_limit=self.max_iterations
         )
         
         return agent
@@ -322,7 +436,9 @@ class Assistant:
         agent = create_graph(
             client=self.client, tools=self.tools,
             yaml_schema=self.prompt, memory=memory,
-            alita_client=self.alita_client
+            alita_client=self.alita_client,
+            steps_limit=self.max_iterations,
+            for_subgraph=self.is_subgraph  # Pass for_subgraph flag to filter PrinterNodes
         )
         #
         return agent
@@ -342,5 +458,7 @@ class Assistant:
                     continue
                 # take only the content of the system message from the openai prompt
                 if isinstance(message, SystemMessage):
-                    return message.content
+                    environment = Environment(undefined=DebugUndefined)
+                    template = environment.from_string(message.content)
+                    return template.render(self.prompt.partial_variables)
         return None