alita-sdk 0.3.149__tar.gz → 0.3.151__tar.gz

{alita_sdk-0.3.149 → alita_sdk-0.3.151}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.149
+Version: 0.3.151
 Summary: SDK for building langchain agents using resources from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedjik@gmail.com>
 License-Expression: Apache-2.0

{alita_sdk-0.3.149 → alita_sdk-0.3.151}/alita_sdk/__init__.py

@@ -9,13 +9,13 @@ This package contains three main modules:
 
 __version__ = "0.3.142"
 
-# Import key components
-from .runtime import *
-from .tools import *
-from .community import *
+__all__ = ["runtime", "tools", "community"]
 
-__all__ = [
-    "runtime",
-    "tools",
-    "community",
-]
+import importlib
+
+def __getattr__(name):
+    if name in __all__:
+        module = importlib.import_module(f".{name}", __name__)
+        globals()[name] = module
+        return module
+    raise AttributeError(f"module {__name__} has no attribute {name}")

alita_sdk-0.3.151/alita_sdk/runtime/__init__.py

@@ -0,0 +1,33 @@
+"""
+Alita SDK Runtime Module
+
+This module contains the core runtime functionality for building langchain agents.
+Includes agents, clients, language models, and utilities.
+"""
+
+import importlib
+
+_modules = [
+    "agents",
+    "clients",
+    "langchain",
+    "llamaindex",
+    "llms",
+    "toolkits",
+    "tools",
+    "utils",
+]
+
+__all__ = _modules + ["get_tools", "get_toolkits"]
+
+def __getattr__(name):
+    if name in _modules:
+        module = importlib.import_module(f".{name}", __name__)
+        globals()[name] = module
+        return module
+    if name in {"get_tools", "get_toolkits"}:
+        toolkits = importlib.import_module(".toolkits.tools", __name__)
+        value = getattr(toolkits, name)
+        globals()[name] = value
+        return value
+    raise AttributeError(f"module {__name__} has no attribute {name}")

alita_sdk-0.3.151/alita_sdk/runtime/clients/__init__.py

@@ -0,0 +1,12 @@
+"""
+Runtime clients package.
+"""
+
+try:
+    from .client import AlitaClient
+    __all__ = ['AlitaClient']
+except ImportError as e:
+    # Handle case where dependencies are not available
+    import logging
+    logging.getLogger(__name__).debug(f"Failed to import AlitaClient: {e}")
+    __all__ = []

{alita_sdk-0.3.149 → alita_sdk-0.3.151}/alita_sdk/runtime/clients/client.py

@@ -9,6 +9,7 @@ from langchain_core.messages import (
     SystemMessage, BaseMessage,
 )
 from langchain_core.tools import ToolException
+from langgraph.store.base import BaseStore
 
 from ..langchain.assistant import Assistant as LangChainAssistant
 # from ..llamaindex.assistant import Assistant as LLamaAssistant
@@ -178,7 +179,7 @@ class AlitaClient:
                     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):
+                    version_details: Optional[dict] = None, store: Optional[BaseStore] = None):
         if tools is None:
             tools = []
         if chat_history is None:
@@ -210,11 +211,11 @@ class AlitaClient:
             app_type = "openai"
         if runtime == 'nonrunnable':
             return LangChainAssistant(self, data, client, chat_history, app_type,
-                                      tools=tools, memory=memory)
+                                      tools=tools, memory=memory, store=store)
         if runtime == 'langchain':
             return LangChainAssistant(self, data, client,
                                       chat_history, app_type,
-                                      tools=tools, memory=memory).runnable()
+                                      tools=tools, memory=memory, store=store).runnable()
         elif runtime == 'llama':
             raise NotImplementedError("LLama runtime is not supported")
 

{alita_sdk-0.3.149 → alita_sdk-0.3.151}/alita_sdk/runtime/langchain/assistant.py

@@ -5,8 +5,9 @@ from typing import Any, Optional
 from langchain.agents import (
     AgentExecutor, create_openai_tools_agent,
     create_json_chat_agent)
+from langgraph.store.base import BaseStore
+
 from .agents.xml_chat import create_xml_chat_agent
-# from ..langchain.mixedAgentRenderes import render_react_text_description_and_args
 from .langraph_agent import create_graph
 from langchain_core.messages import (
     BaseMessage, SystemMessage, HumanMessage
@@ -28,7 +29,8 @@ class Assistant:
                  chat_history: list[BaseMessage] = [],
                  app_type: str = "openai",
                  tools: Optional[list] = [],
-                 memory: Optional[Any] = None):
+                 memory: Optional[Any] = None,
+                 store: Optional[BaseStore] = None):
 
         self.client = copy(client)
         self.client.max_tokens = data['llm_settings']['max_tokens']
@@ -40,6 +42,7 @@ class Assistant:
 
         self.app_type = app_type
         self.memory = memory
+        self.store = store
 
         logger.debug("Data for agent creation: %s", data)
         logger.info("App type: %s", app_type)
@@ -58,7 +61,10 @@ class Assistant:
             raise ToolException("Non-pipeline agents cannot have pipelines as a toolkits. "
                                 "Review toolkits configuration or use pipeline as master agent.")
 
-        self.tools = get_tools(data['tools'], alita_client=alita, llm=self.client)
+        # configure memory store if memory tool is defined
+        memory_tool = next((tool for tool in data['tools'] if tool['type'] == 'memory'), None)
+        self._configure_store(memory_tool)
+        self.tools = get_tools(data['tools'], alita_client=alita, llm=self.client, memory_store=self.store)
         if app_type == "pipeline":
             self.prompt = data['instructions']
         else:
@@ -94,6 +100,18 @@ class Assistant:
             except Exception as e:
                 logger.info(f"Client was created with client setting: temperature - {self.client.temperature} : {self.client.max_tokens}")
 
+    def _configure_store(self, memory_tool: dict | None) -> None:
+        """
+        Configure the memory store based on a memory_tool definition.
+        Only creates a new store if one does not already exist.
+        """
+        if not memory_tool or self.store is not None:
+            return
+        from .store_manager import get_manager
+        conn_str = memory_tool.get('settings', {}).get('connection_string', '')
+        store = get_manager().get_store(conn_str)
+        self.store = store
+
     def runnable(self):
         if self.app_type == 'pipeline':
             return self.pipeline()

alita_sdk-0.3.151/alita_sdk/runtime/langchain/store_manager.py

@@ -0,0 +1,46 @@
+import threading
+import atexit
+import logging
+from psycopg import Connection
+from langgraph.store.postgres import PostgresStore
+
+logger = logging.getLogger(__name__)
+
+class StoreManager:
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._stores = {}
+        return cls._instance
+
+    def get_store(self, conn_str: str) -> PostgresStore:
+        store = self._stores.get(conn_str)
+        if store is None:
+            logger.info(f"Creating new PostgresStore for connection: {conn_str}")
+            conn = Connection.connect(conn_str, autocommit=True, prepare_threshold=0)
+            store = PostgresStore(conn)
+            store.setup()
+            self._stores[conn_str] = store
+        return store
+
+    def shutdown(self) -> None:
+        logger.info("Shutting down StoreManager and closing all stores")
+        for store in list(self._stores.values()):
+            try:
+                conn = getattr(store, 'conn', None)
+                if conn:
+                    conn.close()
+            except Exception:
+                pass
+        self._stores.clear()
+
+_store_manager = StoreManager()
+atexit.register(_store_manager.shutdown)
+
+def get_manager() -> StoreManager:
+    return _store_manager

{alita_sdk-0.3.149 → alita_sdk-0.3.151}/alita_sdk/runtime/llms/preloaded.py

@@ -30,7 +30,11 @@ from langchain_core.language_models import BaseChatModel  # pylint: disable=E040
 from langchain_core.messages import AIMessage, AIMessageChunk  # pylint: disable=E0401
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult  # pylint: disable=E0401
 
-from ..langchain.tools import log
+try:
+    from ..langchain.tools import log
+except ImportError:
+    import logging as _logging
+    log = _logging.getLogger(__name__)
 
 
 class PreloadedEmbeddings(Embeddings):

{alita_sdk-0.3.149 → alita_sdk-0.3.151}/alita_sdk/runtime/toolkits/application.py

@@ -1,4 +1,6 @@
 from typing import List, Any, Optional
+
+from langgraph.store.base import BaseStore
 from pydantic import create_model, BaseModel, Field
 from langchain_community.agent_toolkits.base import BaseToolkit
 from langchain_core.tools import BaseTool
@@ -19,8 +21,8 @@ class ApplicationToolkit(BaseToolkit):
         )
     
     @classmethod
-    def get_toolkit(cls, client: Any, application_id: int, application_version_id: int, app_api_key: str, 
-                    selected_tools: list[str] = []):
+    def get_toolkit(cls, client: Any, application_id: int, application_version_id: int, app_api_key: str,
+                    selected_tools: list[str] = [], store: Optional[BaseStore] = None):
         from ..llms.alita import AlitaChatModel
         
         app_details = client.get_app_details(application_id)
@@ -37,7 +39,7 @@ class ApplicationToolkit(BaseToolkit):
             "temperature": version_details['llm_settings']['temperature'],
         }
 
-        app = client.application(AlitaChatModel(**settings), application_id, application_version_id)
+        app = client.application(AlitaChatModel(**settings), application_id, application_version_id, store=store)
         return cls(tools=[Application(name=app_details.get("name"), 
                                       description=app_details.get("description"), 
                                       application=app, 

{alita_sdk-0.3.149 → alita_sdk-0.3.151}/alita_sdk/runtime/toolkits/tools.py

@@ -1,19 +1,20 @@
 import logging
 
+from langchain_core.tools import ToolException
+from langgraph.store.base import BaseStore
+
 from alita_sdk.tools import get_toolkits as alita_toolkits
 from alita_sdk.tools import get_tools as alita_tools
-
 from .application import ApplicationToolkit
 from .artifact import ArtifactToolkit
 from .datasource import DatasourcesToolkit
 from .prompt import PromptToolkit
 from .subgraph import SubgraphToolkit
 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.mcp_server_tool import McpServerTool
+from ...tools.memory import MemoryToolkit
 
 logger = logging.getLogger(__name__)
 
@@ -24,13 +25,14 @@ def get_toolkits():
         # DatasourcesToolkit.toolkit_config_schema(),
         # ApplicationToolkit.toolkit_config_schema(),
         ArtifactToolkit.toolkit_config_schema(),
+        MemoryToolkit.toolkit_config_schema(),
         VectorStoreToolkit.toolkit_config_schema()
     ]
 
     return core_toolkits + community_toolkits() + alita_toolkits()
 
 
-def get_tools(tools_list: list, alita_client, llm) -> list:
+def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None) -> list:
     prompts = []
     tools = []
 
@@ -65,6 +67,14 @@ def get_tools(tools_list: list, alita_client, llm) -> list:
                 selected_tools=[],
                 llm=llm
             ))
+        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'] == 'artifact':
             tools.extend(ArtifactToolkit.get_toolkit(
                 client=alita_client,

{alita_sdk-0.3.149 → alita_sdk-0.3.151}/alita_sdk/runtime/utils/logging.py

@@ -1,7 +1,12 @@
 import logging
 from functools import wraps
 
-from langchain_core.callbacks import dispatch_custom_event
+try:
+    from langchain_core.callbacks import dispatch_custom_event
+except ImportError:
+    # Fallback stub if langchain_core is unavailable
+    def dispatch_custom_event(name: str, data: dict):  # pragma: no cover
+        pass
 
 
 class StreamlitCallbackHandler(logging.Handler):
@@ -42,6 +47,7 @@ def setup_streamlit_logging(
         formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
         handler.setFormatter(formatter)
         logger.addHandler(handler)
+        logger.setLevel(logging.INFO)
 
     return handler
 

{alita_sdk-0.3.149 → alita_sdk-0.3.151}/alita_sdk/runtime/utils/save_dataframe.py

@@ -6,13 +6,16 @@ from typing import Any, Dict, Optional
 from langchain_core.tools import ToolException
 import pandas as pd
 
-from ..tools.artifact import ArtifactWrapper
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:  # pragma: no cover - avoid heavy imports at runtime
+    from ..tools.artifact import ArtifactWrapper
 
 logger = logging.getLogger(__name__)
 
 
 def save_dataframe_to_artifact(
-    artifacts_wrapper: ArtifactWrapper,
+    artifacts_wrapper: 'ArtifactWrapper',
     df: pd.DataFrame,
     target_file: str,
     csv_options: Optional[Dict[str, Any]] = None,

alita_sdk-0.3.151/alita_sdk/tools/__init__.py

@@ -0,0 +1,205 @@
+import logging
+from importlib import import_module
+
+logger = logging.getLogger(__name__)
+
+# Available tools and toolkits - populated by safe imports
+AVAILABLE_TOOLS = {}
+AVAILABLE_TOOLKITS = {}
+FAILED_IMPORTS = {}
+
+def _safe_import_tool(tool_name, module_path, get_tools_name=None, toolkit_class_name=None):
+    """Safely import a tool module and register available functions/classes."""
+    try:
+        module = __import__(f'alita_sdk.tools.{module_path}', fromlist=[''])
+        
+        imported = {}
+        if get_tools_name and hasattr(module, get_tools_name):
+            imported['get_tools'] = getattr(module, get_tools_name)
+            
+        if toolkit_class_name and hasattr(module, toolkit_class_name):
+            imported['toolkit_class'] = getattr(module, toolkit_class_name)
+            AVAILABLE_TOOLKITS[toolkit_class_name] = getattr(module, toolkit_class_name)
+            
+        if imported:
+            AVAILABLE_TOOLS[tool_name] = imported
+            logger.debug(f"Successfully imported {tool_name}")
+        
+    except Exception as e:
+        FAILED_IMPORTS[tool_name] = str(e)
+        logger.debug(f"Failed to import {tool_name}: {e}")
+
+# Safe imports for all tools
+_safe_import_tool('github', 'github', 'get_tools', 'AlitaGitHubToolkit')
+_safe_import_tool('openapi', 'openapi', 'get_tools')
+_safe_import_tool('jira', 'jira', 'get_tools', 'JiraToolkit')
+_safe_import_tool('confluence', 'confluence', 'get_tools', 'ConfluenceToolkit')
+_safe_import_tool('service_now', 'servicenow', 'get_tools', 'ServiceNowToolkit')
+_safe_import_tool('gitlab', 'gitlab', 'get_tools', 'AlitaGitlabToolkit')
+_safe_import_tool('gitlab_org', 'gitlab_org', 'get_tools', 'AlitaGitlabSpaceToolkit')
+_safe_import_tool('zephyr', 'zephyr', 'get_tools', 'ZephyrToolkit')
+_safe_import_tool('browser', 'browser', 'get_tools', 'BrowserToolkit')
+_safe_import_tool('report_portal', 'report_portal', 'get_tools', 'ReportPortalToolkit')
+_safe_import_tool('bitbucket', 'bitbucket', 'get_tools', 'AlitaBitbucketToolkit')
+_safe_import_tool('testrail', 'testrail', 'get_tools', 'TestrailToolkit')
+_safe_import_tool('testio', 'testio', 'get_tools', 'TestIOToolkit')
+_safe_import_tool('xray_cloud', 'xray', 'get_tools', 'XrayToolkit')
+_safe_import_tool('sharepoint', 'sharepoint', 'get_tools', 'SharepointToolkit')
+_safe_import_tool('qtest', 'qtest', 'get_tools', 'QtestToolkit')
+_safe_import_tool('zephyr_scale', 'zephyr_scale', 'get_tools', 'ZephyrScaleToolkit')
+_safe_import_tool('zephyr_enterprise', 'zephyr_enterprise', 'get_tools', 'ZephyrEnterpriseToolkit')
+_safe_import_tool('ado', 'ado', 'get_tools')
+_safe_import_tool('ado_repos', 'ado.repos', 'get_tools', 'AzureDevOpsReposToolkit')
+_safe_import_tool('ado_plans', 'ado.test_plan', None, 'AzureDevOpsPlansToolkit')
+_safe_import_tool('ado_work_items', 'ado.work_item', None, 'AzureDevOpsWorkItemsToolkit')
+_safe_import_tool('ado_wiki', 'ado.wiki', None, 'AzureDevOpsWikiToolkit')
+_safe_import_tool('rally', 'rally', 'get_tools', 'RallyToolkit')
+_safe_import_tool('sql', 'sql', 'get_tools', 'SQLToolkit')
+_safe_import_tool('sonar', 'code.sonar', 'get_tools', 'SonarToolkit')
+_safe_import_tool('google_places', 'google_places', 'get_tools', 'GooglePlacesToolkit')
+_safe_import_tool('yagmail', 'yagmail', 'get_tools', 'AlitaYagmailToolkit')
+_safe_import_tool('aws', 'cloud.aws', None, 'AWSToolkit')
+_safe_import_tool('azure', 'cloud.azure', None, 'AzureToolkit')
+_safe_import_tool('gcp', 'cloud.gcp', None, 'GCPToolkit')
+_safe_import_tool('k8s', 'cloud.k8s', None, 'KubernetesToolkit')
+_safe_import_tool('custom_open_api', 'custom_open_api', None, 'OpenApiToolkit')
+_safe_import_tool('elastic', 'elastic', None, 'ElasticToolkit')
+_safe_import_tool('keycloak', 'keycloak', None, 'KeycloakToolkit')
+_safe_import_tool('localgit', 'localgit', None, 'AlitaLocalGitToolkit')
+_safe_import_tool('pandas', 'pandas', 'get_tools', 'PandasToolkit')
+_safe_import_tool('azure_search', 'azure_ai.search', 'get_tools', 'AzureSearchToolkit')
+_safe_import_tool('figma', 'figma', 'get_tools', 'FigmaToolkit')
+_safe_import_tool('salesforce', 'salesforce', 'get_tools', 'SalesforceToolkit')
+_safe_import_tool('carrier', 'carrier', 'get_tools', 'AlitaCarrierToolkit')
+_safe_import_tool('ocr', 'ocr', 'get_tools', 'OCRToolkit')
+_safe_import_tool('pptx', 'pptx', 'get_tools', 'PPTXToolkit')
+_safe_import_tool('postman', 'postman', 'get_tools', 'PostmanToolkit')
+
+# Log import summary
+available_count = len(AVAILABLE_TOOLS)
+total_attempted = len(AVAILABLE_TOOLS) + len(FAILED_IMPORTS)
+logger.info(f"Tool imports completed: {available_count}/{total_attempted} successful")
+
+def get_tools(tools_list, alita, llm, *args, **kwargs):
+    tools = []
+    for tool in tools_list:
+        # validate tool name syntax - it cannot be started with _
+        for tool_name in tool.get('settings', {}).get('selected_tools', []):
+            if isinstance(tool_name, str) and tool_name.startswith('_'):
+                raise ValueError(f"Tool name '{tool_name}' from toolkit '{tool.get('type', '')}' cannot start with '_'")
+        
+        tool['settings']['alita'] = alita
+        tool['settings']['llm'] = llm
+        tool_type = tool['type']
+        
+        # Check if tool is available and has get_tools function
+        if tool_type in AVAILABLE_TOOLS and 'get_tools' in AVAILABLE_TOOLS[tool_type]:
+            try:
+                get_tools_func = AVAILABLE_TOOLS[tool_type]['get_tools']
+                
+                # Handle special cases for ADO tools
+                if tool_type in ['ado_boards', 'ado_wiki', 'ado_plans']:
+                    tools.extend(get_tools_func(tool_type, tool))
+                else:
+                    tools.extend(get_tools_func(tool))
+                    
+            except Exception as e:
+                logger.error(f"Error getting tools for {tool_type}: {e}")
+                
+        # Handle ADO repos special case (it might be requested as azure_devops_repos)
+        elif tool_type in ['ado_repos', 'azure_devops_repos'] and 'ado_repos' in AVAILABLE_TOOLS:
+            try:
+                get_tools_func = AVAILABLE_TOOLS['ado_repos']['get_tools']
+                tools.extend(get_tools_func(tool))
+            except Exception as e:
+                logger.error(f"Error getting ADO repos tools: {e}")
+                
+        # Handle custom modules
+        elif tool.get("settings", {}).get("module"):
+            try:
+                settings = tool.get("settings", {})
+                mod = import_module(settings.pop("module"))
+                tkitclass = getattr(mod, settings.pop("class"))
+                toolkit = tkitclass.get_toolkit(**tool["settings"])
+                tools.extend(toolkit.get_tools())
+            except Exception as e:
+                logger.error(f"Error in getting custom toolkit: {e}")
+                
+        else:
+            # Tool not available or not found
+            if tool_type in FAILED_IMPORTS:
+                logger.warning(f"Tool '{tool_type}' is not available: {FAILED_IMPORTS[tool_type]}")
+            else:
+                logger.warning(f"Unknown tool type: {tool_type}")
+                
+    return tools
+
+def get_toolkits():
+    """Return toolkit configurations for all successfully imported toolkits."""
+    toolkit_configs = []
+    
+    for toolkit_name, toolkit_class in AVAILABLE_TOOLKITS.items():
+        try:
+            if hasattr(toolkit_class, 'toolkit_config_schema'):
+                toolkit_configs.append(toolkit_class.toolkit_config_schema())
+            else:
+                logger.debug(f"Toolkit {toolkit_name} does not have toolkit_config_schema method")
+        except Exception as e:
+            logger.error(f"Error getting config schema for {toolkit_name}: {e}")
+    
+    logger.info(f"Successfully loaded {len(toolkit_configs)} toolkit configurations")
+    return toolkit_configs
+
+def get_available_tools():
+    """Return list of available tool types."""
+    return list(AVAILABLE_TOOLS.keys())
+
+def get_failed_imports():
+    """Return dictionary of failed imports and their error messages."""
+    return FAILED_IMPORTS.copy()
+
+def get_available_toolkits():
+    """Return list of available toolkit class names."""
+    return list(AVAILABLE_TOOLKITS.keys())
+
+def diagnose_imports():
+    """Print diagnostic information about tool imports."""
+    available_count = len(AVAILABLE_TOOLS)
+    failed_count = len(FAILED_IMPORTS)
+    total_count = available_count + failed_count
+    
+    print(f"=== Tool Import Diagnostic ===")
+    print(f"Total tools: {total_count}")
+    print(f"Successfully imported: {available_count}")
+    print(f"Failed imports: {failed_count}")
+    print(f"Success rate: {(available_count/total_count*100):.1f}%")
+    
+    if AVAILABLE_TOOLS:
+        print(f"\n✅ Available tools ({len(AVAILABLE_TOOLS)}):")
+        for tool_name in sorted(AVAILABLE_TOOLS.keys()):
+            functions = []
+            if 'get_tools' in AVAILABLE_TOOLS[tool_name]:
+                functions.append('get_tools')
+            if 'toolkit_class' in AVAILABLE_TOOLS[tool_name]:
+                functions.append('toolkit')
+            print(f"  - {tool_name}: {', '.join(functions)}")
+    
+    if FAILED_IMPORTS:
+        print(f"\n❌ Failed imports ({len(FAILED_IMPORTS)}):")
+        for tool_name, error in FAILED_IMPORTS.items():
+            print(f"  - {tool_name}: {error}")
+    
+    if AVAILABLE_TOOLKITS:
+        print(f"\n🔧 Available toolkits ({len(AVAILABLE_TOOLKITS)}):")
+        for toolkit_name in sorted(AVAILABLE_TOOLKITS.keys()):
+            print(f"  - {toolkit_name}")
+
+# Export useful functions
+__all__ = [
+    'get_tools',
+    'get_toolkits',
+    'get_available_tools',
+    'get_failed_imports', 
+    'get_available_toolkits',
+    'diagnose_imports'
+]

alita_sdk-0.3.151/alita_sdk/tools/memory/__init__.py

@@ -0,0 +1,48 @@
+from typing import Optional, List
+
+from langchain_core.tools import BaseToolkit, BaseTool
+from langgraph.store.postgres import PostgresStore
+try:
+    from langmem import create_manage_memory_tool, create_search_memory_tool
+except ImportError:
+    # langmem is optional; define stubs to avoid import errors
+    def create_manage_memory_tool(*args, **kwargs):  # pragma: no cover
+        raise ImportError("langmem is required for MemoryToolkit")
+
+    def create_search_memory_tool(*args, **kwargs):  # pragma: no cover
+        raise ImportError("langmem is required for MemoryToolkit")
+from pydantic import create_model, BaseModel, ConfigDict, Field, SecretStr
+
+name = "memory"
+
+class MemoryToolkit(BaseToolkit):
+    tools: List[BaseTool] = []
+
+
+    @staticmethod
+    def toolkit_config_schema() -> BaseModel:
+        return create_model(
+            name,
+            namespace=(str, Field(description="Memory namespace", json_schema_extra={'toolkit_name': True})),
+            connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
+                                                          default=None,
+                                                          json_schema_extra={'secret': True})),
+            __config__=ConfigDict(json_schema_extra={
+                'metadata': {
+                    "label": "Memory",
+                    "icon_url": "jira-icon.svg",
+                    "categories": ["other"],
+                    "extra_categories": ["long-term memory", "langmem"],
+                }
+            })
+        )
+
+    @classmethod
+    def get_toolkit(cls, namespace: str, store: PostgresStore, **kwargs):
+        return cls(tools=[
+            create_manage_memory_tool(namespace=namespace, store=store),
+            create_search_memory_tool(namespace=namespace, store=store)
+        ])
+
+    def get_tools(self):
+        return self.tools