alita-sdk 0.3.497__py3-none-any.whl → 0.3.515__py3-none-any.whl

alita_sdk/configurations/testio.py

@@ -1,3 +1,4 @@
+import requests
 from pydantic import BaseModel, ConfigDict, Field, SecretStr
 
 
@@ -16,3 +17,85 @@ class TestIOConfiguration(BaseModel):
     )
     endpoint: str = Field(description="TestIO endpoint")
     api_key: SecretStr = Field(description="TestIO API Key")
+
+    @staticmethod
+    def check_connection(settings: dict) -> str | None:
+        """
+        Test the connection to TestIO API.
+
+        Args:
+            settings: Dictionary containing 'endpoint' and 'api_key' (both required)
+
+        Returns:
+            None if connection is successful, error message string otherwise
+        """
+        endpoint = settings.get("endpoint")
+        if endpoint is None or endpoint == "":
+            if endpoint == "":
+                return "Endpoint cannot be empty"
+            return "Endpoint is required"
+
+        # Validate endpoint format
+        if not isinstance(endpoint, str):
+            return "Endpoint must be a string"
+        
+        endpoint = endpoint.strip()
+        if not endpoint:
+            return "Endpoint cannot be empty"
+        if not endpoint.startswith(("http://", "https://")):
+            return "Endpoint must start with http:// or https://"
+
+        # Remove trailing slash for consistency
+        endpoint = endpoint.rstrip("/")
+
+        api_key = settings.get("api_key")
+        if api_key is None:
+            return "API key is required"
+
+        # Extract secret value if it's a SecretStr
+        if hasattr(api_key, "get_secret_value"):
+            api_key = api_key.get_secret_value()
+
+        # Validate API key is not empty
+        if not api_key or not api_key.strip():
+            return "API key cannot be empty"
+
+        # Verification strategy:
+        # Use an auth-required endpoint and a single, explicit auth scheme:
+        #   Authorization: Token <token>
+        url = f"{endpoint}/customer/v2/products"
+
+        try:
+            resp = TestIOConfiguration._get_with_token(url, api_key)
+
+            if resp.status_code == 200:
+                return None  # Connection successful
+            if resp.status_code == 401:
+                return "Invalid token"
+            if resp.status_code == 403:
+                return "Access forbidden - token has no access to /customer/v2/products"
+            if resp.status_code == 404:
+                return "Invalid endpoint. Verify TestIO base endpoint."
+            if resp.status_code == 429:
+                return "Rate limited - please try again later"
+            if 500 <= resp.status_code <= 599:
+                return f"TestIO service error (HTTP {resp.status_code})"
+            return f"Connection failed (HTTP {resp.status_code})"
+
+        except requests.exceptions.Timeout:
+            return "Connection timeout - TestIO did not respond within 10 seconds"
+        except requests.exceptions.ConnectionError:
+            return "Connection error - unable to reach TestIO. Check endpoint URL and network."
+        except requests.exceptions.RequestException as e:
+            return f"Request failed: {str(e)}"
+        except Exception:
+            return "Unexpected error during TestIO connection check"
+
+    @staticmethod
+    def _get_with_token(url: str, token: str) -> requests.Response:
+        """Perform an authenticated GET using `Authorization: Token <token>`."""
+        return requests.get(
+            url,
+            headers={"Authorization": f"Token {token}"},
+            timeout=10,
+        )

alita_sdk/runtime/clients/artifact.py

@@ -38,8 +38,8 @@ class Artifact:
         if len(data) == 0:
             # empty file might be created
             return ""
-        if isinstance(data, dict) and data['error']:
-            return f"{data['error']}. {data['content'] if data['content'] else ''}"
+        if isinstance(data, dict) and data.get('error'):
+            return f"{data['error']}. {data.get('content', '')}"
         detected = chardet.detect(data)
         if detected['encoding'] is not None:
             try:

alita_sdk/runtime/clients/client.py

@@ -20,9 +20,8 @@ from .prompt import AlitaPrompt
 from .datasource import AlitaDataSource
 from .artifact import Artifact
 from ..langchain.chat_message_template import Jinja2TemplatedChatMessagesTemplate
-from ..utils.utils import TOOLKIT_SPLITTER
 from ..utils.mcp_oauth import McpAuthorizationRequired
-from ...tools import get_available_toolkit_models
+from ...tools import get_available_toolkit_models, instantiate_toolkit
 from ...tools.base_indexer_toolkit import IndexTools
 
 logger = logging.getLogger(__name__)
@@ -146,6 +145,19 @@ class AlitaClient:
         data = requests.get(url, headers=self.headers, verify=False).json()
         return data
 
+    def toolkit(self, toolkit_id: int):
+        url = f"{self.base_url}{self.api_path}/tool/prompt_lib/{self.project_id}/{toolkit_id}"
+        response = requests.get(url, headers=self.headers, verify=False)
+        if not response.ok:
+            raise ValueError(f"Failed to fetch toolkit {toolkit_id}: {response.text}")
+        
+        tool_data = response.json()
+        if 'settings' not in tool_data:
+            tool_data['settings'] = {}
+        tool_data['settings']['alita'] = self
+        
+        return instantiate_toolkit(tool_data)
+
     def get_list_of_apps(self):
         apps = []
         limit = 10
@@ -348,7 +360,7 @@ class AlitaClient:
                     application_variables: Optional[dict] = None,
                     version_details: Optional[dict] = None, store: Optional[BaseStore] = None,
                     llm: Optional[ChatOpenAI] = None, mcp_tokens: Optional[dict] = None,
-                    conversation_id: Optional[str] = None):
+                    conversation_id: Optional[str] = None, ignored_mcp_servers: Optional[list] = None):
         if tools is None:
             tools = []
         if chat_history is None:
@@ -397,12 +409,12 @@ class AlitaClient:
         if runtime == 'nonrunnable':
             return LangChainAssistant(self, data, llm, chat_history, app_type,
                                       tools=tools, memory=memory, store=store, mcp_tokens=mcp_tokens,
-                                      conversation_id=conversation_id)
+                                      conversation_id=conversation_id, ignored_mcp_servers=ignored_mcp_servers)
         if runtime == 'langchain':
             return LangChainAssistant(self, data, llm,
                                       chat_history, app_type,
                                       tools=tools, memory=memory, store=store, mcp_tokens=mcp_tokens,
-                                      conversation_id=conversation_id).runnable()
+                                      conversation_id=conversation_id, ignored_mcp_servers=ignored_mcp_servers).runnable()
         elif runtime == 'llama':
             raise NotImplementedError("LLama runtime is not supported")
 
@@ -656,7 +668,8 @@ class AlitaClient:
                       tools: Optional[list] = None, chat_history: Optional[List[Any]] = None,
                       memory=None, runtime='langchain', variables: Optional[list] = None,
                       store: Optional[BaseStore] = None, debug_mode: Optional[bool] = False,
-                      mcp_tokens: Optional[dict] = None, conversation_id: Optional[str] = None):
+                      mcp_tokens: Optional[dict] = None, conversation_id: Optional[str] = None,
+                      ignored_mcp_servers: Optional[list] = None):
         """
         Create a predict-type agent with minimal configuration.
 
@@ -672,6 +685,7 @@ class AlitaClient:
             variables: Optional list of variables for the agent
             store: Optional store for memory
             debug_mode: Enable debug mode for cases when assistant can be initialized without tools
+            ignored_mcp_servers: Optional list of MCP server URLs to ignore (user chose to continue without auth)
 
         Returns:
             Runnable agent ready for execution
@@ -705,7 +719,8 @@ class AlitaClient:
             store=store,
             debug_mode=debug_mode,
             mcp_tokens=mcp_tokens,
-            conversation_id=conversation_id
+            conversation_id=conversation_id,
+            ignored_mcp_servers=ignored_mcp_servers
         ).runnable()
 
     def test_toolkit_tool(self, toolkit_config: dict, tool_name: str, tool_params: dict = None,
@@ -939,7 +954,6 @@ class AlitaClient:
             if target_tool is None:
                 available_tools = []
                 base_available_tools = []
-                full_available_tools = []
 
                 for tool in tools:
                     tool_name_attr = None
@@ -956,10 +970,6 @@ class AlitaClient:
                         if base_name not in base_available_tools:
                             base_available_tools.append(base_name)
 
-                        # Track full names separately
-                        if TOOLKIT_SPLITTER in tool_name_attr:
-                            full_available_tools.append(tool_name_attr)
-
                 # Create comprehensive error message
                 error_msg = f"Tool '{tool_name}' not found in toolkit '{toolkit_config.get('toolkit_name')}'.\n"
 
@@ -967,9 +977,7 @@ class AlitaClient:
                 if toolkit_name in [tool.value for tool in IndexTools]:
                     error_msg += f" Please make sure proper PGVector configuration and embedding model are set in the platform.\n"
 
-                if base_available_tools and full_available_tools:
-                    error_msg += f" Available tools: {base_available_tools} (base names) or {full_available_tools} (full names)"
-                elif base_available_tools:
+                if base_available_tools:
                     error_msg += f" Available tools: {base_available_tools}"
                 elif available_tools:
                     error_msg += f" Available tools: {available_tools}"
@@ -978,10 +986,7 @@ class AlitaClient:
 
                 # Add helpful hint about naming conventions
                 if '___' in tool_name:
-                    error_msg += f" Note: You provided a full name '{tool_name}'. Try using just the base name '{extract_base_tool_name(tool_name)}'."
-                elif full_available_tools:
-                    possible_full_name = create_full_tool_name(tool_name, toolkit_name)
-                    error_msg += f" Note: You provided a base name '{tool_name}'. The full name might be '{possible_full_name}'."
+                    error_msg += f" Note: Tool names no longer use '___' prefixes. Try using just the base name '{extract_base_tool_name(tool_name)}'."
 
                 return {
                     "success": False,

alita_sdk/runtime/clients/sandbox_client.py

@@ -6,6 +6,7 @@ import requests
 from typing import Any
 from json import dumps
 import chardet
+from ...tools import instantiate_toolkit
 
 logger = logging.getLogger(__name__)
 
@@ -184,6 +185,19 @@ class SandboxClient:
         data = requests.get(url, headers=self.headers, verify=False).json()
         return data
 
+    def toolkit(self, toolkit_id: int):
+        url = f"{self.base_url}{self.api_path}/tool/prompt_lib/{self.project_id}/{toolkit_id}"
+        response = requests.get(url, headers=self.headers, verify=False)
+        if not response.ok:
+            raise ValueError(f"Failed to fetch toolkit {toolkit_id}: {response.text}")
+        
+        tool_data = response.json()
+        if 'settings' not in tool_data:
+            tool_data['settings'] = {}
+        tool_data['settings']['alita'] = self
+        
+        return instantiate_toolkit(tool_data)
+
     def get_list_of_apps(self):
         apps = []
         limit = 10

alita_sdk/runtime/langchain/assistant.py

@@ -33,7 +33,8 @@ class Assistant:
                  store: Optional[BaseStore] = None,
                  debug_mode: Optional[bool] = False,
                  mcp_tokens: Optional[dict] = None,
-                 conversation_id: Optional[str] = None):
+                 conversation_id: Optional[str] = None,
+                 ignored_mcp_servers: Optional[list] = None):
 
         self.app_type = app_type
         self.memory = memory
@@ -98,10 +99,55 @@ class Assistant:
             memory_store=self.store,
             debug_mode=debug_mode,
             mcp_tokens=mcp_tokens,
-            conversation_id=conversation_id
+            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']

alita_sdk/runtime/langchain/document_loaders/AlitaJSONLinesLoader.py

@@ -0,0 +1,77 @@
+from .AlitaJSONLoader import AlitaJSONLoader
+import json
+from io import StringIO
+from typing import List, Iterator
+
+from langchain_core.documents import Document
+from langchain_core.tools import ToolException
+
+
+class AlitaJSONLinesLoader(AlitaJSONLoader):
+    """Load local JSONL files (one JSON object per line) using AlitaJSONLoader behavior.
+
+    Behavior:
+    - Supports both `file_path` and `file_content` (bytes or file-like object), same as AlitaJSONLoader.
+    - Treats each non-empty line as an independent JSON object.
+    - Aggregates all parsed JSON objects into a list and feeds them through the same
+      RecursiveJsonSplitter-based chunking used by AlitaJSONLoader.lazy_load.
+    - Returns a list of Documents with chunked JSON content.
+    """
+
+    def __init__(self, **kwargs):
+        # Reuse AlitaJSONLoader initialization logic (file_path / file_content handling, encoding, etc.)
+        super().__init__(**kwargs)
+
+    def _iter_lines(self) -> Iterator[str]:
+        """Yield lines from file_path or file_content, mirroring AlitaJSONLoader sources."""
+        # Prefer file_path if available
+        if hasattr(self, "file_path") and self.file_path:
+            with open(self.file_path, "r", encoding=self.encoding) as f:
+                for line in f:
+                    yield line
+        # Fallback to file_content if available
+        elif hasattr(self, "file_content") and self.file_content:
+            # file_content may be bytes or a file-like object
+            if isinstance(self.file_content, (bytes, bytearray)):
+                text = self.file_content.decode(self.encoding)
+                for line in StringIO(text):
+                    yield line
+            else:
+                # Assume it's a text file-like object positioned at the beginning
+                self.file_content.seek(0)
+                for line in self.file_content:
+                    yield line
+        else:
+            raise ToolException("'file_path' or 'file_content' parameter should be provided.")
+
+    def load(self) -> List[Document]:  # type: ignore[override]
+        """Load JSONL content by delegating each non-empty line to AlitaJSONLoader.
+
+        For each non-empty line in the underlying source (file_path or file_content):
+        - Create a temporary AlitaJSONLoader instance with that line as file_content.
+        - Call lazy_load() on that instance to apply the same RecursiveJsonSplitter logic
+          as for a normal JSON file.
+        - Accumulate all Documents from all lines and return them as a single list.
+        """
+        docs: List[Document] = []
+
+        for raw_line in self._iter_lines():
+            line = raw_line.strip()
+            if not line:
+                continue
+            try:
+                # Instantiate a per-line AlitaJSONLoader using the same configuration
+                line_loader = AlitaJSONLoader(
+                    file_content=line,
+                    file_name=getattr(self, "file_name", str(getattr(self, "file_path", "no_name"))),
+                    encoding=self.encoding,
+                    autodetect_encoding=self.autodetect_encoding,
+                    max_tokens=self.max_tokens,
+                )
+
+                for doc in line_loader.lazy_load():
+                    docs.append(doc)
+            except Exception as e:
+                raise ToolException(f"Error processing JSONL line: {line[:100]}... Error: {e}") from e
+
+        return docs

alita_sdk/runtime/langchain/document_loaders/AlitaJSONLoader.py

@@ -32,6 +32,8 @@ class AlitaJSONLoader(BaseLoader):
             elif hasattr(self, 'file_content') and self.file_content:
                 if isinstance(self.file_content, bytes):
                     return json.loads(self.file_content.decode(self.encoding))
+                elif isinstance(self.file_content, str):
+                    return json.loads(self.file_content)
                 else:
                     return json.load(self.file_content)
             else:
@@ -45,7 +47,6 @@ class AlitaJSONLoader(BaseLoader):
                         try:
                             with open(self.file_path, encoding=encoding.encoding) as f:
                                 return f.read()
-                            break
                         except UnicodeDecodeError:
                             continue
                 elif hasattr(self, 'file_content') and self.file_content:

alita_sdk/runtime/langchain/document_loaders/constants.py

@@ -21,6 +21,7 @@ from .AlitaDocxMammothLoader import AlitaDocxMammothLoader
 from .AlitaExcelLoader import AlitaExcelLoader
 from .AlitaImageLoader import AlitaImageLoader
 from .AlitaJSONLoader import AlitaJSONLoader
+from .AlitaJSONLinesLoader import AlitaJSONLinesLoader
 from .AlitaPDFLoader import AlitaPDFLoader
 from .AlitaPowerPointLoader import AlitaPowerPointLoader
 from .AlitaTextLoader import AlitaTextLoader
@@ -208,7 +209,7 @@ document_loaders_map = {
         'allowed_to_override': DEFAULT_ALLOWED_BASE
     },
     '.jsonl': {
-        'class': AirbyteJSONLoader,
+        'class': AlitaJSONLinesLoader,
         'mime_type': 'application/jsonl',
         'is_multimodal_processing': False,
         'kwargs': {},

alita_sdk/runtime/langchain/langraph_agent.py

@@ -30,7 +30,7 @@ from ..tools.loop import LoopNode
 from ..tools.loop_output import LoopToolNode
 from ..tools.tool import ToolNode
 from ..utils.evaluate import EvaluateTemplate
-from ..utils.utils import clean_string, TOOLKIT_SPLITTER
+from ..utils.utils import clean_string
 from ..tools.router import RouterNode
 
 logger = logging.getLogger(__name__)
@@ -191,7 +191,7 @@ Answer only with step name, no need to add descrip in case none of the steps are
                     additional_info = """### Additoinal info: """
                 additional_info += "{field}: {value}\n".format(field=field, value=state.get(field, ""))
         decision_input.append(HumanMessage(
-            self.prompt.format(steps=self.steps, description=self.description, additional_info=additional_info)))
+            self.prompt.format(steps=self.steps, description=safe_format(self.description, state), additional_info=additional_info)))
         completion = self.client.invoke(decision_input)
         result = clean_string(completion.content.strip())
         logger.info(f"Plan to transition to: {result}")
@@ -483,8 +483,7 @@ def create_graph(
             node_id = clean_string(node['id'])
             toolkit_name = node.get('toolkit_name')
             tool_name = clean_string(node.get('tool', node_id))
-            if toolkit_name:
-                tool_name = f"{clean_string(toolkit_name)}{TOOLKIT_SPLITTER}{tool_name}"
+            # Tool names are now clean (no prefix needed)
             logger.info(f"Node: {node_id} : {node_type} - {tool_name}")
             if node_type in ['function', 'toolkit', 'mcp', 'tool', 'loop', 'loop_from_tool', 'indexer', 'subgraph', 'pipeline', 'agent']:
                 if node_type == 'mcp' and tool_name not in [tool.name for tool in tools]:
@@ -550,8 +549,8 @@ def create_graph(
                             loop_toolkit_name = node.get('loop_toolkit_name')
                             loop_tool_name = node.get('loop_tool')
                             if (loop_toolkit_name and loop_tool_name) or loop_tool_name:
-                                loop_tool_name = f"{clean_string(loop_toolkit_name)}{TOOLKIT_SPLITTER}{loop_tool_name}" if loop_toolkit_name else clean_string(
-                                    loop_tool_name)
+                                # Use clean tool name (no prefix)
+                                loop_tool_name = clean_string(loop_tool_name)
                                 for t in tools:
                                     if t.name == loop_tool_name:
                                         logger.debug(f"Loop tool discovered: {t}")
@@ -609,10 +608,10 @@ def create_graph(
                 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}{TOOLKIT_SPLITTER}{tool}")
+                        # Add tool names directly (no prefix)
+                        tool_names.extend(selected_tools)
                 elif isinstance(connected_tools, list):
-                    # for cases when tools are provided as a list of names with already bound toolkit_name
+                    # Use provided tool names as-is
                     tool_names = connected_tools
                 
                 if tool_names:

alita_sdk/runtime/langchain/utils.py

@@ -208,7 +208,12 @@ def safe_format(template, mapping):
 def create_pydantic_model(model_name: str, variables: dict[str, dict]):
     fields = {}
     for var_name, var_data in variables.items():
-        fields[var_name] = (parse_pydantic_type(var_data['type']), Field(description=var_data.get('description', None)))
+        if 'default' in var_data:
+            # allow user to define if it is required or not
+            fields[var_name] = (parse_pydantic_type(var_data['type']),
+                                Field(description=var_data.get('description', None), default=var_data.get('default')))
+        else:
+            fields[var_name] = (parse_pydantic_type(var_data['type']), Field(description=var_data.get('description', None)))
     return create_model(model_name, **fields)
 
 def parse_pydantic_type(type_name: str):

alita_sdk/runtime/toolkits/artifact.py

@@ -1,6 +1,7 @@
 from typing import List, Any, Literal, Optional
 
-from alita_sdk.tools.utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
+from alita_sdk.tools.utils import clean_string, get_max_toolkit_length
+from alita_sdk.tools.elitea_base import filter_missconfigured_index_tools
 from langchain_community.agent_toolkits.base import BaseToolkit
 from langchain_core.tools import BaseTool
 from pydantic import create_model, BaseModel, ConfigDict, Field
@@ -40,26 +41,34 @@ class ArtifactToolkit(BaseToolkit):
         )
     
     @classmethod
+    @filter_missconfigured_index_tools
     def get_toolkit(cls, client: Any, bucket: str, toolkit_name: Optional[str] = None, selected_tools: list[str] = [], **kwargs):
         if selected_tools is None:
             selected_tools = []
+        
         tools = []
         wrapper_payload = {
             **kwargs,
             **(kwargs.get('pgvector_configuration') or {}),
         }
         artifact_wrapper = ArtifactWrapper(alita=client, bucket=bucket, **wrapper_payload)
-        prefix = clean_string(toolkit_name, cls.toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
+        # Use clean toolkit name for context (max 1000 chars in description)
+        toolkit_context = f" [Toolkit: {clean_string(toolkit_name, 0)}]" if toolkit_name else ''
         available_tools = artifact_wrapper.get_available_tools()
         for tool in available_tools:
             if selected_tools:
                 if tool["name"] not in selected_tools:
                     continue
+            # Add toolkit context to description with character limit
+            description = tool["description"]
+            if toolkit_context and len(description + toolkit_context) <= 1000:
+                description = description + toolkit_context
             tools.append(BaseAction(
                 api_wrapper=artifact_wrapper,
-                name=prefix + tool["name"],
-                description=tool["description"],
-                args_schema=tool["args_schema"]
+                name=tool["name"],
+                description=description,
+                args_schema=tool["args_schema"],
+                metadata={"toolkit_name": toolkit_name} if toolkit_name else {}
             ))
         return cls(tools=tools)
     

alita_sdk/runtime/toolkits/datasource.py

@@ -3,7 +3,7 @@ from pydantic import create_model, BaseModel, Field
 from langchain_community.agent_toolkits.base import BaseToolkit
 from langchain_core.tools import BaseTool, ToolException
 from ..tools.datasource import DatasourcePredict, DatasourceSearch, datasourceToolSchema
-from alita_sdk.tools.utils import clean_string, TOOLKIT_SPLITTER
+from alita_sdk.tools.utils import clean_string
 
 
 class DatasourcesToolkit(BaseToolkit):
@@ -21,21 +21,28 @@ class DatasourcesToolkit(BaseToolkit):
     @classmethod
     def get_toolkit(cls, client: Any, datasource_ids: list[int], toolkit_name: Optional[str] = None, selected_tools: list[str] = []):
         tools = []
-        prefix = clean_string(toolkit_name) + TOOLKIT_SPLITTER if toolkit_name else ''
+        # Use clean toolkit name for context (max 1000 chars in description)
+        toolkit_context = f" [Toolkit: {clean_string(toolkit_name)}]" if toolkit_name else ''
         for datasource_id in datasource_ids:
             datasource = client.datasource(datasource_id)
             ds_name = clean_string(datasource.name)
             if len(ds_name) == 0:
                 raise ToolException(f'Datasource with id {datasource_id} has incorrect name (i.e. special characters, etc.)')
             if len(selected_tools) == 0 or 'chat' in selected_tools:
-                tools.append(DatasourcePredict(name=f'{prefix}chat',
-                                            description=f'Search and summarize. {datasource.description}',
+                description = f'Search and summarize. {datasource.description}'
+                if toolkit_context and len(description + toolkit_context) <= 1000:
+                    description = description + toolkit_context
+                tools.append(DatasourcePredict(name=f'chat',
+                                            description=description,
                                             datasource=datasource, 
                                             args_schema=datasourceToolSchema,
                                             return_type='str'))
             if len(selected_tools) == 0 or 'search' in selected_tools:
-                tools.append(DatasourceSearch(name=f'{prefix}search',
-                                                description=f'Search return results. {datasource.description}',
+                description = f'Search return results. {datasource.description}'
+                if toolkit_context and len(description + toolkit_context) <= 1000:
+                    description = description + toolkit_context
+                tools.append(DatasourceSearch(name=f'search',
+                                                description=description,
                                                 datasource=datasource, 
                                                 args_schema=datasourceToolSchema,
                                                 return_type='str'))