alita-sdk 0.3.256__tar.gz → 0.3.258__tar.gz

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.256
+Version: 0.3.258
 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>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0

alita_sdk-0.3.258/alita_sdk/configurations/github.py

@@ -0,0 +1,87 @@
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
+
+
+class GithubConfiguration(BaseModel):
+    model_config = ConfigDict(
+        json_schema_extra={
+            "metadata": {
+                "label": "GitHub",
+                "icon_url": None,
+                "sections": {
+                    "auth": {
+                        "required": False,
+                        "subsections": [
+                            {
+                                "name": "Token",
+                                "fields": ["access_token"]
+                            },
+                            {
+                                "name": "Password",
+                                "fields": ["username", "password"]
+                            },
+                            {
+                                "name": "App private key",
+                                "fields": ["app_id", "app_private_key"]
+                            }
+                        ]
+                    },
+                },
+                "section": "credentials",
+                "type": "github",
+                "categories": ["code repositories"],
+                "extra_categories": ["github", "git", "repository", "code", "version control"],
+            }
+        }
+    )
+
+    base_url: Optional[str] = Field(description="Base API URL", default="https://api.github.com")
+    app_id: Optional[str] = Field(description="Github APP ID", default=None)
+    app_private_key: Optional[SecretStr] = Field(description="Github APP private key", default=None)
+
+    access_token: Optional[SecretStr] = Field(description="Github Access Token", default=None)
+
+    username: Optional[str] = Field(description="Github Username", default=None)
+    password: Optional[SecretStr] = Field(description="Github Password", default=None)
+
+    @model_validator(mode='before')
+    @classmethod
+    def validate_auth_sections(cls, data):
+        if not isinstance(data, dict):
+            return data
+
+        has_token = bool(data.get('access_token') and str(data.get('access_token')).strip())
+        has_password = bool(
+            data.get('username') and str(data.get('username')).strip() and
+            data.get('password') and str(data.get('password')).strip()
+        )
+        has_app_key = bool(
+            data.get('app_id') and str(data.get('app_id')).strip() and
+            data.get('app_private_key') and str(data.get('app_private_key')).strip()
+        )
+
+        # If any method is partially configured, raise exception
+        if (
+                (data.get('username') and not data.get('password')) or
+                (data.get('password') and not data.get('username')) or
+                (data.get('app_id') and not data.get('app_private_key')) or
+                (data.get('app_private_key') and not data.get('app_id'))
+        ):
+            raise ValueError(
+                "Authentication is misconfigured: both username and password, or both app_id and app_private_key, must be provided together."
+            )
+
+        # If all are missing, allow anonymous
+        if not (has_token or has_password or has_app_key):
+            return data
+
+        # If any method is fully configured
+        if has_token or has_password or has_app_key:
+            return data
+
+        raise ValueError(
+            "Authentication is misconfigured: provide either Token (access_token), "
+            "Password (username + password), App private key (app_id + app_private_key), "
+            "or leave all blank for anonymous access."
+        )

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/runtime/clients/client.py

@@ -364,8 +364,8 @@ class AlitaClient:
         return data.content
 
     def delete_artifact(self, bucket_name, artifact_name):
-        url = f'{self.artifact_url}/{bucket_name}/{quote(artifact_name)}'
-        data = requests.delete(url, headers=self.headers, verify=False)
+        url = f'{self.artifact_url}/{bucket_name}'
+        data = requests.delete(url, headers=self.headers, verify=False, params={'filename': quote(artifact_name)})
         return self._process_requst(data)
 
     def _prepare_messages(self, messages: list[BaseMessage]):

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/ado/__init__.py

@@ -14,11 +14,12 @@ def get_tools(tool_type, tool):
         "limit": tool['settings'].get('limit', 5),
         "toolkit_name": tool.get('toolkit_name', ''),
         # indexer settings
+        "alita": tool['settings'].get('alita', None),
         "llm": tool['settings'].get('llm', None),
         "pgvector_configuration": tool['settings'].get('pgvector_configuration', {}),
         "collection_name": tool['toolkit_name'],
         "doctype": 'doc',
-        "embedding_model": tool['settings'].get('embedding_configuration', {}).get('name', None),
+        "embedding_model": tool['settings'].get('embedding_model'),
         "vectorstore_type": "PGVector"
     }
     if tool_type == 'ado_plans':

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/ado/repos/__init__.py

@@ -23,10 +23,9 @@ def _get_toolkit(tool) -> BaseToolkit:
         active_branch=tool['settings'].get('active_branch', ""),
         toolkit_name=tool['settings'].get('toolkit_name', ""),
         pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
-        collection_name=tool['toolkit_name'],
-        doctype='code',
         embedding_model=tool['settings'].get('embedding_model'),
-        vectorstore_type="PGVector",
+        collection_name=tool['toolkit_name'],
+        alita=tool['settings'].get('alita', None),
     )
 
 def get_toolkit():
@@ -62,11 +61,6 @@ class AzureDevOpsReposToolkit(BaseToolkit):
                     "icon_url": "ado-repos-icon.svg",
                     "categories": ["code repositories"],
                     "extra_categories": ["code", "repository", "version control"],
-                    # "configuration_group": {
-                    #     "name": "ado_repos",
-                    #     "label": "Azure DevOps Repositories",
-                    #     "icon_url": "ado-repos-icon.svg",
-                    # }
                 }}}
         )
 
@@ -95,6 +89,7 @@ class AzureDevOpsReposToolkit(BaseToolkit):
             **kwargs,
             # TODO use ado_repos_configuration fields
             **kwargs['ado_repos_configuration'],
+            **kwargs['ado_repos_configuration']['ado_configuration'],
             **(kwargs.get('pgvector_configuration') or {}),
         }
         azure_devops_repos_wrapper = ReposApiWrapper(**wrapper_payload)

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/ado/repos/repos_wrapper.py

@@ -251,15 +251,6 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
     token: Optional[SecretStr]
     _client: Optional[GitClient] = PrivateAttr()
 
-    llm: Optional[Any] = None
-    # Vector store configuration
-    connection_string: Optional[SecretStr] = None
-    collection_name: Optional[str] = None
-    doctype: Optional[str] = 'code'
-    embedding_model: Optional[str] = "HuggingFaceEmbeddings"
-    embedding_model_params: Optional[dict] = {"model_name": "sentence-transformers/all-MiniLM-L6-v2"}
-    vectorstore_type: Optional[str] = "PGVector"
-
     class Config:
         arbitrary_types_allowed = True
 

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/ado/test_plan/test_plan_wrapper.py

@@ -16,11 +16,6 @@ from pydantic.fields import FieldInfo as Field
 from ..work_item import AzureDevOpsApiWrapper
 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__)
 
 # Input models for Test Plan operations

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/ado/wiki/ado_wrapper.py

@@ -17,11 +17,6 @@ from pydantic.fields import Field
 
 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__)
 
 GetWikiInput = create_model(
@@ -68,13 +63,6 @@ class AzureDevOpsApiWrapper(BaseVectorStoreToolApiWrapper):
     _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)
 

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/ado/work_item/ado_wrapper.py

@@ -16,11 +16,6 @@ from pydantic.fields import Field
 
 from alita_sdk.tools.non_code_indexer_toolkit import NonCodeIndexerToolkit
 
-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__)
 
 create_wi_field = """JSON of the work item fields to create in Azure DevOps, i.e.

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/base_indexer_toolkit.py

@@ -5,7 +5,6 @@ from typing import Any, Optional, List, Literal, Dict, Generator
 from langchain_core.documents import Document
 from pydantic import create_model, Field, SecretStr
 
-# from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
 from .utils.content_parser import process_content_by_type
 from .vector_adapters.VectorStoreAdapter import VectorStoreAdapterFactory
 from ..runtime.tools.vectorstore_base import VectorStoreWrapperBase
@@ -95,7 +94,6 @@ BaseIndexDataParams = create_model(
                          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")),
-    chunking_tool=(Literal[None,'markdown', 'statistical', 'proposal'], Field(description="Name of chunking tool", default=None)),
     chunking_config=(Optional[dict], Field(description="Chunking tool configuration", default_factory=dict)),
 )
 
@@ -162,7 +160,7 @@ class BaseIndexerToolkit(VectorStoreWrapperBase):
         chunking_config = kwargs.get("chunking_config")
         #
         if clean_index:
-            self._clean_index()
+            self._clean_index(collection_suffix)
         #
         documents = self._base_loader(**kwargs)
         documents = self._reduce_duplicates(documents, collection_suffix)
@@ -173,7 +171,7 @@ class BaseIndexerToolkit(VectorStoreWrapperBase):
         return self._save_index(list(documents), collection_suffix=collection_suffix, progress_step=progress_step)
     
     def _apply_loaders_chunkers(self, documents: Generator[Document, None, None], chunking_tool: str=None, chunking_config=None) -> Generator[Document, None, None]:
-        from alita_sdk.tools.chunkers import __confluence_chunkers__ as chunkers
+        from alita_sdk.tools.chunkers import __all__ as chunkers
 
         if chunking_config is None:
             chunking_config = {}

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/bitbucket/tools.py

@@ -134,7 +134,7 @@ class ListBranchesTool(BaseTool):
     name: str = "list_branches_in_repo"
     description: str = """This tool is a wrapper for the Bitbucket API to fetch a list of all branches in the repository. 
     It will return the name of each branch. No input parameters are required."""
-    args_schema: Type[BaseModel] = None
+    args_schema: Type[BaseModel] = create_model("NoInput")
 
     def _run(self):
         try:

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/chunkers/__init__.py

@@ -2,13 +2,15 @@ from .code.codeparser import parse_code_files_for_db
 from .sematic.statistical_chunker import statistical_chunker
 from .sematic.markdown_chunker import markdown_chunker
 from .sematic.proposal_chunker import proposal_chunker
+from .sematic.json_chunker import json_chunker
 from .models import StatisticalChunkerConfig, MarkdownChunkerConfig, ProposalChunkerConfig
 
 __all__ = {
     'code_parser': parse_code_files_for_db,
     'statistical': statistical_chunker,
     'markdown': markdown_chunker,
-    'proposal': proposal_chunker
+    'proposal': proposal_chunker,
+    'json': json_chunker
 }
 
 __confluence_chunkers__ = {

alita_sdk-0.3.258/alita_sdk/tools/chunkers/sematic/json_chunker.py

@@ -0,0 +1,24 @@
+import json
+import logging
+from typing import Generator
+from langchain_text_splitters import RecursiveJsonSplitter
+from langchain_core.documents import Document
+
+def json_chunker(file_content_generator: Generator[Document, None, None], config: dict, *args, **kwargs) -> Generator[Document, None, None]:
+    max_tokens = config.get("max_tokens", 512)
+    for doc in file_content_generator:
+        try:
+            data_dict = json.loads(doc.page_content)
+            chunks = RecursiveJsonSplitter(max_chunk_size=max_tokens).split_json(json_data=data_dict, convert_lists=True)
+            if len(chunks) == 1:
+                yield doc
+                continue
+            chunk_id = 1
+            for chunk in chunks:
+                metadata = doc.metadata.copy()
+                metadata['chunk_id'] = chunk_id
+                chunk_id += 1
+                yield Document(page_content=json.dumps(chunk), metadata=metadata)
+        except Exception as e:
+            logging.error(f"Failed to chunk document: {e}")
+            yield doc

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/jira/api_wrapper.py

@@ -1237,11 +1237,11 @@ class JiraApiWrapper(BaseVectorStoreToolApiWrapper):
         jql = kwargs.get('jql')
         fields_to_extract = kwargs.get('fields_to_extract')
         fields_to_index = kwargs.get('fields_to_index')
-        include_attachments = kwargs.get('include_attachments', False)
         max_total_issues = kwargs.get('max_total_issues', 1000)
 
-        # set values for skipped attachment extensions
+        # set values for skipped attachment extension
         self._skipped_attachment_extensions = kwargs.get('skip_attachment_extensions', [])
+        self._include_attachments = kwargs.get('include_attachments', False)
         self._included_fields = fields_to_extract.copy() if fields_to_extract else []
 
         try:
@@ -1252,7 +1252,7 @@ class JiraApiWrapper(BaseVectorStoreToolApiWrapper):
             if fields_to_extract:
                 fields.extend(fields_to_extract)
 
-            if include_attachments:
+            if self._include_attachments:
                 fields.append('attachment')
 
             # Use provided JQL query or default to all issues
@@ -1292,36 +1292,36 @@ class JiraApiWrapper(BaseVectorStoreToolApiWrapper):
 
         issue_key = base_document.metadata.get('issue_key')
         # get attachments content
-
-        issue = self._client.issue(issue_key, fields="attachment")
-        attachments = issue.get('fields', {}).get('attachment', [])
-        for attachment in attachments:
-            # get extension
-            ext = f".{attachment['filename'].split('.')[-1].lower()}"
-            if ext not in self._skipped_attachment_extensions:
-                attachment_id = f"attach_{attachment['id']}"
-                base_document.metadata.setdefault(IndexerKeywords.DEPENDENT_DOCS.value, []).append(attachment_id)
-                try:
-                    attachment_content = self._client.get_attachment_content(attachment['id'])
-                except Exception as e:
-                    logger.error(f"Failed to download attachment {attachment['filename']} for issue {issue_key}: {str(e)}")
-                    attachment_content = self._client.get(path=f"secure/attachment/{attachment['id']}/{attachment['filename']}", not_json_response=True)
-                content = load_content_from_bytes(attachment_content, ext, llm=self.llm) if ext not in '.pdf' \
-                    else parse_file_content(file_content=attachment_content, file_name=attachment['filename'], llm=self.llm, is_capture_image=True)
-                if not content:
-                    continue
-                yield Document(page_content=content,
-                               metadata={
-                                   'id': attachment_id,
-                                   'issue_key': issue_key,
-                                   'source': f"{self.base_url}/browse/{issue_key}",
-                                   'filename': attachment['filename'],
-                                   'created': attachment['created'],
-                                   'mimeType': attachment['mimeType'],
-                                   'author': attachment.get('author', {}).get('name'),
-                                   IndexerKeywords.PARENT.value: base_document.metadata.get('id', None),
-                                   'type': 'attachment',
-                               })
+        if self._include_attachments:
+            issue = self._client.issue(issue_key, fields="attachment")
+            attachments = issue.get('fields', {}).get('attachment', [])
+            for attachment in attachments:
+                # get extension
+                ext = f".{attachment['filename'].split('.')[-1].lower()}"
+                if ext not in self._skipped_attachment_extensions:
+                    attachment_id = f"attach_{attachment['id']}"
+                    base_document.metadata.setdefault(IndexerKeywords.DEPENDENT_DOCS.value, []).append(attachment_id)
+                    try:
+                        attachment_content = self._client.get_attachment_content(attachment['id'])
+                    except Exception as e:
+                        logger.error(f"Failed to download attachment {attachment['filename']} for issue {issue_key}: {str(e)}")
+                        attachment_content = self._client.get(path=f"secure/attachment/{attachment['id']}/{attachment['filename']}", not_json_response=True)
+                    content = load_content_from_bytes(attachment_content, ext, llm=self.llm) if ext not in '.pdf' \
+                        else parse_file_content(file_content=attachment_content, file_name=attachment['filename'], llm=self.llm, is_capture_image=True)
+                    if not content:
+                        continue
+                    yield Document(page_content=content,
+                                   metadata={
+                                       'id': attachment_id,
+                                       'issue_key': issue_key,
+                                       'source': f"{self.base_url}/browse/{issue_key}",
+                                       'filename': attachment['filename'],
+                                       'created': attachment['created'],
+                                       'mimeType': attachment['mimeType'],
+                                       'author': attachment.get('author', {}).get('name'),
+                                       IndexerKeywords.PARENT.value: base_document.metadata.get('id', None),
+                                       'type': 'attachment',
+                                   })
 
     def _jql_get_tickets(self, jql, fields="*all", start=0, limit=None, expand=None, validate_query=None):
         """
@@ -1430,9 +1430,9 @@ class JiraApiWrapper(BaseVectorStoreToolApiWrapper):
                                     Field(description="Whether to include attachment content in indexing",
                                           default=False)),
             'max_total_issues': (Optional[int], Field(description="Maximum number of issues to index", default=1000)),
-            'skip_attachment_extensions': (Optional[str], Field(
-                description="Comma-separated list of file extensions to skip when processing attachments",
-                default=None)),
+            'skip_attachment_extensions': (Optional[List[str]], Field(
+                description="List of file extensions to skip when processing attachments: i.e. ['.png', '.jpg']",
+                default=[])),
         }
 
     # def index_data(self,

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/qtest/__init__.py

@@ -35,7 +35,7 @@ class QtestToolkit(BaseToolkit):
             name,
             qtest_configuration=(Optional[QtestConfiguration], Field(description="QTest API token", json_schema_extra={
                 'configuration_types': ['qtest']})),
-            qtest_project_id=(int, Field(description="QTest project id", json_schema_extra={'toolkit_name': True,
+            qtest_project_id=(int, Field(default=None, description="QTest project id", json_schema_extra={'toolkit_name': True,
                                                                                             'max_toolkit_length': QtestToolkit.toolkit_max_length})),
             selected_tools=(List[Literal[tuple(selected_tools)]],
                             Field(default=[], json_schema_extra={'args_schemas': selected_tools})),

alita_sdk-0.3.258/alita_sdk/tools/utils/available_tools_decorator.py

@@ -0,0 +1,6 @@
+def extend_with_parent_available_tools(method):
+    def wrapper(self, *args, **kwargs):
+        child_tools = method(self, *args, **kwargs)
+        parent_tools = super(self.__class__, self).get_available_tools()
+        return parent_tools + child_tools
+    return wrapper

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/xray/__init__.py

@@ -17,9 +17,7 @@ name = "xray_cloud"
 def get_tools(tool):
     return XrayToolkit().get_toolkit(
         selected_tools=tool['settings'].get('selected_tools', []),
-        base_url=tool['settings'].get('base_url', None),
-        client_id=tool['settings'].get('client_id', None),
-        client_secret=tool['settings'].get('client_secret', None),
+        xray_configuration=tool['settings'].get('xray_configuration', {}),
         limit=tool['settings'].get('limit', 20),
         verify_ssl=tool['settings'].get('verify_ssl', True),
         toolkit_name=tool.get('toolkit_name'),
@@ -29,7 +27,6 @@ def get_tools(tool):
         pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
         embedding_model=tool['settings'].get('embedding_model'),
         collection_name=str(tool['toolkit_name']),
-        vectorstore_type="PGVector"
     ).get_tools()
 
 
@@ -44,7 +41,7 @@ class XrayToolkit(BaseToolkit):
         return create_model(
             name,
             limit=(Optional[int], Field(description="Limit", default=100)),
-            xray_configuration=(Optional[XrayConfiguration], Field(description="Xray Configuration", json_schema_extra={'configuration_types': ['xray']})),
+            xray_configuration=(XrayConfiguration, Field(description="Xray Configuration", json_schema_extra={'configuration_types': ['xray']})),
             pgvector_configuration=(Optional[PgVectorConfiguration], Field(default=None,
                                                                            description="PgVector Configuration",
                                                                            json_schema_extra={
@@ -71,7 +68,7 @@ class XrayToolkit(BaseToolkit):
         wrapper_payload = {
             **kwargs,
             # Use xray_configuration fields
-            **kwargs.get('xray_configuration', {}),
+            **(kwargs.get('xray_configuration') or {}),
             **(kwargs.get('pgvector_configuration') or {}),
         }
         xray_api_wrapper = XrayApiWrapper(**wrapper_payload)

{alita_sdk-0.3.256 → alita_sdk-0.3.258}/alita_sdk/tools/zephyr_enterprise/api_wrapper.py

@@ -1,14 +1,13 @@
 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
 
 logger = logging.getLogger(__name__)
 
@@ -23,7 +22,7 @@ zql_description = """
                     "folder=\"TestToolkit\"", "name~\"TestToolkit5\"
                     """
 
-class ZephyrApiWrapper(BaseVectorStoreToolApiWrapper):
+class ZephyrApiWrapper(NonCodeIndexerToolkit):
     base_url: str
     token: SecretStr
     _client: Optional[ZephyrClient] = PrivateAttr()
@@ -34,7 +33,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,7 +152,8 @@ class ZephyrApiWrapper(BaseVectorStoreToolApiWrapper):
         Returns a list of fields for index_data args schema.
         """
         return {
-            "zql": (str, Field(description=zql_description, examples=["folder=\"TestToolkit\"", "name~\"TestToolkit5\""]))
+            "zql": (str, Field(description=zql_description, examples=["folder=\"TestToolkit\"", "name~\"TestToolkit5\""])),
+            'chunking_tool': (Literal['json'], Field(description="Name of chunking tool", default='json'))
         }
 
     def _base_loader(self, zql: str, **kwargs) -> Generator[Document, None, None]:
@@ -166,12 +166,17 @@ class ZephyrApiWrapper(BaseVectorStoreToolApiWrapper):
             }
             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
-
-    @extend_with_vector_tools
+    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)
+                document.page_content = test_case_content
+            except Exception as e:
+                logging.error(f"Failed to process document: {e}")
+            yield document
+
+    @extend_with_parent_available_tools
     def get_available_tools(self):
         return [
             {