alita-sdk 0.3.230__tar.gz → 0.3.232__tar.gz

{alita_sdk-0.3.230 → alita_sdk-0.3.232}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.230
+Version: 0.3.232
 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.230 → alita_sdk-0.3.232}/alita_sdk/runtime/tools/vectorstore.py

@@ -8,6 +8,7 @@ from pydantic import BaseModel, model_validator, Field
 from ..langchain.tools.vector import VectorAdapter
 from langchain_core.messages import HumanMessage
 from alita_sdk.tools.elitea_base import BaseToolApiWrapper
+from alita_sdk.tools.vector_adapters.VectorStoreAdapter import VectorStoreAdapterFactory
 from logging import getLogger
 
 from ..utils.logging import dispatch_custom_event
@@ -141,11 +142,14 @@ class VectorStoreWrapper(BaseToolApiWrapper):
     dataset: str = None
     embedding: Any = None
     vectorstore: Any = None
+    # Review usage of old adapter
     vectoradapter: Any = None
     pg_helper: Any = None
     embeddings: Any = None
     process_document_func: Optional[Callable] = None
-    
+    # New adapter for vector database operations
+    vector_adapter: Any = None
+
     @model_validator(mode='before')
     @classmethod
     def validate_toolkit(cls, values):
@@ -170,6 +174,8 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             embeddings=values['embeddings'],
             quota_params=None,
         )
+        # Initialize the new vector adapter
+        values['vector_adapter'] = VectorStoreAdapterFactory.create_adapter(values['vectorstore_type'])
         logger.debug(f"Vectorstore wrapper initialized: {values}")
         return values
 
@@ -196,52 +202,22 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             f"Remove collection '{self.dataset}'",
             tool_name="_remove_collection"
         )
-        from sqlalchemy import text
-        from sqlalchemy.orm import Session
-
-        schema_name = self.vectorstore.collection_name
-        with Session(self.vectorstore.session_maker.bind) as session:
-            drop_schema_query = text(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE;")
-            session.execute(drop_schema_query)
-            session.commit()
-            logger.info(f"Schema '{schema_name}' has been dropped.")
+        self.vector_adapter.remove_collection(self, self.dataset)
         self._log_data(
             f"Collection '{self.dataset}' has been removed. ",
             tool_name="_remove_collection"
         )
 
-    def _get_indexed_ids(self, store):
+    def _get_indexed_ids(self, collection_suffix: Optional[str] = '') -> List[str]:
         """Get all indexed document IDs from vectorstore"""
+        return self.vector_adapter.get_indexed_ids(self, collection_suffix)
 
-        # Check if this is a PGVector store
-        if hasattr(store, 'session_maker') and hasattr(store, 'EmbeddingStore'):
-            return self._get_pgvector_indexed_ids(store)
-        else:
-            # Fall back to Chroma implementation
-            return self._get_chroma_indexed_ids(store)
-
-    def _get_pgvector_indexed_ids(self, store):
-        """Get all indexed document IDs from PGVector"""
-        from sqlalchemy.orm import Session
-
-        try:
-            with Session(store.session_maker.bind) as session:
-                ids = session.query(store.EmbeddingStore.id).all()
-            return [str(id_tuple[0]) for id_tuple in ids]
-        except Exception as e:
-            logger.error(f"Failed to get indexed IDs from PGVector: {str(e)}")
-            return []
+    def list_collections(self) -> List[str]:
+        """List all collections in the vectorstore."""
 
-    def _get_chroma_indexed_ids(self, store):
-        """Get all indexed document IDs from Chroma"""
-        try:
-            data = store.get(include=[])  # Only get IDs, no metadata
-            return data.get('ids', [])
-        except Exception as e:
-            logger.error(f"Failed to get indexed IDs from Chroma: {str(e)}")
-            return []
+        return self.vector_adapter.list_collections(self)
 
-    def _clean_collection(self):
+    def _clean_collection(self, collection_suffix: str = ''):
         """
         Clean the vectorstore collection by deleting all indexed data.
         """
@@ -249,135 +225,28 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             f"Cleaning collection '{self.dataset}'",
             tool_name="_clean_collection"
         )
-        # This logic deletes all data from the vectorstore collection without removal of collection.
-        # Collection itself remains available for future indexing.
-        self.vectoradapter.vectorstore.delete(ids=self._get_indexed_ids(self.vectoradapter.vectorstore))
-
+        self.vector_adapter.clean_collection(self, collection_suffix)
         self._log_data(
             f"Collection '{self.dataset}' has been cleaned. ",
             tool_name="_clean_collection"
         )
 
-    # TODO: refactor to use common method for different vectorstores in a separate vectorstore wrappers
-    def _get_indexed_data(self, store):
+    def _get_indexed_data(self):
         """ Get all indexed data from vectorstore for non-code content """
+        return self.vector_adapter.get_indexed_data(self)
 
-        # Check if this is a PGVector store
-        if hasattr(store, 'session_maker') and hasattr(store, 'EmbeddingStore'):
-            return self._get_pgvector_indexed_data(store)
-        else:
-            # Fall back to original Chroma implementation
-            return self._get_chroma_indexed_data(store)
-
-    def _get_pgvector_indexed_data(self, store):
-        """ Get all indexed data from PGVector for non-code content """
-        from sqlalchemy.orm import Session
-
-        result = {}
-        try:
-            self._log_data("Retrieving already indexed data from PGVector vectorstore",
-                           tool_name="index_documents")
-
-            with Session(store.session_maker.bind) as session:
-                docs = session.query(
-                    store.EmbeddingStore.id,
-                    store.EmbeddingStore.document,
-                    store.EmbeddingStore.cmetadata
-                ).all()
-
-            # Process the retrieved data
-            for doc in docs:
-                db_id = doc.id
-                meta = doc.cmetadata or {}
-
-                # Get document id from metadata
-                doc_id = str(meta.get('id', db_id))
-                dependent_docs = meta.get(IndexerKeywords.DEPENDENT_DOCS.value, [])
-                if dependent_docs:
-                    dependent_docs = [d.strip() for d in dependent_docs.split(';') if d.strip()]
-                parent_id = meta.get(IndexerKeywords.PARENT.value, -1)
-
-                chunk_id = meta.get('chunk_id')
-                if doc_id in result and chunk_id:
-                    # If document with the same id already saved, add db_id for current one as chunk
-                    result[doc_id]['all_chunks'].append(db_id)
-                else:
-                    result[doc_id] = {
-                        'metadata': meta,
-                        'id': db_id,
-                        'all_chunks': [db_id],
-                        IndexerKeywords.DEPENDENT_DOCS.value: dependent_docs,
-                        IndexerKeywords.PARENT.value: parent_id
-                    }
-
-        except Exception as e:
-            logger.error(f"Failed to get indexed data from PGVector: {str(e)}. Continuing with empty index.")
-
-        return result
-
-    def _get_chroma_indexed_data(self, store):
-        """ Get all indexed data from Chroma for non-code content """
-        result = {}
-        try:
-            self._log_data("Retrieving already indexed data from Chroma vectorstore",
-                           tool_name="index_documents")
-            data = store.get(include=['metadatas'])
-
-            # Re-structure data to be more usable
-            for meta, db_id in zip(data['metadatas'], data['ids']):
-                # Get document id from metadata
-                doc_id = str(meta['id'])
-                dependent_docs = meta.get(IndexerKeywords.DEPENDENT_DOCS.value, [])
-                if dependent_docs:
-                    dependent_docs = [d.strip() for d in dependent_docs.split(';') if d.strip()]
-                parent_id = meta.get(IndexerKeywords.PARENT.value, -1)
-
-                chunk_id = meta.get('chunk_id')
-                if doc_id in result and chunk_id:
-                    # If document with the same id already saved, add db_id for current one as chunk
-                    result[doc_id]['all_chunks'].append(db_id)
-                else:
-                    result[doc_id] = {
-                        'metadata': meta,
-                        'id': db_id,
-                        'all_chunks': [db_id],
-                        IndexerKeywords.DEPENDENT_DOCS.value: dependent_docs,
-                        IndexerKeywords.PARENT.value: parent_id
-                    }
-        except Exception as e:
-            logger.error(f"Failed to get indexed data from Chroma: {str(e)}. Continuing with empty index.")
-
-        return result
-
-    def _get_code_indexed_data(self, store) -> Dict[str, Dict[str, Any]]:
+    def _get_code_indexed_data(self) -> Dict[str, Dict[str, Any]]:
         """ Get all indexed data from vectorstore for code content """
+        return self.vector_adapter.get_code_indexed_data(self)
 
-        # get already indexed data
-        result = {}
-        try:
-            self._log_data("Retrieving already indexed code data from vectorstore",
-                           tool_name="index_documents")
-            data = store.get(include=['metadatas'])
-            # re-structure data to be more usable
-            for meta, db_id in zip(data['metadatas'], data['ids']):
-                filename = meta['filename']
-                commit_hash = meta.get('commit_hash')
-                if filename not in result:
-                    result[filename] = {
-                        'commit_hashes': [],
-                        'ids': []
-                    }
-                if commit_hash is not None:
-                    result[filename]['commit_hashes'].append(commit_hash)
-                result[filename]['ids'].append(db_id)
-        except Exception as e:
-            logger.error(f"Failed to get indexed code data from vectorstore: {str(e)}. Continuing with empty index.")
-        return result
+    def _add_to_collection(self, entry_id, new_collection_value):
+        """Add a new collection name to the `collection` key in the `metadata` column."""
+        self.vector_adapter.add_to_collection(self, entry_id, new_collection_value)
 
     def _reduce_duplicates(
             self,
             documents: Generator[Any, None, None],
-            store,
+            collection_suffix: str,
             get_indexed_data: Callable,
             key_fn: Callable,
             compare_fn: Callable,
@@ -386,7 +255,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
     ) -> List[Any]:
         """Generic duplicate reduction logic for documents."""
         self._log_data(log_msg, tool_name="index_documents")
-        indexed_data = get_indexed_data(store)
+        indexed_data = get_indexed_data()
         indexed_keys = set(indexed_data.keys())
         if not indexed_keys:
             self._log_data("Vectorstore is empty, indexing all incoming documents", tool_name="index_documents")
@@ -397,8 +266,15 @@ class VectorStoreWrapper(BaseToolApiWrapper):
 
         for document in documents:
             key = key_fn(document)
-            if key in indexed_keys:
+            if key in indexed_keys and collection_suffix == indexed_data[key]['metadata'].get('collection'):
                 if compare_fn(document, indexed_data[key]):
+                    # Disabled addition of new collection to already indexed documents
+                    # # check metadata.collection and update if needed
+                    # for update_collection_id in remove_ids_fn(indexed_data, key):
+                    #     self._add_to_collection(
+                    #         update_collection_id,
+                    #         collection_suffix
+                    #     )
                     continue
                 final_docs.append(document)
                 docs_to_remove.update(remove_ids_fn(indexed_data, key))
@@ -410,14 +286,14 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 f"Removing {len(docs_to_remove)} documents from vectorstore that are already indexed with different updated_on.",
                 tool_name="index_documents"
             )
-            store.delete(ids=list(docs_to_remove))
+            self.vectorstore.delete(ids=list(docs_to_remove))
 
         return final_docs
 
-    def _reduce_non_code_duplicates(self, documents: Generator[Any, None, None], store) -> List[Any]:
+    def _reduce_non_code_duplicates(self, documents: Generator[Any, None, None], collection_suffix: str) -> List[Any]:
         return self._reduce_duplicates(
             documents,
-            store,
+            collection_suffix,
             self._get_indexed_data,
             lambda doc: doc.metadata.get('id'),
             lambda doc, idx: (
@@ -434,10 +310,10 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             log_msg="Verification of documents to index started"
         )
 
-    def _reduce_code_duplicates(self, documents: Generator[Any, None, None], store) -> List[Any]:
+    def _reduce_code_duplicates(self, documents: Generator[Any, None, None], collection_suffix: str) -> List[Any]:
         return self._reduce_duplicates(
             documents,
-            store,
+            collection_suffix,
             self._get_code_indexed_data,
             lambda doc: doc.metadata.get('filename'),
             lambda doc, idx: (
@@ -449,7 +325,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             log_msg="Verification of code documents to index started"
         )
 
-    def index_documents(self, documents: Generator[Document, None, None], progress_step: int = 20, clean_index: bool = True, is_code: bool = False):
+    def index_documents(self, documents: Generator[Document, None, None], collection_suffix: str, progress_step: int = 20, clean_index: bool = True, is_code: bool = False):
         """ Index documents in the vectorstore.
 
         Args:
@@ -465,7 +341,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             logger.info("Cleaning index before re-indexing all documents.")
             self._log_data("Cleaning index before re-indexing all documents. Previous index will be removed", tool_name="index_documents")
             try:
-                self._clean_collection()
+                self._clean_collection(collection_suffix)
                 self.vectoradapter.persist()
                 self.vectoradapter.vacuum()
                 self._log_data("Previous index has been removed",
@@ -476,8 +352,8 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 documents = list(documents)
         else:
             # remove duplicates based on metadata 'id' and 'updated_on' or 'commit_hash' fields
-            documents = self._reduce_code_duplicates(documents, self.vectoradapter.vectorstore) if is_code \
-                else self._reduce_non_code_duplicates(documents, self.vectoradapter.vectorstore)
+            documents = self._reduce_code_duplicates(documents, collection_suffix) if is_code \
+                else self._reduce_non_code_duplicates(documents, collection_suffix)
 
         if not documents or len(documents) == 0:
             logger.info("No new documents to index after duplicate check.")
@@ -498,6 +374,15 @@ class VectorStoreWrapper(BaseToolApiWrapper):
         logger.debug(self.vectoradapter)
 
         documents = documents + list(dependent_docs_generator)
+
+        # if collection_suffix is provided, add it to metadata of each document
+        if collection_suffix:
+            for doc in documents:
+                if not doc.metadata.get('collection'):
+                    doc.metadata['collection'] = collection_suffix
+                else:
+                    doc.metadata['collection'] += f";{collection_suffix}"
+
         total_docs = len(documents)
         documents_count = 0
         _documents = []
@@ -511,8 +396,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             try:
                 _documents.append(document)
                 if len(_documents) >= self.max_docs_per_add:
-                    add_documents(vectorstore=self.vectoradapter.vectorstore, documents=_documents)
-                    self.vectoradapter.persist()
+                    add_documents(vectorstore=self.vectorstore, documents=_documents)
                     _documents = []
 
                 percent = math.floor((documents_count / total_docs) * 100)
@@ -526,8 +410,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 logger.error(f"Error: {format_exc()}")
                 return {"status": "error", "message": f"Error: {format_exc()}"}
         if _documents:
-            add_documents(vectorstore=self.vectoradapter.vectorstore, documents=_documents)
-            self.vectoradapter.persist()
+            add_documents(vectorstore=self.vectorstore, documents=_documents)
         return {"status": "ok", "message": f"successfully indexed {documents_count} documents"}
 
     def search_documents(self, query:str, doctype: str = 'code', 
@@ -562,7 +445,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 }
                 
             try:
-                document_items = self.vectoradapter.vectorstore.similarity_search_with_score(
+                document_items = self.vectorstore.similarity_search_with_score(
                     query, filter=document_filter, k=search_top
                 )                
                 # Add document results to unique docs
@@ -595,7 +478,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                     }
                 
                 try:
-                    chunk_items = self.vectoradapter.vectorstore.similarity_search_with_score(
+                    chunk_items = self.vectorstore.similarity_search_with_score(
                         query, filter=chunk_filter, k=search_top
                     )
                     
@@ -628,7 +511,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                                 }
                                 
                             try:
-                                fetch_items = self.vectoradapter.vectorstore.similarity_search_with_score(
+                                fetch_items = self.vectorstore.similarity_search_with_score(
                                     query, filter=doc_filter, k=1
                                 )
                                 if fetch_items:
@@ -642,7 +525,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
         else:
             # Default search behavior (unchanged)
             max_search_results = 30 if search_top * 3 > 30 else search_top * 3
-            vector_items = self.vectoradapter.vectorstore.similarity_search_with_score(
+            vector_items = self.vectorstore.similarity_search_with_score(
                 query, filter=filter, k=max_search_results
             )
             
@@ -862,4 +745,5 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 "description": "Get summary of search results using stepback technique",
                 "args_schema": StepBackSearchDocumentsModel
             }
-        ]
+        ]
+

{alita_sdk-0.3.230 → alita_sdk-0.3.232}/alita_sdk/tools/elitea_base.py

@@ -30,13 +30,13 @@ LoaderSchema = create_model(
 # Base Vector Store Schema Models
 BaseIndexParams = create_model(
     "BaseIndexParams",
-    collection_suffix=(Optional[str], Field(description="Optional suffix for collection name (max 7 characters)", default="", max_length=7)),
+    collection_suffix=(str, Field(description="Suffix for collection name (max 7 characters) used to separate datasets", min_length=1, max_length=7)),
     vectorstore_type=(Optional[str], Field(description="Vectorstore type (Chroma, PGVector, Elastic, etc.)", default="PGVector")),
 )
 
 BaseCodeIndexParams = create_model(
     "BaseCodeIndexParams",
-    collection_suffix=(Optional[str], Field(description="Optional suffix for collection name (max 7 characters)", default="", max_length=7)),
+    collection_suffix=(str, Field(description="Suffix for collection name (max 7 characters) used to separate datasets", min_length=1, max_length=7)),
     vectorstore_type=(Optional[str], Field(description="Vectorstore type (Chroma, PGVector, Elastic, etc.)", default="PGVector")),
     branch=(Optional[str], Field(description="Branch to index files from. Defaults to active branch if None.", default=None)),
     whitelist=(Optional[List[str]], Field(description="File extensions or paths to include. Defaults to all files if None.", default=None)),
@@ -51,7 +51,9 @@ RemoveIndexParams = create_model(
 BaseSearchParams = create_model(
     "BaseSearchParams",
     query=(str, Field(description="Query text to search in the index")),
-    collection_suffix=(Optional[str], Field(description="Optional suffix for collection name (max 7 characters)", default="", max_length=7)),
+    collection_suffix=(Optional[str], Field(
+        description="Optional suffix for collection name (max 7 characters). Leave empty to search across all datasets",
+        default="", max_length=7)),
     vectorstore_type=(Optional[str], Field(description="Vectorstore type (Chroma, PGVector, Elastic, etc.)", default="PGVector")),
     filter=(Optional[dict | str], Field(
         description="Filter to apply to the search results. Can be a dictionary or a JSON string.",
@@ -219,6 +221,7 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
     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"
+    _vector_store: Optional[Any] = None
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -297,9 +300,9 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
         collection_suffix = kwargs.get("collection_suffix")
         progress_step = kwargs.get("progress_step")
         clean_index = kwargs.get("clean_index")
-        vs = self._init_vector_store(collection_suffix, embeddings=embedding)
+        vs = self._init_vector_store(embeddings=embedding)
         #
-        return vs.index_documents(docs, progress_step=progress_step, clean_index=clean_index)
+        return vs.index_documents(docs, collection_suffix=collection_suffix, progress_step=progress_step, clean_index=clean_index)
 
     def _process_documents(self, documents: List[Document]) -> Generator[Document, None, None]:
         """
@@ -333,48 +336,37 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
 
 
     # TODO: init store once and re-use the instance
-    def _init_vector_store(self, collection_suffix: str = "", embeddings: Optional[Any] = None):
+    def _init_vector_store(self, embeddings: Optional[Any] = None):
         """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()}"
-
-        # 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,
-            vectorstore_type=self.vectorstore_type,
-            embedding_model=self.embedding_model,
-            embedding_model_params=self.embedding_model_params,
-            vectorstore_params=vectorstore_params,
-            embeddings=embeddings,
-            process_document_func=self._process_documents,
-        )
+        if not self._vector_store:
+            connection_string = self.connection_string.get_secret_value() if self.connection_string else None
+            vectorstore_params = self._adapter.get_vectorstore_params(self.collection_name, connection_string)
+            self._vector_store = VectorStoreWrapper(
+                llm=self.llm,
+                vectorstore_type=self.vectorstore_type,
+                embedding_model=self.embedding_model,
+                embedding_model_params=self.embedding_model_params,
+                vectorstore_params=vectorstore_params,
+                embeddings=embeddings,
+                process_document_func=self._process_documents,
+            )
+        return self._vector_store
 
     def remove_index(self, collection_suffix: str = ""):
         """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"
+        self._init_vector_store()._clean_collection(collection_suffix=collection_suffix)
+        return (f"Collection '{collection_suffix}' 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."""
         vectorstore_wrapper = self._init_vector_store()
-        return self._adapter.list_collections(vectorstore_wrapper, self.collection_name or "")
+        return vectorstore_wrapper.list_collections()
 
     def search_index(self,
                      query: str,
@@ -386,7 +378,14 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
                      extended_search: Optional[List[str]] = None,
                      **kwargs):
         """ Searches indexed documents in the vector store."""
-        vectorstore = self._init_vector_store(collection_suffix)
+        vectorstore = self._init_vector_store()
+        # build filter on top of collection_suffix
+        filter = filter if isinstance(filter, dict) else json.loads(filter)
+        if collection_suffix:
+            filter.update({"collection": {
+                "$eq": collection_suffix.strip()
+            }})
+
         found_docs = vectorstore.search_documents(
             query,
             doctype=self.doctype,
@@ -499,6 +498,8 @@ class BaseVectorStoreToolApiWrapper(BaseToolApiWrapper):
 
 class BaseCodeToolApiWrapper(BaseVectorStoreToolApiWrapper):
 
+    doctype: Optional[str] = 'code'
+
     def _get_files(self):
         raise NotImplementedError("Subclasses should implement this method")
 
@@ -579,22 +580,20 @@ class BaseCodeToolApiWrapper(BaseVectorStoreToolApiWrapper):
         return parse_code_files_for_db(file_content_generator())
     
     def index_data(self,
+                   collection_suffix: str,
                    branch: Optional[str] = None,
                    whitelist: Optional[List[str]] = None,
                    blacklist: Optional[List[str]] = None,
-                   collection_suffix: str = "",
                    **kwargs) -> str:
         """Index repository files in the vector store using code parsing."""
         
-        
-        
         documents = self.loader(
             branch=branch,
             whitelist=whitelist,
             blacklist=blacklist
         )
-        vectorstore = self._init_vector_store(collection_suffix)
-        return vectorstore.index_documents(documents, clean_index=False, is_code=True)
+        vectorstore = self._init_vector_store()
+        return vectorstore.index_documents(documents, collection_suffix=collection_suffix, clean_index=False, is_code=True)
 
     def _get_vector_search_tools(self):
         """

{alita_sdk-0.3.230 → alita_sdk-0.3.232}/alita_sdk/tools/github/api_wrapper.py

@@ -37,9 +37,9 @@ class AlitaGitHubAPIWrapper(BaseCodeToolApiWrapper):
     Wrapper for GitHub API that integrates both REST and GraphQL functionality.
     """
     # Authentication config
-    github_access_token: Optional[str] = None
-    github_username: Optional[str] = None
-    github_password: Optional[str] = None
+    github_access_token: Optional[SecretStr] = None
+    github_username: Optional[SecretStr] = None
+    github_password: Optional[SecretStr] = None
     github_app_id: Optional[str] = None
     github_app_private_key: Optional[str] = None
     github_base_url: Optional[str] = None
@@ -49,19 +49,9 @@ class AlitaGitHubAPIWrapper(BaseCodeToolApiWrapper):
     active_branch: Optional[str] = None
     github_base_branch: Optional[str] = None
 
-    # Add LLM instance
-    llm: Optional[Any] = None
     # Alita instance
     alita: Optional[Any] = None
     
-    # Vector store configuration
-    connection_string: Optional[SecretStr] = None
-    collection_name: Optional[str] = None
-    doctype: Optional[str] = 'code'  # GitHub uses 'code' doctype
-    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"
-
     # Client instances - renamed without leading underscores and marked as exclude=True
     github_client_instance: Optional[GitHubClient] = Field(default=None, exclude=True)
     graphql_client_instance: Optional[GraphQLClientWrapper] = Field(default=None, exclude=True)
@@ -84,12 +74,12 @@ class AlitaGitHubAPIWrapper(BaseCodeToolApiWrapper):
         from langchain.utils import get_from_dict_or_env
 
         # Get all authentication values
-        github_access_token = get_from_dict_or_env(values, "github_access_token", "GITHUB_ACCESS_TOKEN", default='')
-        github_username = get_from_dict_or_env(values, "github_username", "GITHUB_USERNAME", default='')
-        github_password = get_from_dict_or_env(values, "github_password", "GITHUB_PASSWORD", default='')
-        github_app_id = get_from_dict_or_env(values, "github_app_id", "GITHUB_APP_ID", default='')
-        github_app_private_key = get_from_dict_or_env(values, "github_app_private_key", "GITHUB_APP_PRIVATE_KEY", default='')
-        github_base_url = get_from_dict_or_env(values, "github_base_url", "GITHUB_BASE_URL", default='https://api.github.com')
+        github_access_token = get_from_dict_or_env(values, ["access_token", "github_access_token"], "GITHUB_ACCESS_TOKEN", default='')
+        github_username = get_from_dict_or_env(values, ["username", "github_username"], "GITHUB_USERNAME", default='')
+        github_password = get_from_dict_or_env(values, ["password", "github_password"], "GITHUB_PASSWORD", default='')
+        github_app_id = get_from_dict_or_env(values, ["app_id", "github_app_id"], "GITHUB_APP_ID", default='')
+        github_app_private_key = get_from_dict_or_env(values, ["app_private_key", "github_app_private_key"], "GITHUB_APP_PRIVATE_KEY", default='')
+        github_base_url = get_from_dict_or_env(values, ["base_url", "github_base_url"], "GITHUB_BASE_URL", default='https://api.github.com')
 
         auth_config = GitHubAuthConfig(
             github_access_token=github_access_token,