alita-sdk 0.3.215__py3-none-any.whl → 0.3.217__py3-none-any.whl

alita_sdk/tools/__init__.py

@@ -80,6 +80,7 @@ _safe_import_tool('pptx', 'pptx', 'get_tools', 'PPTXToolkit')
 _safe_import_tool('postman', 'postman', 'get_tools', 'PostmanToolkit')
 _safe_import_tool('memory', 'memory', 'get_tools', 'MemoryToolkit')
 _safe_import_tool('zephyr_squad', 'zephyr_squad', 'get_tools', 'ZephyrSquadToolkit')
+_safe_import_tool('zephyr_essential', 'zephyr_essential', 'get_tools', 'ZephyrEssentialToolkit')
 _safe_import_tool('slack', 'slack', 'get_tools', 'SlackToolkit')
 _safe_import_tool('bigquery', 'google.bigquery', 'get_tools', 'BigQueryToolkit')
 _safe_import_tool('delta_lake', 'aws.delta_lake', 'get_tools', 'DeltaLakeToolkit')

alita_sdk/tools/elitea_base.py

@@ -1,5 +1,6 @@
 import ast
 import fnmatch
+import json
 import logging
 import traceback
 from typing import Any, Optional, List, Literal, Dict, Generator
@@ -326,6 +327,7 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
                             yield processed_doc
 
 
+    # TODO: init store once and re-use the instance
     def _init_vector_store(self, collection_suffix: str = "", embeddings: Optional[Any] = None):
         """ Initializes the vector store wrapper with the provided parameters."""
         try:
@@ -375,6 +377,13 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
 
         self._init_vector_store(collection_suffix)._clean_collection()
 
+    def list_collections(self):
+        """
+            Lists all collections in the vector store
+        """
+        vector_client = self._init_vector_store().vectoradapter.vectorstore._client
+        return ','.join([collection.name for collection in vector_client.list_collections()])
+
     def search_index(self,
                      query: str,
                      collection_suffix: str = "",
@@ -386,17 +395,18 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
                      **kwargs):
         """ Searches indexed documents in the vector store."""
         vectorstore = self._init_vector_store(collection_suffix)
-        return vectorstore.search_documents(
-            query, 
-            doctype=self.doctype, 
-            filter=filter, 
-            cut_off=cut_off, 
-            search_top=search_top, 
+        found_docs = vectorstore.search_documents(
+            query,
+            doctype=self.doctype,
+            filter=filter,
+            cut_off=cut_off,
+            search_top=search_top,
             reranker=reranker,
-            full_text_search=full_text_search, 
-            reranking_config=reranking_config, 
+            full_text_search=full_text_search,
+            reranking_config=reranking_config,
             extended_search=extended_search
         )
+        return f"Found {len(found_docs)} documents matching the query\n{json.dumps(found_docs, indent=4)}" if found_docs else "No documents found matching the query."
 
     def stepback_search_index(self,
                      query: str,
@@ -410,17 +420,18 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
                      **kwargs):
         """ Searches indexed documents in the vector store."""
         vectorstore = self._init_vector_store(collection_suffix)
-        return vectorstore.stepback_search(
-            query, 
-            messages, 
-            self.doctype, 
-            filter=filter, 
-            cut_off=cut_off, 
+        found_docs = vectorstore.stepback_search(
+            query,
+            messages,
+            self.doctype,
+            filter=filter,
+            cut_off=cut_off,
             search_top=search_top,
-            full_text_search=full_text_search, 
-            reranking_config=reranking_config, 
+            full_text_search=full_text_search,
+            reranking_config=reranking_config,
             extended_search=extended_search
         )
+        return f"Found {len(found_docs)} documents matching the query\n{json.dumps(found_docs, indent=4)}" if found_docs else "No documents found matching the query."
 
     def stepback_summary_index(self,
                      query: str,
@@ -483,6 +494,14 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
                 "description": self.remove_index.__doc__,
                 "args_schema": RemoveIndexParams
             },
+            {
+                "name": "list_collections",
+                "mode": "list_collections",
+                "ref": self.list_collections,
+                "description": self.list_collections.__doc__,
+                "args_schema": create_model("ListCollectionsParams")  # No parameters
+            },
+
         ]
 
 

alita_sdk/tools/zephyr_enterprise/__init__.py

@@ -13,7 +13,15 @@ def get_tools(tool):
         selected_tools=tool['settings'].get('selected_tools', []),
         base_url=tool['settings']['base_url'],
         token=tool['settings']['token'],
-        toolkit_name=tool.get('toolkit_name')
+        toolkit_name=tool.get('toolkit_name'),
+        llm=tool['settings'].get('llm', None),
+
+        # indexer settings
+        connection_string=tool['settings'].get('connection_string', None),
+        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        embedding_model="HuggingFaceEmbeddings",
+        embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
+        vectorstore_type="PGVector"
     ).get_tools()
 
 class ZephyrEnterpriseToolkit(BaseToolkit):
@@ -29,6 +37,10 @@ class ZephyrEnterpriseToolkit(BaseToolkit):
             name,
             base_url=(str, Field(description="Zephyr Enterprise base URL", json_schema_extra={'toolkit_name': True, 'max_toolkit_length': ZephyrEnterpriseToolkit.toolkit_max_length })),
             token=(SecretStr, Field(description="API token", json_schema_extra={'secret': True})),
+            # indexer settings
+            connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
+                                                          default=None,
+                                                          json_schema_extra={'secret': True})),
             selected_tools=(List[Literal[tuple(selected_tools)]], []),
             __config__=ConfigDict(json_schema_extra={
                 'metadata': {

alita_sdk/tools/zephyr_essential/__init__.py

@@ -0,0 +1,70 @@
+from typing import List, Literal, Optional
+
+from langchain_core.tools import BaseToolkit, BaseTool
+from pydantic import create_model, BaseModel, Field, SecretStr
+
+from .api_wrapper import ZephyrEssentialApiWrapper
+from ..base.tool import BaseAction
+from ..utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
+
+name = "zephyr_essential"
+
+def get_tools(tool):
+    return ZephyrEssentialToolkit().get_toolkit(
+        selected_tools=tool['settings'].get('selected_tools', []),
+        token=tool['settings']["token"],
+        toolkit_name=tool.get('toolkit_name'),
+        llm = tool['settings'].get('llm', None),
+
+        # indexer settings
+        connection_string = tool['settings'].get('connection_string', None),
+        collection_name = f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        embedding_model = "HuggingFaceEmbeddings",
+        embedding_model_params = {"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
+        vectorstore_type = "PGVector"
+    ).get_tools()
+
+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)),
+            selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
+            # indexer settings
+            connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
+                                                          default=None,
+                                                          json_schema_extra={'secret': True})),
+            __config__={'json_schema_extra': {'metadata': {"label": "Zephyr Essential", "icon_url": "zephyr.svg",
+                            "categories": ["test management"],
+                            "extra_categories": ["test automation", "test case management", "test planning"]
+                        }}}
+        )
+
+    @classmethod
+    def get_toolkit(cls, selected_tools: list[str] | None = None, toolkit_name: Optional[str] = None, **kwargs):
+        zephyr_api_wrapper = ZephyrEssentialApiWrapper(**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
+            tools.append(BaseAction(
+                api_wrapper=zephyr_api_wrapper,
+                name=prefix + tool["name"],
+                description=tool["description"],
+                args_schema=tool["args_schema"]
+            ))
+        return cls(tools=tools)
+
+    def get_tools(self):
+        return self.tools
+