alita-sdk 0.3.206__py3-none-any.whl → 0.3.208__py3-none-any.whl

alita_sdk/runtime/utils/toolkit_runtime.py

@@ -0,0 +1,147 @@
+"""
+Toolkit runtime utilities for event dispatching and execution context.
+This module provides tools with the ability to dispatch custom events during execution.
+"""
+
+import sys
+import logging
+from typing import Dict, Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def dispatch_custom_event(event_type: str, data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """
+    Dispatch a custom event from within a toolkit tool execution.
+    
+    This function can be called by toolkit tools to send events back to the runtime
+    for monitoring, logging, or other purposes.
+    
+    Args:
+        event_type: Type of the event (e.g., "progress", "warning", "info")
+        data: Event data dictionary
+        
+    Returns:
+        Event dictionary if successful, None if no executor context available
+        
+    Example:
+        ```python
+        from alita_sdk.runtime.utils.toolkit_runtime import dispatch_custom_event
+        
+        def my_tool_function(param1, param2):
+            # Dispatch a progress event
+            dispatch_custom_event("progress", {
+                "message": "Processing started",
+                "step": 1,
+                "total_steps": 3
+            })
+            
+            # Do some work
+            result = process_data(param1, param2)
+            
+            # Dispatch completion event
+            dispatch_custom_event("completion", {
+                "message": "Processing completed",
+                "result_size": len(result)
+            })
+            
+            return result
+        ```
+    """
+    try:
+        # Try to get the current executor context
+        if hasattr(sys.modules[__name__], 'toolkit_dispatch_context'):
+            context = sys.modules[__name__].toolkit_dispatch_context
+            return context.dispatch_custom_event(event_type, data)
+        else:
+            # No executor context available - this is normal when not in test mode
+            logger.debug(f"No toolkit executor context available for event: {event_type}")
+            return None
+    except Exception as e:
+        logger.warning(f"Error dispatching custom event {event_type}: {e}")
+        return None
+
+
+def get_executor_context():
+    """
+    Get the current toolkit executor context if available.
+    
+    Returns:
+        ToolkitExecutor context or None if not in execution context
+    """
+    try:
+        if hasattr(sys.modules[__name__], 'toolkit_dispatch_context'):
+            return sys.modules[__name__].toolkit_dispatch_context.executor
+        return None
+    except Exception:
+        return None
+
+
+def is_in_test_mode() -> bool:
+    """
+    Check if the toolkit is currently running in test mode.
+    
+    Returns:
+        True if running in test mode with executor context, False otherwise
+    """
+    return get_executor_context() is not None
+
+
+class ToolkitRuntimeContext:
+    """
+    Context manager for toolkit runtime execution.
+    
+    This can be used by tools that need to perform setup/cleanup operations
+    when running in test mode vs normal execution.
+    """
+    
+    def __init__(self, tool_name: str):
+        self.tool_name = tool_name
+        self.executor = get_executor_context()
+        
+    def __enter__(self):
+        if self.executor:
+            dispatch_custom_event("tool_start", {
+                "tool_name": self.tool_name,
+                "message": f"Starting execution of {self.tool_name}"
+            })
+        return self
+        
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.executor:
+            if exc_type is None:
+                dispatch_custom_event("tool_end", {
+                    "tool_name": self.tool_name,
+                    "message": f"Completed execution of {self.tool_name}",
+                    "success": True
+                })
+            else:
+                dispatch_custom_event("tool_error", {
+                    "tool_name": self.tool_name,
+                    "message": f"Error in {self.tool_name}: {exc_val}",
+                    "success": False,
+                    "error_type": exc_type.__name__ if exc_type else "Unknown"
+                })
+        return False  # Don't suppress exceptions
+    
+    def dispatch_progress(self, message: str, step: int = None, total_steps: int = None, **kwargs):
+        """Convenience method for dispatching progress events."""
+        data = {"message": message, "tool_name": self.tool_name}
+        if step is not None:
+            data["step"] = step
+        if total_steps is not None:
+            data["total_steps"] = total_steps
+        data.update(kwargs)
+        dispatch_custom_event("progress", data)
+    
+    def dispatch_info(self, message: str, **kwargs):
+        """Convenience method for dispatching info events."""
+        data = {"message": message, "tool_name": self.tool_name}
+        data.update(kwargs)
+        dispatch_custom_event("info", data)
+    
+    def dispatch_warning(self, message: str, **kwargs):
+        """Convenience method for dispatching warning events."""
+        data = {"message": message, "tool_name": self.tool_name}
+        data.update(kwargs)
+        dispatch_custom_event("warning", data)

alita_sdk/runtime/utils/toolkit_utils.py

@@ -0,0 +1,157 @@
+"""
+Toolkit utilities for instantiating and managing toolkits.
+This module provides toolkit management functions that are not tied to any specific interface.
+"""
+
+import logging
+import random
+from typing import Dict, Any, Optional, List
+
+logger = logging.getLogger(__name__)
+
+
+def instantiate_toolkit_with_client(toolkit_config: Dict[str, Any], 
+                                   llm_client: Any, 
+                                   alita_client: Optional[Any] = None) -> List[Any]:
+    """
+    Instantiate a toolkit with LLM client support.
+    
+    This is a variant of instantiate_toolkit that includes LLM client support
+    for toolkits that require LLM capabilities.
+    
+    Args:
+        toolkit_config: Configuration dictionary for the toolkit
+        llm_client: LLM client instance for tools that need LLM capabilities
+        client: Optional additional client instance
+    
+    Returns:
+        List of instantiated tools from the toolkit
+        
+    Raises:
+        ValueError: If required configuration or client is missing
+        Exception: If toolkit instantiation fails
+    """
+    try:
+        from ..toolkits.tools import get_tools
+        
+        toolkit_name = toolkit_config.get('toolkit_name')
+        if not toolkit_name:
+            raise ValueError("toolkit_name is required in configuration")
+        
+        if not llm_client:
+            raise ValueError("LLM client is required but not provided")
+        
+        settings = toolkit_config.get('settings', {})
+        
+        # Log the configuration being used
+        logger.info(f"Instantiating toolkit {toolkit_name} with LLM client")
+        logger.debug(f"Toolkit {toolkit_name} configuration: {toolkit_config}")
+        
+        # Create a tool configuration dict with required fields
+        tool_config = {
+            'id': toolkit_config.get('id', random.randint(1, 1000000)),
+            'type': toolkit_config.get('type', toolkit_name.lower()),
+            'settings': settings,
+            'toolkit_name': toolkit_name
+        }
+        
+        # Get tools using the toolkit configuration with clients
+        # Parameter order: get_tools(tools_list, alita_client, llm, memory_store)
+        tools = get_tools([tool_config], alita_client, llm_client)
+        
+        if not tools:
+            logger.warning(f"No tools returned for toolkit {toolkit_name}")
+            return []
+        
+        logger.info(f"Successfully instantiated toolkit {toolkit_name} with {len(tools)} tools")
+        return tools
+            
+    except Exception as e:
+        logger.error(f"Error instantiating toolkit {toolkit_name} with client: {str(e)}")
+        raise
+
+
+def get_toolkit_tools(toolkit_instance: Any) -> List[Any]:
+    """
+    Extract tools from an instantiated toolkit instance.
+    
+    This function provides a standardized way to get tools from various
+    toolkit implementations that might have different interfaces.
+    
+    Args:
+        toolkit_instance: An instantiated toolkit object
+        
+    Returns:
+        List of tools from the toolkit
+        
+    Raises:
+        ValueError: If no tools can be extracted from the toolkit
+    """
+    try:
+        # Try different methods to get tools from the toolkit
+        if hasattr(toolkit_instance, 'get_tools'):
+            tools = toolkit_instance.get_tools()
+        elif hasattr(toolkit_instance, 'tools'):
+            tools = toolkit_instance.tools
+        elif hasattr(toolkit_instance, '_tools'):
+            tools = toolkit_instance._tools
+        else:
+            raise ValueError("Could not find tools in the toolkit instance")
+        
+        if not tools:
+            logger.warning("Toolkit instance returned empty tools list")
+            return []
+        
+        logger.info(f"Extracted {len(tools)} tools from toolkit instance")
+        return tools
+        
+    except Exception as e:
+        logger.error(f"Error extracting tools from toolkit: {str(e)}")
+        raise
+
+
+def find_tool_by_name(tools: List[Any], tool_name: str) -> Optional[Any]:
+    """
+    Find a specific tool by name from a list of tools.
+    
+    Args:
+        tools: List of tool instances
+        tool_name: Name of the tool to find
+        
+    Returns:
+        The tool instance if found, None otherwise
+    """
+    for tool in tools:
+        # Check various attributes that might contain the tool name
+        if hasattr(tool, 'name') and tool.name == tool_name:
+            return tool
+        elif hasattr(tool, 'func') and hasattr(tool.func, '__name__') and tool.func.__name__ == tool_name:
+            return tool
+        elif hasattr(tool, '__name__') and tool.__name__ == tool_name:
+            return tool
+    
+    return None
+
+
+def get_tool_names(tools: List[Any]) -> List[str]:
+    """
+    Extract tool names from a list of tools.
+    
+    Args:
+        tools: List of tool instances
+        
+    Returns:
+        List of tool names
+    """
+    tool_names = []
+    for tool in tools:
+        if hasattr(tool, 'name'):
+            tool_names.append(tool.name)
+        elif hasattr(tool, 'func') and hasattr(tool.func, '__name__'):
+            tool_names.append(tool.func.__name__)
+        elif hasattr(tool, '__name__'):
+            tool_names.append(tool.__name__)
+        else:
+            tool_names.append(str(tool))
+    
+    return tool_names

alita_sdk/tools/ado/test_plan/test_plan_wrapper.py

@@ -1,7 +1,10 @@
 import json
 import logging
-from typing import Optional
+from typing import Any, Dict, Generator, List, Optional
 
+from langchain_core.documents import Document
+
+from alita_sdk.tools.elitea_base import BaseIndexParams
 from azure.devops.connection import Connection
 from azure.devops.v7_0.test_plan.models import TestPlanCreateParams, TestSuiteCreateParams, \
     SuiteTestCaseCreateUpdateParameters
@@ -13,7 +16,11 @@ from pydantic.fields import FieldInfo as Field
 import xml.etree.ElementTree as ET
 
 from ..work_item import AzureDevOpsApiWrapper
-from ...elitea_base import BaseToolApiWrapper
+from ...elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
+try:
+    from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
+except ImportError:
+    from alita_sdk.langchain.interfaces.llm_processor import get_embeddings
 
 logger = logging.getLogger(__name__)
 
@@ -96,7 +103,6 @@ TestCaseAddModel = create_model(
     suite_id=(int, Field(description="ID of the test suite to which test cases are to be added"))
 )
 
-
 test_steps_description = """Json or XML array string with test steps. 
     Json example: [{"stepNumber": 1, "action": "Some action", "expectedResult": "Some expectation"},...]
     XML example: 
@@ -158,7 +164,19 @@ TestCasesGetModel = create_model(
     suite_id=(int, Field(description="ID of the test suite for which test cases are requested"))
 )
 
-class TestPlanApiWrapper(BaseToolApiWrapper):
+# Schema for indexing ADO Wiki pages into vector store
+indexData = create_model(
+    "indexData",
+    __base__=BaseIndexParams,
+    plan_id=(int, Field(description="ID of the test plan for which test cases are requested")),
+    suite_ids=(list[int], Field(description="List of test suite IDs for which test cases are requested (can be empty)")),
+    progress_step=(Optional[int], Field(default=None, ge=0, le=100,
+                         description="Optional step size for progress reporting during indexing")),
+    clean_index=(Optional[bool], Field(default=False,
+                       description="Optional flag to enforce clean existing index before indexing new data")),
+)
+
+class TestPlanApiWrapper(BaseVectorStoreToolApiWrapper):
     __test__ = False
     organization_url: str
     project: str
@@ -166,6 +184,13 @@ class TestPlanApiWrapper(BaseToolApiWrapper):
     limit: Optional[int] = 5
     _client: Optional[TestPlanClient] = PrivateAttr()
 
+    llm: Any = None
+    connection_string: Optional[SecretStr] = None
+    collection_name: Optional[str] = None
+    embedding_model: Optional[str] = "HuggingFaceEmbeddings"
+    embedding_model_params: Optional[Dict[str, Any]] = {"model_name": "sentence-transformers/all-MiniLM-L6-v2"}
+    vectorstore_type: Optional[str] = "PGVector"
+
     class Config:
         arbitrary_types_allowed = True
 
@@ -250,8 +275,10 @@ class TestPlanApiWrapper(BaseToolApiWrapper):
         try:
             if isinstance(suite_test_case_create_update_parameters, str):
                 suite_test_case_create_update_parameters = json.loads(suite_test_case_create_update_parameters)
-            suite_test_case_create_update_params_obj = [SuiteTestCaseCreateUpdateParameters(**param) for param in suite_test_case_create_update_parameters]
-            test_cases = self._client.add_test_cases_to_suite(suite_test_case_create_update_params_obj, self.project, plan_id, suite_id)
+            suite_test_case_create_update_params_obj = [SuiteTestCaseCreateUpdateParameters(**param) for param in
+                                                        suite_test_case_create_update_parameters]
+            test_cases = self._client.add_test_cases_to_suite(suite_test_case_create_update_params_obj, self.project,
+                                                              plan_id, suite_id)
             return [test_case.as_dict() for test_case in test_cases]
         except Exception as e:
             logger.error(f"Error adding test case: {e}")
@@ -268,10 +295,11 @@ class TestPlanApiWrapper(BaseToolApiWrapper):
             test_steps=test_case['test_steps'],
             test_steps_format=test_case['test_steps_format']) for test_case in test_cases]
 
-
-    def create_test_case(self, plan_id: int, suite_id: int, title: str, description: str, test_steps: str, test_steps_format: str = 'json'):
+    def create_test_case(self, plan_id: int, suite_id: int, title: str, description: str, test_steps: str,
+                         test_steps_format: str = 'json'):
         """Creates a new test case in specified suite in Azure DevOps."""
-        work_item_wrapper = AzureDevOpsApiWrapper(organization_url=self.organization_url, token=self.token.get_secret_value(), project=self.project)
+        work_item_wrapper = AzureDevOpsApiWrapper(organization_url=self.organization_url,
+                                                  token=self.token.get_secret_value(), project=self.project)
         if test_steps_format == 'json':
             steps_xml = self.get_test_steps_xml(json.loads(test_steps))
         elif test_steps_format == 'xml':
@@ -279,8 +307,9 @@ class TestPlanApiWrapper(BaseToolApiWrapper):
         else:
             return ToolException("Unknown test steps format: " + test_steps_format)
         work_item_json = self.build_ado_test_case(title, description, steps_xml)
-        created_work_item_id = work_item_wrapper.create_work_item(work_item_json=json.dumps(work_item_json), wi_type="Test Case")['id']
-        return self.add_test_case([{"work_item":{"id":created_work_item_id}}], plan_id, suite_id)
+        created_work_item_id = \
+        work_item_wrapper.create_work_item(work_item_json=json.dumps(work_item_json), wi_type="Test Case")['id']
+        return self.add_test_case([{"work_item": {"id": created_work_item_id}}], plan_id, suite_id)
 
     def build_ado_test_case(self, title, description, steps_xml):
         """
@@ -355,6 +384,42 @@ class TestPlanApiWrapper(BaseToolApiWrapper):
             logger.error(f"Error getting test cases: {e}")
             return ToolException(f"Error getting test cases: {e}")
 
+    def index_data(self,
+            plan_id: str,
+            suite_ids: list[str] = [],
+            collection_suffix: str = '',
+            progress_step: int = None,
+            clean_index: bool = False
+    ):
+        """Load ADO TestCases into the vector store."""
+        docs = self._base_loader(plan_id, suite_ids)
+        embedding = get_embeddings(self.embedding_model, self.embedding_model_params)
+        vs = self._init_vector_store(collection_suffix, embeddings=embedding)
+        return vs.index_documents(docs, progress_step=progress_step, clean_index=clean_index)
+
+    def _base_loader(self, plan_id: str, suite_ids: Optional[list[str]] = []) -> Generator[Document, None, None]:
+        cases = []
+        for sid in suite_ids:
+            cases.extend(self.get_test_cases(plan_id, sid))
+        #
+        for case in cases:
+            field_dicts = case.get('work_item', {}).get('work_item_fields', [])
+            data = {k: v for d in field_dicts for k, v in d.items()}
+            yield Document(
+                page_content=data.get('Microsoft.VSTS.TCM.Steps', ''),
+                metadata={
+                    'id': case.get('work_item', {}).get('id', ''),
+                    'title': case.get('work_item', {}).get('name', ''),
+                    'plan_id': case.get('test_plan', {}).get('id', ''),
+                    'suite_id': case.get('test_suite', {}).get('id', ''),
+                    'description': data.get('System.Description', ''),
+                    'updated_on': data.get('System.Rev', ''),
+                })
+
+    def _process_document(self, document: Document) -> Generator[Document, None, None]:
+        yield document
+
+    @extend_with_vector_tools
     def get_available_tools(self):
         """Return a list of available tools."""
         return [
@@ -423,5 +488,11 @@ class TestPlanApiWrapper(BaseToolApiWrapper):
                 "description": self.get_test_cases.__doc__,
                 "args_schema": TestCasesGetModel,
                 "ref": self.get_test_cases,
+            },
+            {
+                "name": "index_data",
+                "ref": self.index_data,
+                "description": self.index_data.__doc__,
+                "args_schema": indexData,
             }
         ]

alita_sdk/tools/ado/wiki/ado_wrapper.py

@@ -1,6 +1,10 @@
 import logging
-from typing import Optional
+from typing import Any, Dict, Generator, List, Optional
 
+from alita_sdk.tools.elitea_base import BaseIndexParams
+from langchain_core.documents import Document
+
+from ...elitea_base import BaseVectorStoreToolApiWrapper, extend_with_vector_tools
 from azure.devops.connection import Connection
 from azure.devops.exceptions import AzureDevOpsServiceError
 from azure.devops.v7_0.core import CoreClient
@@ -12,6 +16,10 @@ from msrest.authentication import BasicAuthentication
 from pydantic import create_model, PrivateAttr, SecretStr
 from pydantic import model_validator
 from pydantic.fields import Field
+try:
+    from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
+except ImportError:
+    from alita_sdk.langchain.interfaces.llm_processor import get_embeddings
 
 from ...elitea_base import BaseToolApiWrapper
 
@@ -52,14 +60,32 @@ RenamePageInput = create_model(
     version_type=(Optional[str], Field(description="Version type (branch, tag, or commit). Determines how Id is interpreted", default="branch"))
 )
 
+# Schema for indexing ADO Wiki pages into vector store
+indexData = create_model(
+    "indexData",
+    __base__=BaseIndexParams,
+    wiki_identifier=(str, Field(description="Wiki identifier to index, e.g., 'ABCProject.wiki'")),
+    progress_step=(Optional[int], Field(default=None, ge=0, le=100,
+                         description="Optional step size for progress reporting during indexing")),
+    clean_index=(Optional[bool], Field(default=False,
+                       description="Optional flag to enforce clean existing index before indexing new data")),
+)
+
 
-class AzureDevOpsApiWrapper(BaseToolApiWrapper):
+class AzureDevOpsApiWrapper(BaseVectorStoreToolApiWrapper):
     organization_url: str
     project: str
     token: SecretStr
     _client: Optional[WikiClient] = PrivateAttr()  # Private attribute for the wiki client
     _core_client: Optional[CoreClient] = PrivateAttr()  # Private attribute for the CoreClient client
 
+    llm: Any = None
+    connection_string: Optional[SecretStr] = None
+    collection_name: Optional[str] = None
+    embedding_model: Optional[str] = "HuggingFaceEmbeddings"
+    embedding_model_params: Optional[Dict[str, Any]] = {"model_name": "sentence-transformers/all-MiniLM-L6-v2"}
+    vectorstore_type: Optional[str] = "PGVector"
+
     class Config:
         arbitrary_types_allowed = True  # Allow arbitrary types (e.g., WorkItemTrackingClient)
 
@@ -216,6 +242,34 @@ class AzureDevOpsApiWrapper(BaseToolApiWrapper):
             logger.error(f"Unable to modify wiki page: {str(e)}")
             return ToolException(f"Unable to modify wiki page: {str(e)}")
 
+    def index_data(
+            self,
+            wiki_identifier: str,
+            collection_suffix: str = '',
+            progress_step: int = None,
+            clean_index: bool = False
+    ):
+        """Load ADO Wiki pages into the vector store."""
+        docs = self._base_loader(wiki_identifier)
+        embedding = get_embeddings(self.embedding_model, self.embedding_model_params)
+        vs = self._init_vector_store(collection_suffix, embeddings=embedding)
+        return vs.index_documents(docs, progress_step=progress_step, clean_index=clean_index)
+
+    def _base_loader(self, wiki_identifier: str) -> Generator[Document, None, None]:
+        pages = self._client.get_pages_batch(pages_batch_request={}, project=self.project, wiki_identifier=wiki_identifier)
+        #
+        for page in pages:
+            content = self._client.get_page_by_id(project=self.project, wiki_identifier=wiki_identifier, id=page.id, include_content=True).page.content
+            yield Document(page_content=content, metadata={
+                'id': page.id, 
+                'path': page.path,
+                'updated_on': ''
+            })
+
+    def _process_document(self, document: Document) -> Generator[Document, None, None]:
+        yield document
+
+    @extend_with_vector_tools
     def get_available_tools(self):
         """Return a list of available tools."""
         return [
@@ -260,5 +314,11 @@ class AzureDevOpsApiWrapper(BaseToolApiWrapper):
                 "description": self.rename_wiki_page.__doc__,
                 "args_schema": RenamePageInput,
                 "ref": self.rename_wiki_page,
+            },
+            {
+                "name": "index_data",
+                "ref": self.index_data,
+                "description": self.index_data.__doc__,
+                "args_schema": indexData,
             }
         ]

alita_sdk/tools/chunkers/sematic/markdown_chunker.py

@@ -1,12 +1,13 @@
 from typing import Generator
 from langchain.schema import Document
+from langchain_core.documents import Document
 from langchain_text_splitters import MarkdownHeaderTextSplitter
 from langchain.text_splitter import TokenTextSplitter
 from ..utils import tiktoken_length
 from copy import deepcopy as copy
 
 
-def markdown_chunker(file_content_generator: Generator[Document, None, None], config: dict, *args, **kwargs) -> Generator[str, None, None]:
+def markdown_chunker(file_content_generator: Generator[Document, None, None], config: dict, *args, **kwargs) -> Generator[Document, None, None]:
     strip_header = config.get("strip_header", False)
     return_each_line = config.get("return_each_line", False)
     headers_to_split_on = config.get("headers_to_split_on", [])

alita_sdk/tools/memory/__init__.py

@@ -1,7 +1,7 @@
 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:
@@ -15,13 +15,37 @@ from pydantic import create_model, BaseModel, ConfigDict, Field, SecretStr
 
 name = "memory"
 
-def get_tools(tool):
-    return MemoryToolkit().get_toolkit(
-        namespace=tool['settings'].get('namespace', str(tool['id'])),
-        username=tool['settings'].get('username', ''),
-        store=tool['settings'].get('store', None),
-        toolkit_name=tool.get('toolkit_name', '')
-).get_tools()
+def get_tools(tools_list: list, alita_client, llm, memory_store=None):
+    """
+    Get memory tools for the provided tool configurations.
+    
+    Args:
+        tools_list: List of tool configurations
+        alita_client: Alita client instance
+        llm: LLM client instance
+        memory_store: Optional memory store instance
+    
+    Returns:
+        List of memory tools
+    """
+    all_tools = []
+    
+    for tool in tools_list:
+        if tool.get('type') == 'memory' or tool.get('toolkit_name') == 'memory':
+            try:
+                toolkit_instance = MemoryToolkit().get_toolkit(
+                    namespace=tool['settings'].get('namespace', str(tool['id'])),
+                    username=tool['settings'].get('username', ''),
+                    store=tool['settings'].get('store', memory_store),
+                    toolkit_name=tool.get('toolkit_name', '')
+                )
+                all_tools.extend(toolkit_instance.get_tools())
+            except Exception as e:
+                print(f"DEBUG: Error in memory toolkit get_tools: {e}")
+                print(f"DEBUG: Tool config: {tool}")
+                raise
+    
+    return all_tools
 
 class MemoryToolkit(BaseToolkit):
     tools: List[BaseTool] = []
@@ -30,7 +54,7 @@ class MemoryToolkit(BaseToolkit):
     @staticmethod
     def toolkit_config_schema() -> BaseModel:
         return create_model(
-            name,
+            'MemoryConfig',
             namespace=(str, Field(description="Memory namespace", json_schema_extra={'toolkit_name': True})),
             username=(Optional[str], Field(description="Username", default='Tester', json_schema_extra={'hidden': True})),
             connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
@@ -48,7 +72,27 @@ class MemoryToolkit(BaseToolkit):
         )
 
     @classmethod
-    def get_toolkit(cls, namespace: str, store: PostgresStore, **kwargs):
+    def get_toolkit(cls, namespace: str, store=None, **kwargs):
+        """
+        Get toolkit with memory tools.
+        
+        Args:
+            namespace: Memory namespace
+            store: PostgresStore instance (imported dynamically)
+            **kwargs: Additional arguments
+        """
+        try:
+            from langgraph.store.postgres import PostgresStore
+        except ImportError:
+            raise ImportError(
+                "PostgreSQL dependencies (psycopg) are required for MemoryToolkit. "
+                "Install with: pip install psycopg[binary]"
+            )
+        
+        # Validate store type
+        if store is not None and not isinstance(store, PostgresStore):
+            raise TypeError(f"Expected PostgresStore, got {type(store)}")
+        
         return cls(tools=[
             create_manage_memory_tool(namespace=namespace, store=store),
             create_search_memory_tool(namespace=namespace, store=store)