alita-sdk 0.3.220__tar.gz → 0.3.222__tar.gz

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.220
+Version: 0.3.222
 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.220 → alita_sdk-0.3.222}/alita_sdk/runtime/tools/artifact.py

@@ -90,7 +90,6 @@ class ArtifactWrapper(BaseVectorStoreToolApiWrapper):
                 yield Document(page_content=str(value), metadata=metadata)
         else:
             document.page_content = json.dumps(str(page_content))
-            yield document
 
     @extend_with_vector_tools
     def get_available_tools(self):

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/bitbucket/__init__.py

@@ -27,7 +27,7 @@ def get_tools(tool):
         llm=tool['settings'].get('llm', None),
         alita=tool['settings'].get('alita', None),
         connection_string=tool['settings'].get('connection_string', None),
-        collection_name=str(tool['id']),
+        collection_name=str(tool['toolkit_name']),
         doctype='code',
         embedding_model="HuggingFaceEmbeddings",
         embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
@@ -60,6 +60,10 @@ class AlitaBitbucketToolkit(BaseToolkit):
             connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
                                                           default=None,
                                                           json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__=ConfigDict(json_schema_extra=
             {

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/confluence/__init__.py

@@ -27,7 +27,7 @@ def get_tools(tool):
         toolkit_name=tool.get('toolkit_name'),
         # indexer settings
         connection_string = tool['settings'].get('connection_string', None),
-        collection_name=str(tool['id']),
+        collection_name=str(tool['toolkit_name']),
         doctype='doc',
         embedding_model="HuggingFaceEmbeddings",
         embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
@@ -84,6 +84,11 @@ class ConfluenceToolkit(BaseToolkit):
             connection_string = (Optional[SecretStr], Field(description="Connection string for vectorstore",
                                                             default=None,
                                                             json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
+
             __config__=ConfigDict(json_schema_extra={
                 'metadata': {
                     "label": "Confluence",

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/elitea_base.py

@@ -12,6 +12,7 @@ from pydantic import BaseModel, create_model, Field, SecretStr
 from alita_sdk.runtime.langchain.interfaces.llm_processor import get_embeddings
 from .chunkers import markdown_chunker
 from .utils import TOOLKIT_SPLITTER
+from .vector_adapters.VectorStoreAdapter import VectorStoreAdapterFactory
 from ..runtime.utils.utils import IndexerKeywords
 
 logger = logging.getLogger(__name__)
@@ -219,6 +220,10 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
     embedding_model_params: Optional[Dict[str, Any]] = {"model_name": "sentence-transformers/all-MiniLM-L6-v2"}
     vectorstore_type: Optional[str] = "PGVector"
 
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._adapter = VectorStoreAdapterFactory.create_adapter(self.vectorstore_type)
+
     def _index_tool_params(self, **kwargs) -> dict[str, tuple[type, Field]]:
         """
         Returns a list of fields for index_data args schema.
@@ -329,36 +334,24 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
 
     # 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."""
+        """Initializes the vector store wrapper with the provided parameters."""
         try:
             from alita_sdk.runtime.tools.vectorstore import VectorStoreWrapper
         except ImportError:
             from alita_sdk.runtime.tools.vectorstore import VectorStoreWrapper
-        
+
         # Validate collection_suffix length
         if collection_suffix and len(collection_suffix.strip()) > 7:
             raise ToolException("collection_suffix must be 7 characters or less")
-        
+
         # Create collection name with suffix if provided
         collection_name = str(self.collection_name)
         if collection_suffix and collection_suffix.strip():
             collection_name = f"{self.collection_name}_{collection_suffix.strip()}"
-        
-        if self.vectorstore_type == 'PGVector':
-            vectorstore_params = {
-                "use_jsonb": True,
-                "collection_name": collection_name,
-                "create_extension": True,
-                "alita_sdk_options": {
-                    "target_schema": collection_name,
-                },
-                "connection_string": self.connection_string.get_secret_value()
-            }
-        elif self.vectorstore_type == 'Chroma':
-            vectorstore_params = {
-                "collection_name": collection_name,
-                "persist_directory": "./indexer_db"
-            }
+
+        # Get database-specific parameters using adapter
+        connection_string = self.connection_string.get_secret_value() if self.connection_string else None
+        vectorstore_params = self._adapter.get_vectorstore_params(collection_name, connection_string)
 
         return VectorStoreWrapper(
             llm=self.llm,
@@ -371,36 +364,17 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
         )
 
     def remove_index(self, collection_suffix: str = ""):
-        """
-            Cleans the indexed data in the collection
-        """
-
-        self._init_vector_store(collection_suffix)._remove_collection()
+        """Cleans the indexed data in the collection."""
+        vectorstore_wrapper = self._init_vector_store(collection_suffix)
+        collection_name = f"{self.collection_name}_{collection_suffix}" if collection_suffix else str(self.collection_name)
+        self._adapter.remove_collection(vectorstore_wrapper, collection_name)
+        return (f"Collection '{collection_name}' has been removed from the vector store.\n"
+                f"Available collections: {self.list_collections()}")
 
     def list_collections(self):
-        """
-            Lists all collections in the vector store
-        """
-        if self.vectorstore_type.lower() == 'PGVector'.lower():
-            from sqlalchemy import text
-            from sqlalchemy.orm import Session
-
-            # schema_name = self.vectorstore.collection_name
-            with Session(self._init_vector_store().vectorstore.session_maker.bind) as session:
-                get_collections = text("""
-                    SELECT table_schema
-                    FROM information_schema.columns
-                    WHERE udt_name = 'vector';
-                """)
-
-                # Execute the raw SQL query
-                result = session.execute(get_collections)
-
-                # Fetch all rows from the result
-                docs = result.fetchall()
-            return str(docs)
-        vector_client = self._init_vector_store().vectoradapter.vectorstore._client
-        return ','.join([collection.name for collection in vector_client.list_collections()])
+        """Lists all collections in the vector store."""
+        vectorstore_wrapper = self._init_vector_store()
+        return self._adapter.list_collections(vectorstore_wrapper)
 
     def search_index(self,
                      query: str,

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/figma/__init__.py

@@ -22,7 +22,7 @@ def get_tools(tool):
             # indexer settings
             llm=tool['settings'].get('llm', None),
             connection_string = tool['settings'].get('connection_string', None),
-            collection_name=str(tool['id']),
+            collection_name=str(tool['toolkit_name']),
             doctype='doc',
             embedding_model="HuggingFaceEmbeddings",
             embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
@@ -57,6 +57,10 @@ class FigmaToolkit(BaseToolkit):
             connection_string = (Optional[SecretStr], Field(description="Connection string for vectorstore",
                                                             default=None,
                                                             json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
             __config__=ConfigDict(
                 json_schema_extra={
                      "metadata": {

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/github/__init__.py

@@ -25,7 +25,7 @@ def _get_toolkit(tool) -> BaseToolkit:
         llm=tool['settings'].get('llm', None),
         alita=tool['settings'].get('alita', None),
         connection_string=tool['settings'].get('connection_string', None),
-        collection_name=str(tool['id']),
+        collection_name=str(tool['toolkit_name']),
         doctype='code',
         embedding_model="HuggingFaceEmbeddings",
         embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
@@ -96,6 +96,10 @@ class AlitaGitHubToolkit(BaseToolkit):
                                                             default=None,
                                                             json_schema_extra={'secret': True})),
             
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
+
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools}))
         )
 
@@ -123,4 +127,3 @@ class AlitaGitHubToolkit(BaseToolkit):
 
     def get_tools(self):
         return self.tools
-

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/gitlab/__init__.py

@@ -7,7 +7,6 @@ from pydantic import create_model, BaseModel, ConfigDict, SecretStr
 from pydantic.fields import Field
 
 from .api_wrapper import GitLabAPIWrapper
-# from .tools import __all__
 from ..utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
 
 name = "gitlab"
@@ -24,7 +23,7 @@ def get_tools(tool):
         llm=tool['settings'].get('llm', None),
         alita=tool['settings'].get('alita', None),
         connection_string=tool['settings'].get('connection_string', None),
-        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        collection_name=str(tool['toolkit_name']),
         doctype='code',
         embedding_model="HuggingFaceEmbeddings",
         embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
@@ -51,6 +50,10 @@ class AlitaGitlabToolkit(BaseToolkit):
             connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
                                                           default=None,
                                                           json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__=ConfigDict(json_schema_extra={
                 'metadata': {

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/jira/__init__.py

@@ -25,7 +25,7 @@ def get_tools(tool):
         llm=tool['settings'].get('llm', None),
         alita=tool['settings'].get('alita', None),
         connection_string=tool['settings'].get('connection_string', None),
-        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        collection_name=str(tool['toolkit_name']),
         embedding_model="HuggingFaceEmbeddings",
         embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
         vectorstore_type="PGVector",
@@ -81,6 +81,16 @@ class JiraToolkit(BaseToolkit):
             )),
             verify_ssl=(bool, Field(description="Verify SSL", default=True)),
             additional_fields=(Optional[str], Field(description="Additional fields", default="")),
+
+            # indexer settings
+            connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
+                                                          default=None,
+                                                          json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
+
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__=ConfigDict(json_schema_extra={
                 'metadata': {

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/sharepoint/__init__.py

@@ -18,7 +18,7 @@ def get_tools(tool):
         llm=tool['settings'].get('llm'),
         # indexer settings
         connection_string=tool['settings'].get('connection_string', None),
-        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        collection_name=str(tool['toolkit_name']),
         embedding_model="HuggingFaceEmbeddings",
         embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
         vectorstore_type="PGVector")
@@ -43,6 +43,10 @@ class SharepointToolkit(BaseToolkit):
             connection_string = (Optional[SecretStr], Field(description="Connection string for vectorstore",
                                                             default=None,
                                                             json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
             __config__=ConfigDict(json_schema_extra={
                 'metadata': {
                     "label": "Sharepoint", "icon_url": "sharepoint.svg",

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/sharepoint/api_wrapper.py

@@ -172,7 +172,6 @@ class SharepointApiWrapper(BaseVectorStoreToolApiWrapper):
                 yield Document(page_content=str(content), metadata=new_metadata)
         else:
             document.page_content = str(doc_content)
-            yield document
 
     @extend_with_vector_tools
     def get_available_tools(self):

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/testrail/__init__.py

@@ -21,7 +21,7 @@ def get_tools(tool):
 
         # indexer settings
         connection_string=tool['settings'].get('connection_string', None),
-        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        collection_name=f"{tool.get('toolkit_name')}",
         embedding_model="HuggingFaceEmbeddings",
         embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
         vectorstore_type="PGVector"
@@ -55,6 +55,10 @@ class TestrailToolkit(BaseToolkit):
             connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
                                                           default=None,
                                                           json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__=ConfigDict(json_schema_extra={'metadata':
                                                          {"label": "Testrail", "icon_url": "testrail-icon.svg",

alita_sdk-0.3.222/alita_sdk/tools/vector_adapters/VectorStoreAdapter.py

@@ -0,0 +1,91 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+
+
+class VectorStoreAdapter(ABC):
+    """Abstract base class for vector store adapters."""
+
+    @abstractmethod
+    def get_vectorstore_params(self, collection_name: str, connection_string: Optional[str] = None) -> Dict[str, Any]:
+        """Get vector store specific parameters."""
+        pass
+
+    @abstractmethod
+    def list_collections(self, vectorstore_wrapper) -> str:
+        """List all collections in the vector store."""
+        pass
+
+    @abstractmethod
+    def remove_collection(self, vectorstore_wrapper, collection_name: str):
+        """Remove a collection from the vector store."""
+        pass
+
+
+class PGVectorAdapter(VectorStoreAdapter):
+    """Adapter for PGVector database operations."""
+
+    def get_vectorstore_params(self, collection_name: str, connection_string: Optional[str] = None) -> Dict[str, Any]:
+        return {
+            "use_jsonb": True,
+            "collection_name": collection_name,
+            "create_extension": True,
+            "alita_sdk_options": {
+                "target_schema": collection_name,
+            },
+            "connection_string": connection_string
+        }
+
+    def list_collections(self, vectorstore_wrapper) -> str:
+        from sqlalchemy import text
+        from sqlalchemy.orm import Session
+
+        with Session(vectorstore_wrapper.vectorstore.session_maker.bind) as session:
+            get_collections = text("""
+                SELECT table_schema
+                FROM information_schema.columns
+                WHERE udt_name = 'vector';
+            """)
+            result = session.execute(get_collections)
+            docs = result.fetchall()
+        return str(docs)
+
+    def remove_collection(self, vectorstore_wrapper, collection_name: str):
+        vectorstore_wrapper._remove_collection()
+
+
+class ChromaAdapter(VectorStoreAdapter):
+    """Adapter for Chroma database operations."""
+
+    def get_vectorstore_params(self, collection_name: str, connection_string: Optional[str] = None) -> Dict[str, Any]:
+        return {
+            "collection_name": collection_name,
+            "persist_directory": "./indexer_db"
+        }
+
+    def list_collections(self, vectorstore_wrapper) -> str:
+        vector_client = vectorstore_wrapper.vectoradapter.vectorstore._client
+        return ','.join([collection.name for collection in vector_client.list_collections()])
+
+    def remove_collection(self, vectorstore_wrapper, collection_name: str):
+        vectorstore_wrapper.vectoradapter.vectorstore.delete_collection()
+
+
+class VectorStoreAdapterFactory:
+    """Factory for creating vector store adapters."""
+
+    _adapters = {
+        'PGVector': PGVectorAdapter,
+        'Chroma': ChromaAdapter,
+    }
+
+    @classmethod
+    def create_adapter(cls, vectorstore_type: str) -> VectorStoreAdapter:
+        adapter_class = cls._adapters.get(vectorstore_type)
+        if not adapter_class:
+            raise ValueError(f"Unsupported vectorstore type: {vectorstore_type}")
+        return adapter_class()
+
+    @classmethod
+    def register_adapter(cls, vectorstore_type: str, adapter_class: type):
+        """Register a new adapter for a vector store type."""
+        cls._adapters[vectorstore_type] = adapter_class

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/xray/__init__.py

@@ -24,7 +24,7 @@ def get_tools(tool):
 
         # indexer settings
         connection_string=tool['settings'].get('connection_string', None),
-        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        collection_name=str(tool['toolkit_name']),
         embedding_model="HuggingFaceEmbeddings",
         embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
         vectorstore_type="PGVector"
@@ -45,6 +45,16 @@ class XrayToolkit(BaseToolkit):
             client_id=(str, Field(description="Client ID")),
             client_secret=(SecretStr, Field(description="Client secret", json_schema_extra={'secret': True})),
             limit=(Optional[int], Field(description="Limit", default=100)),
+
+            # indexer settings
+            connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
+                                                          default=None,
+                                                          json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
+
             selected_tools=(List[Literal[tuple(selected_tools)]], Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __config__={'json_schema_extra':
                             {

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/zephyr_enterprise/__init__.py

@@ -18,7 +18,7 @@ def get_tools(tool):
 
         # indexer settings
         connection_string=tool['settings'].get('connection_string', None),
-        collection_name=f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        collection_name=str(tool['toolkit_name']),
         embedding_model="HuggingFaceEmbeddings",
         embedding_model_params={"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
         vectorstore_type="PGVector"
@@ -41,6 +41,10 @@ class ZephyrEnterpriseToolkit(BaseToolkit):
             connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
                                                           default=None,
                                                           json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
             selected_tools=(List[Literal[tuple(selected_tools)]], []),
             __config__=ConfigDict(json_schema_extra={
                 'metadata': {
@@ -71,4 +75,3 @@ class ZephyrEnterpriseToolkit(BaseToolkit):
 
     def get_tools(self) -> List[BaseTool]:
         return self.tools
-

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/zephyr_enterprise/api_wrapper.py

@@ -1,15 +1,29 @@
 import logging
-from typing import Optional, List
+from importlib.metadata import metadata
+from operator import ifloordiv
+from typing import Optional, List, Generator
 
 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
+from ..elitea_base import BaseToolApiWrapper, BaseVectorStoreToolApiWrapper, extend_with_vector_tools
 
 logger = logging.getLogger(__name__)
 
-class ZephyrApiWrapper(BaseToolApiWrapper):
+zql_description = """
+                    ZQL query to search for test cases. 
+                    Supported: estimatedTime, testcaseId, creator, release,
+                    project, priority, altId, version,
+                    versionId, automated, folder, contents,
+                    name, comment, tag
+
+                    It has to follow the syntax in examples:
+                    "folder=\"TestToolkit\"", "name~\"TestToolkit5\"
+                    """
+
+class ZephyrApiWrapper(BaseVectorStoreToolApiWrapper):
     base_url: str
     token: SecretStr
     _client: Optional[ZephyrClient] = PrivateAttr()
@@ -39,7 +53,7 @@ class ZephyrApiWrapper(BaseToolApiWrapper):
             return ToolException(f"Unable to retrieve Zephyr entities: {e}")
 
 
-    def get_testcases_by_zql(self, zql: str):
+    def get_testcases_by_zql(self, zql: str, return_as_list: bool = False):
 
         """
         Retrieve testcases by zql.
@@ -50,8 +64,12 @@ class ZephyrApiWrapper(BaseToolApiWrapper):
             testcases = self._client.get_testcases_by_zql(zql)
             parsed_test_cases = []
             if testcases['resultSize'] == 0:
+                if return_as_list:
+                    return []
                 return "No test cases found for the provided ZQL query."
             logger.info(f"Retrieved test cases: {testcases}")
+            if return_as_list:
+                return [test_case['testcase'] for test_case in testcases['results']]
             for test_case in testcases['results']:
                 parsed_test_cases.append(f"Test case ID: {test_case.get('id')}, Test case: {test_case['testcase']}")
             return "\n".join(parsed_test_cases)
@@ -69,6 +87,20 @@ class ZephyrApiWrapper(BaseToolApiWrapper):
         except Exception as e:
             return ToolException(f"Unable to retrieve Zephyr entities: {e}")
 
+    def get_test_case_steps(self, testcase_tree_id: str):
+        last_version = self.get_last_version(testcase_tree_id)
+        return self._client.get_testcase_steps(last_version)
+
+    def get_last_version(self, testcase_tree_id: str):
+        # get test case id
+        test_case_id = self._client.get_test_case(testcase_tree_id)['testcase']['testcaseId']
+        logger.info(f"Test case id: {test_case_id}")
+        # get test case version id
+        return self._client.get_testcase_versions(test_case_id)[-1]['id']
+
+    def get_test_case_steps_by_version(self, test_case_version_id: str):
+        return self._client.get_testcase_steps(test_case_version_id)
+
     def add_steps(self, testcase_tree_id: str, steps: Optional[List[dict]] = []):
 
         """
@@ -81,14 +113,10 @@ class ZephyrApiWrapper(BaseToolApiWrapper):
             if not steps:
                 return ToolException("Steps cannot be empty.")
 
-            # get test case id
-            test_case_id = self._client.get_test_case(testcase_tree_id)['testcase']['testcaseId']
-            logger.info(f"Test case id: {test_case_id}")
-            # get test case version id
-            test_case_version_id = self._client.get_testcase_versions(test_case_id)[-1]['id']
+            test_case_version_id = self.get_last_version(testcase_tree_id)
             logger.info(f"Test case version id: {test_case_version_id}")
             # get test case steps
-            test_case_steps = self._client.get_testcase_steps(test_case_version_id)
+            test_case_steps = self.get_test_case_steps_by_version(test_case_version_id)
             logger.info(f"Test case steps: {test_case_steps}")
             # add steps
             actual_steps = test_case_steps['steps'] if test_case_steps else None
@@ -120,6 +148,30 @@ class ZephyrApiWrapper(BaseToolApiWrapper):
         except Exception as e:
             return ToolException(f"Unable to retrieve Zephyr entities: {e}")
 
+    def _index_tool_params(self, **kwargs) -> dict[str, tuple[type, Field]]:
+        """
+        Returns a list of fields for index_data args schema.
+        """
+        return {
+            "zql": (str, Field(description=zql_description, examples=["folder=\"TestToolkit\"", "name~\"TestToolkit5\""))
+        }
+
+    def _base_loader(self, zql: str) -> Generator[Document, None, None]:
+        test_cases = self.get_testcases_by_zql(zql)
+        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"
+            }
+            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 get_available_tools(self):
         return [
             {
@@ -188,16 +240,7 @@ class ZephyrApiWrapper(BaseToolApiWrapper):
                 "description": self.get_testcases_by_zql.__doc__,
                 "args_schema": create_model(
                     "TestCaseByZqlModel",
-                    zql=(str, Field(description="""
-                    ZQL query to search for test cases. 
-                    Supported: estimatedTime, testcaseId, creator, release,
-                    project, priority, altId, version,
-                    versionId, automated, folder, contents,
-                    name, comment, tag
-                    
-                    It has to follow the syntax in examples:
-                    "folder=\"TestToolkit\"", "name~\"TestToolkit5\"
-                    """, examples=["folder=\"TestToolkit\"", "name~\"TestToolkit5\""]))
+                    zql=(str, Field(description=zql_description, examples=["folder=\"TestToolkit\"", "name~\"TestToolkit5\""]))
                 ),
                 "ref": self.get_testcases_by_zql,
             }

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/zephyr_essential/__init__.py

@@ -18,7 +18,7 @@ def get_tools(tool):
 
         # indexer settings
         connection_string = tool['settings'].get('connection_string', None),
-        collection_name = f"{tool.get('toolkit_name')}_{str(tool['id'])}",
+        collection_name=str(tool['toolkit_name']),
         embedding_model = "HuggingFaceEmbeddings",
         embedding_model_params = {"model_name": "sentence-transformers/all-MiniLM-L6-v2"},
         vectorstore_type = "PGVector"
@@ -41,6 +41,10 @@ class ZephyrEssentialToolkit(BaseToolkit):
             connection_string=(Optional[SecretStr], Field(description="Connection string for vectorstore",
                                                           default=None,
                                                           json_schema_extra={'secret': True})),
+
+            # embedder settings
+            embedding_model=(str, Field(description="Embedding model: i.e. 'HuggingFaceEmbeddings', etc.", default="HuggingFaceEmbeddings")),
+            embedding_model_params=(dict, Field(description="Embedding model parameters: i.e. `{'model_name': 'sentence-transformers/all-MiniLM-L6-v2'}", default={"model_name": "sentence-transformers/all-MiniLM-L6-v2"})),
             __config__={'json_schema_extra': {'metadata': {"label": "Zephyr Essential", "icon_url": "zephyr.svg",
                             "categories": ["test management"],
                             "extra_categories": ["test automation", "test case management", "test planning"]
@@ -67,4 +71,3 @@ class ZephyrEssentialToolkit(BaseToolkit):
 
     def get_tools(self):
         return self.tools
-

{alita_sdk-0.3.220 → alita_sdk-0.3.222}/alita_sdk/tools/zephyr_essential/api_wrapper.py

@@ -253,7 +253,6 @@ class ZephyrEssentialApiWrapper(BaseVectorStoreToolApiWrapper):
                 base_data['test_case_content'] = additional_content
 
             document.page_content = json.dumps(base_data)
-            yield document
         except json.JSONDecodeError as e:
             raise ToolException(f"Failed to decode JSON from document: {e}")