alita-sdk 0.3.257__py3-none-any.whl → 0.3.562__py3-none-any.whl

alita_sdk/tools/xray/api_wrapper.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import hashlib
-from typing import Any, Dict, Generator, List, Optional
+from typing import Any, Dict, Generator, List, Optional, Literal
 
 import requests
 from langchain_core.documents import Document
@@ -9,12 +9,9 @@ from langchain_core.tools import ToolException
 from pydantic import PrivateAttr, SecretStr, create_model, model_validator, Field
 from python_graphql_client import GraphqlClient
 
-from ..elitea_base import (
-    BaseVectorStoreToolApiWrapper,
-    extend_with_vector_tools,
-)
+from ..non_code_indexer_toolkit import NonCodeIndexerToolkit
+from ..utils.available_tools_decorator import extend_with_parent_available_tools
 from ...runtime.utils.utils import IndexerKeywords
-from ..utils.content_parser import parse_file_content, load_content_from_bytes
 
 try:
     from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
@@ -31,7 +28,7 @@ _get_tests_query = """query GetTests($jql: String!, $limit:Int!, $start: Int)
         limit
         results {
             issueId
-            jira(fields: ["key", "summary", "created", "updated", "assignee.displayName", "reporter.displayName"])
+            jira(fields: ["key", "summary", "description", "created", "updated", "assignee.displayName", "reporter.displayName"])
             projectId
             testType {
                 name
@@ -107,7 +104,7 @@ XrayCreateTest = create_model(
 
 XrayCreateTests = create_model(
     "XrayCreateTests",
-    graphql_mutations=(list[str], Field(description="list of GraphQL mutations:\n" + _graphql_mutation_description))
+    graphql_mutations=(List[str], Field(description="list of GraphQL mutations:\n" + _graphql_mutation_description))
 )
 
 def _parse_tests(test_results) -> List[Any]:
@@ -120,7 +117,7 @@ def _parse_tests(test_results) -> List[Any]:
     return test_results
 
 
-class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
+class XrayApiWrapper(NonCodeIndexerToolkit):
     _default_base_url: str = 'https://xray.cloud.getxray.app'
     base_url: str = ""
     client_id: str = None
@@ -147,7 +144,7 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
         client_id = values['client_id']
         client_secret = values['client_secret']
         # Authenticate to get the token
-        values['base_url'] = values.get('base_url', '') or cls._default_base_url
+        values['base_url'] = values.get('base_url', '') or cls._default_base_url.default
         auth_url = f"{values['base_url']}/api/v1/authenticate"
         auth_data = {
             "client_id": client_id,
@@ -168,7 +165,7 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
                 return ToolException(f"Please, check you credentials ({values['client_id']} / {masked_secret}). Unable")
             else:
                 return ToolException(f"Authentication failed: {str(e)}")
-        return values
+        return super().validate_toolkit(values)
 
     def __init__(self, **data):
         super().__init__(**data)
@@ -333,6 +330,7 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
 
             for test in tests_data:
                 page_content = ""
+                content_structure = {}
                 test_type_name = test.get("testType", {}).get("name", "").lower()
 
                 attachment_ids = []
@@ -359,19 +357,16 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
                     content_structure = {"steps": steps_content}
                     if attachment_ids:
                         content_structure["attachment_ids"] = sorted(attachment_ids)
-                    page_content = json.dumps(content_structure, indent=2)
 
                 elif test_type_name == "cucumber" and test.get("gherkin"):
                     content_structure = {"gherkin": test["gherkin"]}
                     if attachment_ids:
                         content_structure["attachment_ids"] = sorted(attachment_ids)
-                    page_content = json.dumps(content_structure, indent=2)
 
                 elif test.get("unstructured"):
                     content_structure = {"unstructured": test["unstructured"]}
                     if attachment_ids:
                         content_structure["attachment_ids"] = sorted(attachment_ids)
-                    page_content = json.dumps(content_structure, indent=2)
 
                 metadata = {"doctype": self.doctype}
 
@@ -382,7 +377,12 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
 
                     if "created" in jira_data:
                         metadata["created_on"] = jira_data["created"]
-                    
+
+                    if jira_data.get("description"):
+                        content_structure["description"] = jira_data.get("description")
+
+                    page_content = json.dumps(content_structure if content_structure.items() else "", indent=2)
+
                     content_hash = hashlib.sha256(page_content.encode('utf-8')).hexdigest()[:16]
                     metadata["updated_on"] = content_hash
 
@@ -407,11 +407,13 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
                         if "attachments" in step and step["attachments"]:
                             for attachment in step["attachments"]:
                                 if attachment and "id" in attachment and "filename" in attachment:
+                                    attachment['step_id'] = step['id']
                                     attachments_data.append(attachment)
                     if attachments_data:
                         metadata["_attachments_data"] = attachments_data
 
-                yield Document(page_content=page_content, metadata=metadata)
+                metadata[IndexerKeywords.CONTENT_IN_BYTES.value] = page_content.encode('utf-8')
+                yield Document(page_content='', metadata=metadata)
                 
         except Exception as e:
             logger.error(f"Error processing test data: {e}")
@@ -430,14 +432,7 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
             Generator[Document, None, None]: A generator yielding processed Document objects with metadata.
         """
         try:
-            if not getattr(self, '_include_attachments', False):
-                yield document
-                return
-                
             attachments_data = document.metadata.get("_attachments_data", [])
-            if not attachments_data:
-                yield document
-                return
                 
             issue_id = document.metadata.get("id")
             
@@ -458,44 +453,33 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
                 ).append(attachment_id)
 
                 try:
-                    content = self._process_attachment(attachment)
-                    if not content or content.startswith("Attachment processing failed"):
-                        logger.warning(f"Skipping attachment {filename} due to processing failure")
-                        continue
+                    attachment_metadata = {
+                        'id': str(attachment_id),
+                        'issue_key': document.metadata.get('key', ''),
+                        'issueId': str(issue_id),
+                        'projectId': document.metadata.get('projectId', ''),
+                        'source': f"xray_test_{issue_id}",
+                        'filename': filename,
+                        'download_link': attachment.get('downloadLink', ''),
+                        'entity_type': 'test_case_attachment',
+                        'step_id': attachment.get('step_id', ''),
+                        'key': document.metadata.get('key', ''),
+                        IndexerKeywords.PARENT.value: document.metadata.get('id', str(issue_id)),
+                        'type': 'attachment',
+                        'doctype': self.doctype,
+                    }
+                    yield from self._process_attachment(attachment, attachment_metadata)
                 except Exception as e:
                     logger.error(f"Failed to process attachment {filename}: {str(e)}")
                     continue
-
-                attachment_metadata = {
-                    'id': str(attachment_id),
-                    'issue_key': document.metadata.get('key', ''),
-                    'issueId': str(issue_id),
-                    'projectId': document.metadata.get('projectId', ''),
-                    'source': f"xray_test_{issue_id}",
-                    'filename': filename,
-                    'download_link': attachment.get('downloadLink', ''),
-                    'entity_type': 'test_case_attachment',
-                    'key': document.metadata.get('key', ''),
-                    IndexerKeywords.PARENT.value: document.metadata.get('id', str(issue_id)),
-                    'type': 'attachment',
-                    'doctype': self.doctype,
-                }
-
-                yield Document(
-                    page_content=content,
-                    metadata=attachment_metadata
-                )
             
             if "_attachments_data" in document.metadata:
                 del document.metadata["_attachments_data"]
             
-            yield document
-            
         except Exception as e:
             logger.error(f"Error processing document for attachments: {e}")
-            yield document
 
-    def _process_attachment(self, attachment: Dict[str, Any]) -> str:
+    def _process_attachment(self, attachment: Dict[str, Any], attachment_metadata) -> Generator[Document, None, None]:
         """
         Processes an attachment to extract its content.
 
@@ -508,38 +492,17 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
         try:
             download_link = attachment.get('downloadLink')
             filename = attachment.get('filename', '')
-            
-            if not download_link:
-                return f"Attachment: {filename} (no download link available)"
 
             try:
                 auth_token = self._ensure_auth_token()
                 headers = {'Authorization': f'Bearer {auth_token}'}
                 response = requests.get(download_link, headers=headers, timeout=30)
                 response.raise_for_status()
-                
-                ext = f".{filename.split('.')[-1].lower()}" if filename and '.' in filename else ""
-                
-                if ext == '.pdf':
-                    content = parse_file_content(
-                        file_content=response.content,
-                        file_name=filename,
-                        llm=self.llm,
-                        is_capture_image=True
-                    )
-                else:
-                    content = load_content_from_bytes(
-                        response.content,
-                        ext,
-                        llm=self.llm
-                    )
-
-                if content:
-                    return f"filename: {filename}\ncontent: {content}"
-                else:
-                    logger.warning(f"No content extracted from attachment {filename}")
-                    return f"filename: {filename}\ncontent: [No extractable content]"
-                
+
+                yield from self._load_attachment(content=response.content,
+                                                 file_name=filename,
+                                                 attachment_metadata=attachment_metadata)
+
             except requests.RequestException as req_e:
                 logger.error(f"Unable to download attachment {filename} with existing token: {req_e}")
                 
@@ -560,23 +523,13 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
                         fresh_headers = {'Authorization': f'Bearer {fresh_token}'}
                         response = requests.get(download_link, headers=fresh_headers, timeout=60)
                         response.raise_for_status()
-                        
-                        ext = f".{filename.split('.')[-1].lower()}" if filename and '.' in filename else ""
-                        content = parse_file_content(
-                            file_content=response.content,
-                            file_name=filename,
-                            llm=self.llm,
-                            is_capture_image=True
-                        ) if ext == '.pdf' else load_content_from_bytes(response.content, ext, llm=self.llm)
-                        
-                        if content:
-                            return f"filename: {filename}\ncontent: {content}"
-                        else:
-                            return f"filename: {filename}\ncontent: [Content extraction failed after re-auth]"
+
+                        yield from self._load_attachment(content=response.content,
+                                                         file_name=filename,
+                                                         attachment_metadata=attachment_metadata)
                             
                     except Exception as reauth_e:
                         logger.error(f"Re-authentication and retry failed for {filename}: {reauth_e}")
-                        return f"Attachment: {filename} (download failed: {str(req_e)}, re-auth failed: {str(reauth_e)})"
                 else:
                     try:
                         auth_token = self._ensure_auth_token()
@@ -587,34 +540,29 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
                         }
                         response = requests.get(download_link, headers=fallback_headers, timeout=60)
                         response.raise_for_status()
-                        
-                        ext = f".{filename.split('.')[-1].lower()}" if filename and '.' in filename else ""
-                        content = parse_file_content(
-                            file_content=response.content,
-                            file_name=filename,
-                            llm=self.llm,
-                            is_capture_image=True
-                        ) if ext == '.pdf' else load_content_from_bytes(response.content, ext, llm=self.llm)
-                        
-                        if content:
-                            return f"filename: {filename}\ncontent: {content}"
-                        else:
-                            return f"filename: {filename}\ncontent: [Content extraction failed after fallback]"
+
+                        yield from self._load_attachment(content=response.content,
+                                                         file_name=filename,
+                                                         attachment_metadata=attachment_metadata)
                             
                     except Exception as fallback_e:
                         logger.error(f"Fallback download also failed for {filename}: {fallback_e}")
-                        return f"Attachment: {filename} (download failed: {str(req_e)}, fallback failed: {str(fallback_e)})"
                     
             except Exception as parse_e:
                 logger.error(f"Unable to parse attachment {filename}: {parse_e}")
-                return f"Attachment: {filename} (parsing failed: {str(parse_e)})"
                 
         except Exception as e:
             logger.error(f"Error processing attachment: {e}")
-            return f"Attachment processing failed: {str(e)}"
+
+    def _load_attachment(self, content, file_name, attachment_metadata) -> Generator[Document, None, None]:
+        attachment_metadata[IndexerKeywords.CONTENT_IN_BYTES.value] = content
+        attachment_metadata[IndexerKeywords.CONTENT_FILE_NAME.value] = file_name
+        yield Document(page_content='', metadata=attachment_metadata)
 
     def _index_tool_params(self, **kwargs) -> dict[str, tuple[type, Field]]:
         return {
+            'chunking_tool': (Literal['json', ''],
+                              Field(description="Name of chunking tool for base document", default='json')),
             'jql': (Optional[str], Field(description="""JQL query for searching test cases in Xray.
 
             Standard JQL query syntax for filtering Xray test cases. Examples:
@@ -684,9 +632,9 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
         except Exception as e:
             raise ToolException(f"Unable to execute GraphQL due to error: {str(e)}")
 
-    @extend_with_vector_tools
+    @extend_with_parent_available_tools
     def get_available_tools(self):
-        tools = [
+        return [
             {
                 "name": "get_tests",
                 "description": self.get_tests.__doc__,
@@ -711,7 +659,4 @@ class XrayApiWrapper(BaseVectorStoreToolApiWrapper):
                 "args_schema": XrayGrapql,
                 "ref": self.execute_graphql,
             }
-        ]
-
-        tools.extend(self._get_vector_search_tools())
-        return tools
+        ]

alita_sdk/tools/yagmail/__init__.py

@@ -34,7 +34,7 @@ class AlitaYagmailToolkit(BaseToolkit):
         )
 
     @classmethod
-    def get_toolkit(cls, selected_tools: list[str] | None = None, **kwargs):
+    def get_toolkit(cls, selected_tools: list[str] | None = None, toolkit_name: Optional[str] = None, **kwargs):
         if selected_tools is None:
             selected_tools = []
         yagmail_wrapper = YagmailWrapper(**kwargs)
@@ -44,11 +44,16 @@ class AlitaYagmailToolkit(BaseToolkit):
             if selected_tools:
                 if tool["name"] not in selected_tools:
                     continue
+            description = tool["description"]
+            if toolkit_name:
+                description = f"Toolkit: {toolkit_name}\n{description}"
+            description = description[:1000]
             tools.append(BaseAction(
                 api_wrapper=yagmail_wrapper,
                 name=tool["name"],
-                description=tool["description"],
-                args_schema=tool["args_schema"]
+                description=description,
+                args_schema=tool["args_schema"],
+                metadata={"toolkit_name": toolkit_name} if toolkit_name else {}
             ))
         return cls(tools=tools)
 

alita_sdk/tools/zephyr/__init__.py

@@ -7,7 +7,8 @@ from pydantic import create_model, BaseModel, Field, SecretStr
 
 from ..base.tool import BaseAction
 from .api_wrapper import ZephyrV1ApiWrapper
-from ..utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
+from ..elitea_base import filter_missconfigured_index_tools
+from ..utils import clean_string, get_max_toolkit_length
 
 name = "zephyr"
 
@@ -22,15 +23,13 @@ def get_tools(tool):
 
 class ZephyrToolkit(BaseToolkit):
     tools: List[BaseTool] = []
-    toolkit_max_length: int = 0
 
     @staticmethod
     def toolkit_config_schema() -> BaseModel:
         selected_tools = {x['name']: x['args_schema'].schema() for x in ZephyrV1ApiWrapper.model_construct().get_available_tools()}
-        ZephyrToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         return create_model(
             name,
-            base_url=(str, Field(description="Base URL", json_schema_extra={'toolkit_name': True, 'max_toolkit_length': ZephyrToolkit.toolkit_max_length})),
+            base_url=(str, Field(description="Base URL")),
             username=(str, Field(description="Username")),
             password=(SecretStr, Field(description="Password", json_schema_extra={'secret': True})),
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
@@ -45,20 +44,25 @@ class ZephyrToolkit(BaseToolkit):
         )
 
     @classmethod
+    @filter_missconfigured_index_tools
     def get_toolkit(cls, selected_tools: list[str] | None = None, toolkit_name: Optional[str] = None, **kwargs):
         zephyr_api_wrapper = ZephyrV1ApiWrapper(**kwargs)
-        prefix = clean_string(toolkit_name, cls.toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
         available_tools = zephyr_api_wrapper.get_available_tools()
         tools = []
         for tool in available_tools:
             if selected_tools:
                 if tool["name"] not in selected_tools:
                     continue
+            description = tool["description"]
+            if toolkit_name:
+                description = f"Toolkit: {toolkit_name}\n{description}"
+            description = description[:1000]
             tools.append(BaseAction(
                 api_wrapper=zephyr_api_wrapper,
                 name=tool["name"],
-                description=tool["description"],
-                args_schema=tool["args_schema"]
+                description=description,
+                args_schema=tool["args_schema"],
+                metadata={"toolkit_name": toolkit_name} if toolkit_name else {}
             ))
         return cls(tools=tools)
 

alita_sdk/tools/zephyr_enterprise/__init__.py

@@ -4,7 +4,8 @@ from typing import List, Literal, Optional
 
 from .api_wrapper import ZephyrApiWrapper
 from ..base.tool import BaseAction
-from ..utils import clean_string, get_max_toolkit_length, TOOLKIT_SPLITTER
+from ..elitea_base import filter_missconfigured_index_tools
+from ..utils import clean_string, get_max_toolkit_length
 from ...configurations.pgvector import PgVectorConfiguration
 from ...configurations.zephyr_enterprise import ZephyrEnterpriseConfiguration
 
@@ -27,23 +28,22 @@ def get_tools(tool):
 
 class ZephyrEnterpriseToolkit(BaseToolkit):
     tools: List[BaseTool] = []
-    toolkit_max_length: int = 0
 
     @staticmethod
     def toolkit_config_schema() -> BaseModel:
         selected_tools = {x['name']: x['args_schema'].schema() for x in
                           ZephyrApiWrapper.model_construct().get_available_tools()}
-        ZephyrEnterpriseToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         return create_model(
             name,
-            zephyr_configuration=(Optional[ZephyrEnterpriseConfiguration], Field(description="Zephyr Configuration", json_schema_extra={'configuration_types': ['zephyr-enterprise']})),
+            zephyr_configuration=(ZephyrEnterpriseConfiguration, Field(description="Zephyr Configuration", json_schema_extra={'configuration_types': ['zephyr_enterprise']})),
             pgvector_configuration=(Optional[PgVectorConfiguration], Field(description="PgVector Configuration",
                                                                            json_schema_extra={
                                                                                'configuration_types': ['pgvector']},
                                                                            default=None)),
             # embedder settings
             embedding_model=(Optional[str], Field(default=None, description="Embedding configuration.", json_schema_extra={'configuration_model': 'embedding'})),
-            selected_tools=(List[Literal[tuple(selected_tools)]], []),
+            selected_tools=(List[Literal[tuple(selected_tools)]],
+                            Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__=ConfigDict(json_schema_extra={
                 'metadata': {
                     "label": "Zephyr Enterprise", "icon_url": "zephyr.svg",
@@ -53,6 +53,7 @@ class ZephyrEnterpriseToolkit(BaseToolkit):
         )
 
     @classmethod
+    @filter_missconfigured_index_tools
     def get_toolkit(cls, selected_tools: list[str] | None = None, toolkit_name: Optional[str] = None, **kwargs):
         if selected_tools is None:
             selected_tools = []
@@ -64,17 +65,22 @@ class ZephyrEnterpriseToolkit(BaseToolkit):
             **(kwargs.get('embedding_configuration') or {}),
         }
         zephyr_api_wrapper = ZephyrApiWrapper(**wrapper_payload)
-        prefix = clean_string(toolkit_name, cls.toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
         available_tools = zephyr_api_wrapper.get_available_tools()
         tools = []
         for tool in available_tools:
             if selected_tools and tool["name"] not in selected_tools:
                 continue
+            description = tool["description"]
+            if toolkit_name:
+                description = f"Toolkit: {toolkit_name}\n{description}"
+            description = description + "\nZephyr Enterprise instance: " + zephyr_api_wrapper.base_url
+            description = description[:1000]
             tools.append(BaseAction(
                 api_wrapper=zephyr_api_wrapper,
-                name=prefix + tool["name"],
-                description=tool["description"] + "\nZephyr Enterprise instance: " + zephyr_api_wrapper.base_url,
-                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/tools/zephyr_enterprise/api_wrapper.py

@@ -1,14 +1,16 @@
+import json
 import logging
-from importlib.metadata import metadata
-from operator import ifloordiv
-from typing import Optional, List, Generator
+from typing import Optional, List, Generator, Literal
 
 from langchain_core.tools import ToolException
 from pydantic import create_model, model_validator, PrivateAttr, Field, SecretStr
 
 from langchain_core.documents import Document
 from .zephyr_enterprise import ZephyrClient
-from ..elitea_base import BaseToolApiWrapper, BaseVectorStoreToolApiWrapper, extend_with_vector_tools
+from ..non_code_indexer_toolkit import NonCodeIndexerToolkit
+from ..utils.available_tools_decorator import extend_with_parent_available_tools
+from ..utils.content_parser import file_extension_by_chunker
+from ...runtime.utils.utils import IndexerKeywords
 
 logger = logging.getLogger(__name__)
 
@@ -23,7 +25,7 @@ zql_description = """
                     "folder=\"TestToolkit\"", "name~\"TestToolkit5\"
                     """
 
-class ZephyrApiWrapper(BaseVectorStoreToolApiWrapper):
+class ZephyrApiWrapper(NonCodeIndexerToolkit):
     base_url: str
     token: SecretStr
     _client: Optional[ZephyrClient] = PrivateAttr()
@@ -34,7 +36,7 @@ class ZephyrApiWrapper(BaseVectorStoreToolApiWrapper):
         base_url = values.get('base_url')
         token = values.get('token')
         cls._client = ZephyrClient(base_url=base_url, token=token)
-        return values
+        return super().validate_toolkit(values)
 
     def get_test_case(self, testcase_id: str):
 
@@ -153,25 +155,38 @@ class ZephyrApiWrapper(BaseVectorStoreToolApiWrapper):
         Returns a list of fields for index_data args schema.
         """
         return {
+            'chunking_tool': (Literal['json', ''], Field(description="Name of chunking tool", default='json')),
             "zql": (str, Field(description=zql_description, examples=["folder=\"TestToolkit\"", "name~\"TestToolkit5\""]))
+
         }
 
     def _base_loader(self, zql: str, **kwargs) -> Generator[Document, None, None]:
-        test_cases = self.get_testcases_by_zql(zql)
+        self._chunking_tool = kwargs.get('chunking_tool', None)
+        test_cases = self.get_testcases_by_zql(zql=zql, return_as_list=True)
         for test_case in test_cases:
             metadata = {
-                ("updated_on" if k == "lastModifiedOn" else "id" if k == "testcaseId" else k): str(v)
-                for k, v in test_case
-                if k != "id"
+                "updated_on": str(test_case.get("lastModifiedOn")),
+                "id": str(test_case.get("id")),
+                "name": test_case.get("name"),
+                "testcaseId": str(test_case.get("testcaseId")),
+                "projectId": test_case.get("projectId"),
+                "projectName": test_case.get("projectName"),
+                "testcaseType": test_case.get("testcaseType"),
             }
             yield Document(page_content='', metadata=metadata)
 
-    def _process_document(self, document: Document) -> Generator[Document, None, None]:
-        id = document.metadata['id']
-        test_case_content = self.get_test_case_steps(id)
-        document.page_content = test_case_content
+    def _extend_data(self, documents: Generator[Document, None, None]) -> Generator[Document, None, None]:
+        for document in documents:
+            try:
+                id = document.metadata['id']
+                test_case_content = self.get_test_case_steps(id)
+                page_content = json.dumps(test_case_content)
+                document.metadata[IndexerKeywords.CONTENT_IN_BYTES.value] = page_content.encode('utf-8')
+            except Exception as e:
+                logging.error(f"Failed to process document: {e}")
+            yield document
 
-    @extend_with_vector_tools
+    @extend_with_parent_available_tools
     def get_available_tools(self):
         return [
             {

alita_sdk/tools/zephyr_essential/__init__.py

@@ -5,15 +5,17 @@ from pydantic import create_model, BaseModel, Field
 
 from .api_wrapper import ZephyrEssentialApiWrapper
 from ..base.tool import BaseAction
-from ..utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
+from ..elitea_base import filter_missconfigured_index_tools
+from ..utils import clean_string, get_max_toolkit_length
 from ...configurations.pgvector import PgVectorConfiguration
+from ...configurations.zephyr_essential import ZephyrEssentialConfiguration
 
 name = "zephyr_essential"
 
 def get_tools(tool):
     return ZephyrEssentialToolkit().get_toolkit(
         selected_tools=tool['settings'].get('selected_tools', []),
-        token=tool['settings']["token"],
+        zephyr_essential_configuration=tool['settings']['zephyr_essential_configuration'],
         toolkit_name=tool.get('toolkit_name'),
         llm = tool['settings'].get('llm', None),
         alita=tool['settings'].get('alita', None),
@@ -27,16 +29,13 @@ def get_tools(tool):
 
 class ZephyrEssentialToolkit(BaseToolkit):
     tools: List[BaseTool] = []
-    toolkit_max_length: int = 0
 
     @staticmethod
     def toolkit_config_schema() -> BaseModel:
         selected_tools = {x['name']: x['args_schema'].schema() for x in ZephyrEssentialApiWrapper.model_construct().get_available_tools()}
-        ZephyrEssentialToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         return create_model(
             name,
-            token=(str, Field(description="Bearer api token")),
-            base_url=(Optional[str], Field(description="Zephyr Essential base url", default=None)),
+            zephyr_essential_configuration=(ZephyrEssentialConfiguration, Field(description="Zephyr Essential Configuration", json_schema_extra={'configuration_types': ['zephyr_essential']})),
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             pgvector_configuration=(Optional[PgVectorConfiguration], Field(default=None,
                                                                            description="PgVector Configuration",
@@ -51,26 +50,32 @@ class ZephyrEssentialToolkit(BaseToolkit):
         )
 
     @classmethod
+    @filter_missconfigured_index_tools
     def get_toolkit(cls, selected_tools: list[str] | None = None, toolkit_name: Optional[str] = None, **kwargs):
         if selected_tools is None:
             selected_tools = []
         wrapper_payload = {
             **kwargs,
+            **kwargs.get('zephyr_essential_configuration', {}),
             **(kwargs.get('pgvector_configuration') or {}),
         }
         zephyr_api_wrapper = ZephyrEssentialApiWrapper(**wrapper_payload)
-        prefix = clean_string(toolkit_name, cls.toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
         available_tools = zephyr_api_wrapper.get_available_tools()
         tools = []
         for tool in available_tools:
             if selected_tools:
                 if tool["name"] not in selected_tools:
                     continue
+            description = tool["description"]
+            if toolkit_name:
+                description = f"Toolkit: {toolkit_name}\n{description}"
+            description = description[:1000]
             tools.append(BaseAction(
                 api_wrapper=zephyr_api_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)