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

alita_sdk/tools/vector_adapters/VectorStoreAdapter.py

@@ -2,6 +2,8 @@ from abc import ABC, abstractmethod
 from typing import Any, Dict, Optional, List
 from logging import getLogger
 
+from ...runtime.utils.utils import IndexerKeywords
+
 logger = getLogger(__name__)
 
 
@@ -24,13 +26,13 @@ class VectorStoreAdapter(ABC):
         pass
 
     @abstractmethod
-    def get_indexed_ids(self, vectorstore_wrapper, collection_suffix: Optional[str] = '') -> List[str]:
+    def get_indexed_ids(self, vectorstore_wrapper, index_name: Optional[str] = '') -> List[str]:
         """Get all indexed document IDs from vectorstore"""
         pass
 
     @abstractmethod
-    def clean_collection(self, vectorstore_wrapper, collection_suffix: str = ''):
-        """Clean the vectorstore collection by deleting all indexed data."""
+    def clean_collection(self, vectorstore_wrapper, index_name: str = '', including_index_meta: bool = False):
+        """Clean the vectorstore collection by deleting all indexed data. If including_index_meta is True, skip the index_meta records."""
         pass
 
     @abstractmethod
@@ -39,7 +41,7 @@ class VectorStoreAdapter(ABC):
         pass
 
     @abstractmethod
-    def get_code_indexed_data(self, vectorstore_wrapper, collection_suffix) -> Dict[str, Dict[str, Any]]:
+    def get_code_indexed_data(self, vectorstore_wrapper, index_name) -> Dict[str, Dict[str, Any]]:
         """Get all indexed data from vectorstore for code content"""
         pass
 
@@ -48,15 +50,26 @@ class VectorStoreAdapter(ABC):
         """Add a new collection name to the metadata"""
         pass
 
+    @abstractmethod
+    def get_index_meta(self, vectorstore_wrapper, index_name: str) -> List[Dict[str, Any]]:
+        """Get all index_meta entries 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]:
+        try:
+            from tools import this  # pylint: disable=E0401,C0415
+            worker_config = this.for_module("indexer_worker").descriptor.config
+        except:  # pylint: disable=W0702
+            worker_config = {}
+        #
         return {
             "use_jsonb": True,
             "collection_name": collection_name,
-            "create_extension": True,
+            "create_extension": worker_config.get("pgvector_create_extension", True),
             "alita_sdk_options": {
                 "target_schema": collection_name,
             },
@@ -93,20 +106,25 @@ class PGVectorAdapter(VectorStoreAdapter):
             session.commit()
             logger.info(f"Schema '{schema_name}' has been dropped.")
 
-    def get_indexed_ids(self, vectorstore_wrapper, collection_suffix: Optional[str] = '') -> List[str]:
+    def get_indexed_ids(self, vectorstore_wrapper, index_name: Optional[str] = '') -> List[str]:
         """Get all indexed document IDs from PGVector"""
         from sqlalchemy.orm import Session
-        from sqlalchemy import func
+        from sqlalchemy import func, or_
 
         store = vectorstore_wrapper.vectorstore
         try:
             with Session(store.session_maker.bind) as session:
                 # Start building the query
                 query = session.query(store.EmbeddingStore.id)
-                # Apply filter only if collection_suffix is provided
-                if collection_suffix:
+                # Apply filter only if index_name is provided
+                if index_name:
                     query = query.filter(
-                        func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'collection') == collection_suffix
+                        func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'collection') == index_name,
+                        or_(
+                            func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'type').is_(None),
+                            func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata,
+                                                         'type') != IndexerKeywords.INDEX_META_TYPE.value
+                        )
                     )
                 ids = query.all()
                 return [str(id_tuple[0]) for id_tuple in ids]
@@ -114,25 +132,37 @@ class PGVectorAdapter(VectorStoreAdapter):
             logger.error(f"Failed to get indexed IDs from PGVector: {str(e)}")
             return []
 
-    def clean_collection(self, vectorstore_wrapper, collection_suffix: str = ''):
-        """Clean the vectorstore collection by deleting all indexed data."""
-        # This logic deletes all data from the vectorstore collection without removal of collection.
-        # Collection itself remains available for future indexing.
-        vectorstore_wrapper.vectorstore.delete(ids=self.get_indexed_ids(vectorstore_wrapper, collection_suffix))
+    def clean_collection(self, vectorstore_wrapper, index_name: str = '', including_index_meta: bool = False):
+        """Clean the vectorstore collection by deleting all indexed data. If including_index_meta is True, skip the index_meta records."""
+        from sqlalchemy.orm import Session
+        from sqlalchemy import func, or_
+        store = vectorstore_wrapper.vectorstore
+        with Session(store.session_maker.bind) as session:
+            if including_index_meta:
+                session.query(store.EmbeddingStore).filter(
+                    func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'collection') == index_name
+                ).delete(synchronize_session=False)
+            else:
+                session.query(store.EmbeddingStore).filter(
+                    func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'collection') == index_name,
+                    or_(func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'type').is_(None),
+                        func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'type') != IndexerKeywords.INDEX_META_TYPE.value)
+                ).delete(synchronize_session=False)
+            session.commit()
 
     def is_vectorstore_type(self, vectorstore) -> bool:
         """Check if the vectorstore is a PGVector store."""
         return hasattr(vectorstore, 'session_maker') and hasattr(vectorstore, 'EmbeddingStore')
 
-    def get_indexed_data(self, vectorstore_wrapper, collection_suffix: str)-> Dict[str, Dict[str, Any]]:
-        """Get all indexed data from PGVector for non-code content per collection_suffix."""
+    def get_indexed_data(self, vectorstore_wrapper, index_name: str)-> Dict[str, Dict[str, Any]]:
+        """Get all indexed data from PGVector for non-code content per index_name."""
         from sqlalchemy.orm import Session
         from sqlalchemy import func
         from ...runtime.utils.utils import IndexerKeywords
 
         result = {}
         try:
-            vectorstore_wrapper._log_data("Retrieving already indexed data from PGVector vectorstore",
+            vectorstore_wrapper._log_tool_event("Retrieving already indexed data from PGVector vectorstore",
                            tool_name="get_indexed_data")
             store = vectorstore_wrapper.vectorstore
             with Session(store.session_maker.bind) as session:
@@ -141,7 +171,7 @@ class PGVectorAdapter(VectorStoreAdapter):
                     store.EmbeddingStore.document,
                     store.EmbeddingStore.cmetadata
                 ).filter(
-                    func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'collection') == collection_suffix
+                    func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'collection') == index_name
                 ).all()
 
             # Process the retrieved data
@@ -174,14 +204,14 @@ class PGVectorAdapter(VectorStoreAdapter):
 
         return result
 
-    def get_code_indexed_data(self, vectorstore_wrapper, collection_suffix: str) -> Dict[str, Dict[str, Any]]:
+    def get_code_indexed_data(self, vectorstore_wrapper, index_name: str) -> Dict[str, Dict[str, Any]]:
         """Get all indexed code data from PGVector per collection suffix."""
         from sqlalchemy.orm import Session
         from sqlalchemy import func
 
         result = {}
         try:
-            vectorstore_wrapper._log_data("Retrieving already indexed code data from PGVector vectorstore",
+            vectorstore_wrapper._log_tool_event(message="Retrieving already indexed code data from PGVector vectorstore",
                            tool_name="index_code_data")
             store = vectorstore_wrapper.vectorstore
             with (Session(store.session_maker.bind) as session):
@@ -189,7 +219,7 @@ class PGVectorAdapter(VectorStoreAdapter):
                     store.EmbeddingStore.id,
                     store.EmbeddingStore.cmetadata
                 ).filter(
-                    func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'collection') == collection_suffix
+                    func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'collection') == index_name
                 ).all()
 
             for db_id, meta in docs:
@@ -259,6 +289,29 @@ class PGVectorAdapter(VectorStoreAdapter):
         except Exception as e:
             logger.error(f"Failed to update collection for entry ID {entry_id}: {str(e)}")
 
+    def get_index_meta(self, vectorstore_wrapper, index_name: str) -> List[Dict[str, Any]]:
+        from sqlalchemy.orm import Session
+        from sqlalchemy import func
+
+        store = vectorstore_wrapper.vectorstore
+        try:
+            with Session(store.session_maker.bind) as session:
+                meta = session.query(
+                    store.EmbeddingStore.id,
+                    store.EmbeddingStore.document,
+                    store.EmbeddingStore.cmetadata
+                ).filter(
+                    store.EmbeddingStore.cmetadata['type'].astext == IndexerKeywords.INDEX_META_TYPE.value,
+                    func.jsonb_extract_path_text(store.EmbeddingStore.cmetadata, 'collection') == index_name
+                ).all()
+                result = []
+                for id, document, cmetadata in meta:
+                    result.append({"id": id, "content": document, "metadata": cmetadata})
+                return result
+        except Exception as e:
+            logger.error(f"Failed to get index_meta from PGVector: {str(e)}")
+            raise e
+
 
 class ChromaAdapter(VectorStoreAdapter):
     """Adapter for Chroma database operations."""
@@ -276,7 +329,7 @@ class ChromaAdapter(VectorStoreAdapter):
     def remove_collection(self, vectorstore_wrapper, collection_name: str):
         vectorstore_wrapper.vectorstore.delete_collection()
 
-    def get_indexed_ids(self, vectorstore_wrapper, collection_suffix: Optional[str] = '') -> List[str]:
+    def get_indexed_ids(self, vectorstore_wrapper, index_name: Optional[str] = '') -> List[str]:
         """Get all indexed document IDs from Chroma"""
         try:
             data = vectorstore_wrapper.vectorstore.get(include=[])  # Only get IDs, no metadata
@@ -285,9 +338,9 @@ class ChromaAdapter(VectorStoreAdapter):
             logger.error(f"Failed to get indexed IDs from Chroma: {str(e)}")
             return []
 
-    def clean_collection(self, vectorstore_wrapper, collection_suffix: str = ''):
-        """Clean the vectorstore collection by deleting all indexed data."""
-        vectorstore_wrapper.vectorstore.delete(ids=self.get_indexed_ids(vectorstore_wrapper, collection_suffix))
+    def clean_collection(self, vectorstore_wrapper, index_name: str = '', including_index_meta: bool = False):
+        """Clean the vectorstore collection by deleting all indexed data. including_index_meta is ignored."""
+        vectorstore_wrapper.vectorstore.delete(ids=self.get_indexed_ids(vectorstore_wrapper, index_name))
 
     def get_indexed_data(self, vectorstore_wrapper):
         """Get all indexed data from Chroma for non-code content"""
@@ -325,7 +378,7 @@ class ChromaAdapter(VectorStoreAdapter):
 
         return result
 
-    def get_code_indexed_data(self, vectorstore_wrapper, collection_suffix) -> Dict[str, Dict[str, Any]]:
+    def get_code_indexed_data(self, vectorstore_wrapper, index_name) -> Dict[str, Dict[str, Any]]:
         """Get all indexed code data from Chroma."""
         result = {}
         try:
@@ -355,6 +408,9 @@ class ChromaAdapter(VectorStoreAdapter):
         # This is a simplified implementation - in practice, you might need more complex logic
         logger.warning("add_to_collection for Chroma is not fully implemented yet")
 
+    def get_index_meta(self, vectorstore_wrapper, index_name: str) -> List[Dict[str, Any]]:
+        logger.warning("get_index_meta for Chroma is not implemented yet")
+
 
 class VectorStoreAdapterFactory:
     """Factory for creating vector store adapters."""

alita_sdk/tools/xray/__init__.py

@@ -7,7 +7,8 @@ from pydantic import create_model, BaseModel, Field
 
 from .api_wrapper import XrayApiWrapper
 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.xray import XrayConfiguration
 
@@ -17,34 +18,30 @@ 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'),
+        llm=tool['settings'].get('llm', None),
         alita=tool['settings'].get('alita', None),
 
         # indexer settings
         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()
 
 
 class XrayToolkit(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 XrayApiWrapper.model_construct().get_available_tools()}
-        XrayToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
         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={
@@ -57,7 +54,7 @@ class XrayToolkit(BaseToolkit):
                             {
                                 'metadata': {
                                     "label": "XRAY cloud", "icon_url": "xray.svg",
-                                "categories": ["test management"],
+                                    "categories": ["test management"],
                                     "extra_categories": ["test automation", "test case management", "test planning"]
                                 }
                             }
@@ -65,28 +62,34 @@ class XrayToolkit(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,
             # Use xray_configuration fields
-            **kwargs.get('xray_configuration', {}),
+            **(kwargs.get('xray_configuration') or {}),
             **(kwargs.get('pgvector_configuration') or {}),
         }
         xray_api_wrapper = XrayApiWrapper(**wrapper_payload)
-        prefix = clean_string(toolkit_name, cls.toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
         available_tools = xray_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 + "\nXray instance: " + xray_api_wrapper.base_url
+            description = description[:1000]
             tools.append(BaseAction(
                 api_wrapper=xray_api_wrapper,
-                name=prefix + tool["name"],
-                description=tool["description"] + "\nXray instance: " + xray_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)