alita-sdk 0.3.205__py3-none-any.whl → 0.3.207__py3-none-any.whl

alita_sdk/runtime/clients/client.py

@@ -10,6 +10,7 @@ from langchain_core.messages import (
 )
 from langchain_core.tools import ToolException
 from langgraph.store.base import BaseStore
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 
 from ..langchain.assistant import Assistant as LangChainAssistant
 # from ..llamaindex.assistant import Assistant as LLamaAssistant
@@ -37,6 +38,7 @@ class AlitaClient:
 
         self.base_url = base_url.rstrip('/')
         self.api_path = '/api/v1'
+        self.llm_path = '/llm/v1'
         self.project_id = project_id
         self.auth_token = auth_token
         self.headers = {
@@ -152,6 +154,35 @@ class AlitaClient:
             return resp.json()
         return []
 
+    def get_llm(self, model_name: str, model_config: dict) -> ChatOpenAI:
+        """
+        Get a ChatOpenAI model instance based on the model name and configuration.
+        
+        Args:
+            model_name: Name of the model to retrieve
+            model_config: Configuration parameters for the model
+            
+        Returns:
+            An instance of ChatOpenAI configured with the provided parameters.
+        """
+        if not model_name:
+            raise ValueError("Model name must be provided")
+        
+        logger.info(f"Creating ChatOpenAI model: {model_name} with config: {model_config}")
+        
+        return ChatOpenAI(
+            base_url=f"{self.base_url}{self.llm_path}",
+            model=model_name,
+            api_key=self.auth_token,
+            stream_usage=model_config.get("stream_usage", True),
+            max_tokens=model_config.get("max_tokens", None),
+            top_p=model_config.get("top_p"),
+            temperature=model_config.get("temperature"),
+            max_retries=model_config.get("max_retries", 3),
+            seed=model_config.get("seed", None),
+        )
+         
+    
     def get_app_version_details(self, application_id: int, application_version_id: int) -> dict:
         url = f"{self.application_versions}/{application_id}/{application_version_id}"
         if self.configurations:
@@ -177,11 +208,12 @@ class AlitaClient:
         logger.info(f"Unsecret response: {data}")
         return data.get('value', None)
 
-    def application(self, client: Any, application_id: int, application_version_id: int,
+    def application(self, application_id: int, application_version_id: int,
                     tools: Optional[list] = None, chat_history: Optional[List[Any]] = None,
                     app_type=None, memory=None, runtime='langchain',
                     application_variables: Optional[dict] = None,
-                    version_details: Optional[dict] = None, store: Optional[BaseStore] = None):
+                    version_details: Optional[dict] = None, store: Optional[BaseStore] = None,
+                    llm: Optional[ChatOpenAI] = None):
         if tools is None:
             tools = []
         if chat_history is None:
@@ -200,7 +232,15 @@ class AlitaClient:
             for var in data.get('variables', {}):
                 if var['name'] in application_variables:
                     var.update(application_variables[var['name']])
-
+        if llm is None:
+            llm = self.get_llm(
+                model_name=data['llm_settings']['model_name'],
+                model_config = {
+                    "max_tokens": data['llm_settings']['max_tokens'],
+                    "top_p": data['llm_settings']['top_p'],
+                    "temperature": data['llm_settings']['temperature']
+                }
+            )
         if not app_type:
             app_type = data.get("agent_type", "react")
         if app_type == "alita":
@@ -212,10 +252,10 @@ class AlitaClient:
         elif app_type == 'autogen':
             app_type = "openai"
         if runtime == 'nonrunnable':
-            return LangChainAssistant(self, data, client, chat_history, app_type,
+            return LangChainAssistant(self, data, llm, chat_history, app_type,
                                       tools=tools, memory=memory, store=store)
         if runtime == 'langchain':
-            return LangChainAssistant(self, data, client,
+            return LangChainAssistant(self, data, llm,
                                       chat_history, app_type,
                                       tools=tools, memory=memory, store=store).runnable()
         elif runtime == 'llama':
@@ -433,15 +473,15 @@ class AlitaClient:
             logger.warning(f"Error: Could not determine user ID for MCP tool: {e}")
             return None
 
-    def predict_agent(self, client: Any, instructions: str = "You are a helpful assistant.",
+    def predict_agent(self, llm: ChatOpenAI, instructions: str = "You are a helpful assistant.",
                       tools: Optional[list] = None, chat_history: Optional[List[Any]] = None,
                       memory=None, runtime='langchain', variables: Optional[list] = None,
                       store: Optional[BaseStore] = None):
         """
         Create a predict-type agent with minimal configuration.
-        
+
         Args:
-            client: The LLM client to use
+            llm: The LLM to use
             instructions: System instructions for the agent
             tools: Optional list of tools to provide to the agent
             chat_history: Optional chat history
@@ -449,7 +489,7 @@ class AlitaClient:
             runtime: Runtime type (default: 'langchain')
             variables: Optional list of variables for the agent
             store: Optional store for memory
-            
+
         Returns:
             Runnable agent ready for execution
         """
@@ -459,7 +499,7 @@ class AlitaClient:
             chat_history = []
         if variables is None:
             variables = []
-            
+
         # Create a minimal data structure for predict agent
         # All LLM settings are taken from the passed client instance
         agent_data = {
@@ -467,6 +507,269 @@ class AlitaClient:
             'tools': tools,  # Tools are handled separately in predict agents
             'variables': variables
         }
-        return LangChainAssistant(self, agent_data, client,
+        return LangChainAssistant(self, agent_data, llm,
                                   chat_history, "predict", memory=memory, store=store).runnable()
+
+    def test_toolkit_tool(self, toolkit_config: dict, tool_name: str, tool_params: dict = None, 
+                          runtime_config: dict = None, llm_model: str = None, 
+                          llm_config: dict = None) -> dict:
+        """
+        Test a single tool from a toolkit with given parameters and runtime callbacks.
+        
+        This method initializes a toolkit, calls a specific tool, and supports runtime
+        callbacks for event dispatching, enabling tools to send custom events back to
+        the platform during execution.
+        
+        Args:
+            toolkit_config: Configuration dictionary for the toolkit containing:
+                - toolkit_name: Name of the toolkit (e.g., 'github', 'jira')
+                - settings: Dictionary containing toolkit-specific settings
+            tool_name: Name of the specific tool to call
+            tool_params: Parameters to pass to the tool (default: empty dict)
+            runtime_config: Runtime configuration with callbacks for events, containing:
+                - callbacks: List of callback handlers for event processing
+                - configurable: Additional configuration parameters
+                - tags: Tags for the execution
+            llm_model: Name of the LLM model to use (default: 'gpt-4o-mini')
+            llm_config: Configuration for the LLM containing:
+                - max_tokens: Maximum tokens for response (default: 1000)
+                - temperature: Temperature for response generation (default: 0.1)
+                - top_p: Top-p value for response generation (default: 1.0)
+            
+        Returns:
+            Dictionary containing:
+                - success: Boolean indicating if the operation was successful
+                - result: The actual result from the tool (if successful)
+                - error: Error message (if unsuccessful)
+                - tool_name: Name of the executed tool
+                - toolkit_config: Original toolkit configuration
+                - events_dispatched: List of custom events dispatched during execution
+                - llm_model: LLM model used for the test
+                - execution_time_seconds: Time taken to execute the tool in seconds
+            
+        Example:
+            >>> from langchain_core.callbacks import BaseCallbackHandler
+            >>> 
+            >>> class TestCallback(BaseCallbackHandler):
+            ...     def __init__(self):
+            ...         self.events = []
+            ...     def on_custom_event(self, name, data, **kwargs):
+            ...         self.events.append({'name': name, 'data': data})
+            >>> 
+            >>> callback = TestCallback()
+            >>> runtime_config = {'callbacks': [callback]}
+            >>> 
+            >>> config = {
+            ...     'toolkit_name': 'github',
+            ...     'settings': {'github_token': 'your_token'}
+            ... }
+            >>> result = client.test_toolkit_tool(
+            ...     config, 'get_repository_info', 
+            ...     {'repo_name': 'alita'}, runtime_config,
+            ...     llm_model='gpt-4o-mini',
+            ...     llm_config={'temperature': 0.1}
+            ... )
+        """
+        if tool_params is None:
+            tool_params = {}
+        if llm_model is None:
+            llm_model = 'gpt-4o-mini'
+        if llm_config is None:
+            llm_config = {
+                'max_tokens': 1024,
+                'temperature': 0.1,
+                'top_p': 1.0
+            }
+            
+        try:
+            from ..utils.toolkit_utils import instantiate_toolkit_with_client
+            from langchain_core.runnables import RunnableConfig
+            import logging
+            import time
+            
+            logger = logging.getLogger(__name__)
+            logger.info(f"Testing tool '{tool_name}' from toolkit '{toolkit_config.get('toolkit_name')}' with LLM '{llm_model}'")
+            
+            # Create RunnableConfig for callback support
+            config = None
+            callbacks = []
+            events_dispatched = []
+            
+            if runtime_config:
+                callbacks = runtime_config.get('callbacks', [])
+                if callbacks:
+                    config = RunnableConfig(
+                        callbacks=callbacks,
+                        configurable=runtime_config.get('configurable', {}),
+                        tags=runtime_config.get('tags', [])
+                    )
+            
+            # Create LLM instance using the client's get_llm method
+            try:
+                llm = self.get_llm(llm_model, llm_config)
+                logger.info(f"Created LLM instance: {llm_model} with config: {llm_config}")
+            except Exception as llm_error:
+                logger.error(f"Failed to create LLM instance: {str(llm_error)}")
+                return {
+                    "success": False,
+                    "error": f"Failed to create LLM instance '{llm_model}': {str(llm_error)}",
+                    "tool_name": tool_name,
+                    "toolkit_config": toolkit_config,
+                    "llm_model": llm_model,
+                    "events_dispatched": events_dispatched,
+                    "execution_time_seconds": 0.0
+                }
+            
+            # Instantiate the toolkit with client and LLM support
+            tools = instantiate_toolkit_with_client(toolkit_config, llm, self)
+            
+            if not tools:
+                return {
+                    "success": False,
+                    "error": f"Failed to instantiate toolkit '{toolkit_config.get('toolkit_name')}' or no tools found",
+                    "tool_name": tool_name,
+                    "toolkit_config": toolkit_config,
+                    "llm_model": llm_model,
+                    "events_dispatched": events_dispatched,
+                    "execution_time_seconds": 0.0
+                }
+            
+            # Find the specific tool
+            target_tool = None
+            for tool in tools:
+                if hasattr(tool, 'name') and tool.name == tool_name:
+                    target_tool = tool
+                    break
+                elif hasattr(tool, 'func') and hasattr(tool.func, '__name__') and tool.func.__name__ == tool_name:
+                    target_tool = tool
+                    break
+            
+            if target_tool is None:
+                available_tools = []
+                for tool in tools:
+                    if hasattr(tool, 'name'):
+                        available_tools.append(tool.name)
+                    elif hasattr(tool, 'func') and hasattr(tool.func, '__name__'):
+                        available_tools.append(tool.func.__name__)
+                
+                return {
+                    "success": False,
+                    "error": f"Tool '{tool_name}' not found. Available tools: {available_tools}",
+                    "tool_name": tool_name,
+                    "toolkit_config": toolkit_config,
+                    "llm_model": llm_model,
+                    "events_dispatched": events_dispatched,
+                    "execution_time_seconds": 0.0
+                }
+            
+            # Execute the tool with callback support
+            try:
+                logger.info(f"Executing tool '{tool_name}' with parameters: {tool_params}")
+                
+                # Start timing the tool execution
+                start_time = time.time()
+                
+                # Different tools might have different invocation patterns
+                if hasattr(target_tool, 'invoke'):
+                    # Use config for tools that support RunnableConfig
+                    if config is not None:
+                        result = target_tool.invoke(tool_params, config=config)
+                    else:
+                        result = target_tool.invoke(tool_params)
+                elif hasattr(target_tool, 'run'):
+                    result = target_tool.run(tool_params)
+                elif callable(target_tool):
+                    result = target_tool(**tool_params)
+                else:
+                    execution_time = time.time() - start_time
+                    return {
+                        "success": False,
+                        "error": f"Tool '{tool_name}' is not callable",
+                        "tool_name": tool_name,
+                        "toolkit_config": toolkit_config,
+                        "llm_model": llm_model,
+                        "events_dispatched": events_dispatched,
+                        "execution_time_seconds": execution_time
+                    }
+                
+                # Calculate execution time
+                execution_time = time.time() - start_time
+                
+                # Extract events from callbacks if they support it
+                for callback in callbacks:
+                    if hasattr(callback, 'events'):
+                        events_dispatched.extend(callback.events)
+                    elif hasattr(callback, 'get_events'):
+                        events_dispatched.extend(callback.get_events())
+                    elif hasattr(callback, 'dispatched_events'):
+                        events_dispatched.extend(callback.dispatched_events)
+                
+                logger.info(f"Tool '{tool_name}' executed successfully in {execution_time:.3f} seconds")
+                
+                return {
+                    "success": True,
+                    "result": result,
+                    "tool_name": tool_name,
+                    "toolkit_config": toolkit_config,
+                    "llm_model": llm_model,
+                    "events_dispatched": events_dispatched,
+                    "execution_time_seconds": execution_time
+                }
+                
+            except Exception as tool_error:
+                # Calculate execution time even for failed executions
+                execution_time = time.time() - start_time
+                logger.error(f"Error executing tool '{tool_name}' after {execution_time:.3f} seconds: {str(tool_error)}")
+                
+                # Still collect events even if tool execution failed
+                for callback in callbacks:
+                    if hasattr(callback, 'events'):
+                        events_dispatched.extend(callback.events)
+                    elif hasattr(callback, 'get_events'):
+                        events_dispatched.extend(callback.get_events())
+                    elif hasattr(callback, 'dispatched_events'):
+                        events_dispatched.extend(callback.dispatched_events)
+                
+                return {
+                    "success": False,
+                    "error": f"Tool execution failed: {str(tool_error)}",
+                    "tool_name": tool_name,
+                    "toolkit_config": toolkit_config,
+                    "llm_model": llm_model,
+                    "events_dispatched": events_dispatched,
+                    "execution_time_seconds": execution_time
+                }
+                
+        except Exception as e:
+            logger = logging.getLogger(__name__)
+            logger.error(f"Error in test_toolkit_tool: {str(e)}")
+            return {
+                "success": False,
+                "error": f"Method execution failed: {str(e)}",
+                "tool_name": tool_name,
+                "toolkit_config": toolkit_config,
+                "llm_model": llm_model if 'llm_model' in locals() else None,
+                "events_dispatched": [],
+                "execution_time_seconds": 0.0
+            }
+    
+    def _get_real_user_id(self) -> str:
+        """Extract the real user ID from the auth token for MCP tool calls."""
+        try:
+            import base64
+            import json
+            # Assuming JWT token, extract user ID from payload
+            # This is a basic implementation - adjust based on your token format
+            token_parts = self.auth_token.split('.')
+            if len(token_parts) >= 2:
+                payload_part = token_parts[1]
+                # Add padding if needed
+                padding = len(payload_part) % 4
+                if padding:
+                    payload_part += '=' * (4 - padding)
+                payload = json.loads(base64.b64decode(payload_part))
+                return payload.get('user_id') or payload.get('sub') or payload.get('uid')
+        except Exception as e:
+            logger.error(f"Error extracting user ID from token: {e}")
+        return None
         

alita_sdk/runtime/langchain/assistant.py

@@ -36,29 +36,30 @@ class Assistant:
 
         logger.debug("Data for agent creation: %s", data)
         logger.info("App type: %s", app_type)
-
+        
+        self.client = client
         # For predict agents, use the client as-is since it's already configured
-        if app_type == "predict":
-            self.client = client
-        else:
-            # For other agent types, configure client from llm_settings
-            self.client = copy(client)
-            self.client.max_tokens = data['llm_settings']['max_tokens']
-            self.client.temperature = data['llm_settings']['temperature']
-            self.client.top_p = data['llm_settings']['top_p']
-            self.client.top_k = data['llm_settings']['top_k']
-            self.client.model_name = data['llm_settings']['model_name']
-            self.client.integration_uid = data['llm_settings']['integration_uid']
+        # if app_type == "predict":
+        #     self.client = client
+        # else:
+        #     # For other agent types, configure client from llm_settings
+        #     self.client = copy(client)
+        #     self.client.max_tokens = data['llm_settings']['max_tokens']
+        #     self.client.temperature = data['llm_settings']['temperature']
+        #     self.client.top_p = data['llm_settings']['top_p']
+        #     self.client.top_k = data['llm_settings']['top_k']
+        #     self.client.model_name = data['llm_settings']['model_name']
+        #     self.client.integration_uid = data['llm_settings']['integration_uid']
 
-            model_type = data["llm_settings"]["indexer_config"]["ai_model"]
-            model_params = data["llm_settings"]["indexer_config"]["ai_model_params"]
-            #
-            target_pkg, target_name = model_type.rsplit(".", 1)
-            target_cls = getattr(
-                importlib.import_module(target_pkg),
-                target_name
-            )
-            self.client = target_cls(**model_params)
+        #     model_type = data["llm_settings"]["indexer_config"]["ai_model"]
+        #     model_params = data["llm_settings"]["indexer_config"]["ai_model_params"]
+        #     #
+        #     target_pkg, target_name = model_type.rsplit(".", 1)
+        #     target_cls = getattr(
+        #         importlib.import_module(target_pkg),
+        #         target_name
+        #     )
+        #     self.client = target_cls(**model_params)
         # validate agents compatibility: non-pipeline agents cannot have pipelines as toolkits
         if app_type not in ["pipeline", "predict"]:
             tools_to_check = data.get('tools', [])

alita_sdk/runtime/langchain/interfaces/llm_processor.py

@@ -50,9 +50,6 @@ def get_model(model_type: str, model_params: dict):
         return get_llm(model_type)(**model_params)
     if model_type == "PreloadedChatModel":
         return PreloadedChatModel(**model_params)
-    if model_type == "Alita":
-        from ...llms.alita import AlitaClient
-        return AlitaClient(**model_params)
     if model_type in chat_models:
         model = getattr(
             __import__("langchain_community.chat_models", fromlist=[model_type]),
@@ -185,7 +182,7 @@ def add_documents(vectorstore, documents):
         texts.append(document.page_content)
         for key in document.metadata:
             if isinstance(document.metadata[key], list):
-                document.metadata[key] = "; ".join(document.metadata[key])
+                document.metadata[key] = "; ".join([str(val) for val in document.metadata[key]])
             if isinstance(document.metadata[key], dict):
                 document.metadata[key] = dumps(document.metadata[key])
         metadata.append(document.metadata)

alita_sdk/runtime/langchain/langraph_agent.py

@@ -500,7 +500,12 @@ def create_graph(
                 }
                 
                 # Check if tools should be bound to this LLM node
-                tool_names = node.get('tool_names', []) if isinstance(node.get('tool_names'), list) else []
+                connected_tools = node.get('tool_names', {})
+                tool_names = []
+                if isinstance(connected_tools, dict):
+                    for toolkit, selected_tools in connected_tools.items():
+                        for tool in selected_tools:
+                            tool_names.append(f"{toolkit}___{tool}")
                 
                 # Filter tools if specific tool names are provided
                 available_tools = []

alita_sdk/runtime/langchain/store_manager.py

@@ -3,9 +3,6 @@ import atexit
 import logging
 from urllib.parse import urlparse, unquote
 
-from psycopg import Connection
-from langgraph.store.postgres import PostgresStore
-
 logger = logging.getLogger(__name__)
 
 class StoreManager:
@@ -37,7 +34,10 @@ class StoreManager:
             "dbname": parsed.path.lstrip("/") if parsed.path else None
         }
 
-    def get_store(self, conn_str: str) -> PostgresStore:
+    def get_store(self, conn_str: str):
+        from psycopg import Connection
+        from langgraph.store.postgres import PostgresStore
+        
         store = self._stores.get(conn_str)
         if store is None:
             logger.info(f"Creating new PostgresStore for connection: {conn_str}")

alita_sdk/runtime/toolkits/application.py

@@ -21,25 +21,20 @@ class ApplicationToolkit(BaseToolkit):
         )
     
     @classmethod
-    def get_toolkit(cls, client: Any, application_id: int, application_version_id: int, app_api_key: str,
+    def get_toolkit(cls, client: 'AlitaClient', application_id: int, application_version_id: int,
                     selected_tools: list[str] = [], store: Optional[BaseStore] = None):
-        from ..llms.alita import AlitaChatModel
         
         app_details = client.get_app_details(application_id)
         version_details = client.get_app_version_details(application_id, application_version_id)
-        settings = {
-            "deployment": client.base_url,
-            "model": version_details['llm_settings']['model_name'],
-            "api_key": app_api_key,
-            "project_id": client.project_id,
-            "integration_uid": version_details['llm_settings']['integration_uid'],
+        model_settings = {
             "max_tokens": version_details['llm_settings']['max_tokens'],
             "top_p": version_details['llm_settings']['top_p'],
-            "top_k": version_details['llm_settings']['top_k'],
             "temperature": version_details['llm_settings']['temperature'],
         }
 
-        app = client.application(AlitaChatModel(**settings), application_id, application_version_id, store=store)
+        app = client.application(application_id, application_version_id, store=store, 
+                                 llm=client.get_llm(version_details['llm_settings']['model_name'], 
+                                                    model_settings))
         return cls(tools=[Application(name=app_details.get("name"), 
                                       description=app_details.get("description"), 
                                       application=app, 

alita_sdk/runtime/toolkits/tools.py

@@ -14,18 +14,15 @@ from .vectorstore import VectorStoreToolkit
 from ..tools.mcp_server_tool import McpServerTool
 # Import community tools
 from ...community import get_toolkits as community_toolkits, get_tools as community_tools
-# from ...tools.memory import MemoryToolkit
+from ...tools.memory import MemoryToolkit
 
 logger = logging.getLogger(__name__)
 
 
 def get_toolkits():
     core_toolkits = [
-        # PromptToolkit.toolkit_config_schema(),
-        # DatasourcesToolkit.toolkit_config_schema(),
-        # ApplicationToolkit.toolkit_config_schema(),
         ArtifactToolkit.toolkit_config_schema(),
-        # MemoryToolkit.toolkit_config_schema(),
+        MemoryToolkit.toolkit_config_schema(),
         VectorStoreToolkit.toolkit_config_schema()
     ]
 
@@ -37,12 +34,7 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
     tools = []
 
     for tool in tools_list:
-        if tool['type'] == 'prompt':
-            prompts.append([
-                int(tool['settings']['prompt_id']),
-                int(tool['settings']['prompt_version_id'])
-            ])
-        elif tool['type'] == 'datasource':
+        if tool['type'] == 'datasource':
             tools.extend(DatasourcesToolkit.get_toolkit(
                 alita_client,
                 datasource_ids=[int(tool['settings']['datasource_id'])],
@@ -54,7 +46,6 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                 alita_client,
                 application_id=int(tool['settings']['application_id']),
                 application_version_id=int(tool['settings']['application_version_id']),
-                app_api_key=alita_client.auth_token,
                 selected_tools=[]
             ).get_tools())
         elif tool['type'] == 'application' and tool.get('agent_type', '') == 'pipeline':
@@ -67,15 +58,14 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                 selected_tools=[],
                 llm=llm
             ))
-            # move on tools level
-        # elif tool['type'] == 'memory':
-        #     if memory_store is None:
-        #         raise ToolException(f"Memory store is not provided for memory tool: {tool['name']}")
-        #     tools += MemoryToolkit.get_toolkit(
-        #         namespace=tool['settings'].get('namespace', str(tool['id'])),
-        #         store=memory_store,
-        #         toolkit_name=tool.get('toolkit_name', '')
-        #     ).get_tools()
+        elif tool['type'] == 'memory':
+            if memory_store is None:
+                raise ToolException(f"Memory store is not provided for memory tool: {tool.get('name', tool.get('toolkit_name', 'unknown'))}")
+            tools += MemoryToolkit.get_toolkit(
+                namespace=tool['settings'].get('namespace', str(tool['id'])),
+                store=memory_store,
+                toolkit_name=tool.get('toolkit_name', '')
+            ).get_tools()
         elif tool['type'] == 'artifact':
             tools.extend(ArtifactToolkit.get_toolkit(
                 client=alita_client,