alita-sdk 0.3.214__py3-none-any.whl → 0.3.216__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
@@ -350,9 +351,7 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
                 "alita_sdk_options": {
                     "target_schema": collection_name,
                 },
-                # "connection_string": self.connection_string.get_secret_value()
-                # 'postgresql+psycopg://project_23_user:Rxu4QtM2InLVNnm62GX7@pgvector:5432/project_23'
-                "connection_string": 'postgresql+psycopg://postgres:yourpassword@localhost:5432/postgres'
+                "connection_string": self.connection_string.get_secret_value()
             }
         elif self.vectorstore_type == 'Chroma':
             vectorstore_params = {
@@ -377,6 +376,12 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
 
         self._init_vector_store(collection_suffix)._clean_collection()
 
+    def list_collections(self):
+        """
+            Lists all collections in the vector store
+        """
+        return ','.join([collection.name for collection in self.vectoradapter.vectorstore._client.list_collections()])
+
     def search_index(self,
                      query: str,
                      collection_suffix: str = "",
@@ -388,17 +393,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,
@@ -412,17 +418,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,
@@ -485,6 +492,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_essential/__init__.py

@@ -0,0 +1,58 @@
+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')
+    ).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})),
+            __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
+