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

alita_sdk/runtime/tools/tool.py

@@ -86,7 +86,9 @@ Answer must be JSON only extractable by JSON.LOADS."""
         else:
             input_[-1].content += self.unstructured_output
             completion = self.client.invoke(input_, config=config)
-            result = _extract_json(completion.content.strip())
+            from ..langchain.utils import extract_text_from_completion
+            content_text = extract_text_from_completion(completion)
+            result = _extract_json(content_text.strip())
             logger.info(f"ToolNode tool params: {result}")
         try:
             # handler for application added as a tool

alita_sdk/runtime/tools/vectorstore.py

@@ -1,7 +1,7 @@
 import json
 import math
 import types
-from typing import Any, Optional, List, Dict, Callable, Generator
+from typing import Any, Optional, List, Dict, Callable, Generator, OrderedDict
 
 from langchain_core.documents import Document
 from pydantic import BaseModel, model_validator, Field
@@ -12,10 +12,11 @@ from alita_sdk.tools.vector_adapters.VectorStoreAdapter import VectorStoreAdapte
 from logging import getLogger
 
 from ..utils.logging import dispatch_custom_event
-from ..utils.utils import IndexerKeywords
+from ..langchain.utils import extract_text_from_completion
 
 logger = getLogger(__name__)
 
+
 class IndexDocumentsModel(BaseModel):
     documents: Any = Field(description="Generator of documents to index")
 
@@ -73,6 +74,10 @@ class StepBackSearchDocumentsModel(BaseModel):
         }""",
         default=None
     )
+    extended_search: Optional[List[str]] = Field(
+        description="List of chunk types to search for (title, summary, propositions, keywords, documents)",
+        default=None
+    )
     reranking_config: Optional[Dict[str, Dict[str, Any]]] = Field(
         description="""Reranking configuration. Example:
         {
@@ -87,10 +92,6 @@ class StepBackSearchDocumentsModel(BaseModel):
         }""",
         default=None
     )
-    extended_search: Optional[List[str]] = Field(
-        description="List of chunk types to search for (title, summary, propositions, keywords, documents)",
-        default=None
-    )
 
 STEPBACK_PROMPT = """Your task is to convert provided question into a more generic question that will be used for similarity search.
 Remove all not important words, question words, but save all names, dates and acronym as in original question.
@@ -138,7 +139,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
     embedding_model_params: dict
     vectorstore_type: str
     vectorstore_params: dict
-    max_docs_per_add: int = 100
+    max_docs_per_add: int = 20
     dataset: str = None
     embedding: Any = None
     vectorstore: Any = None
@@ -208,16 +209,33 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             tool_name="_remove_collection"
         )
 
-    def _get_indexed_ids(self, collection_suffix: Optional[str] = '') -> List[str]:
+    def _get_indexed_ids(self, index_name: Optional[str] = '') -> List[str]:
         """Get all indexed document IDs from vectorstore"""
-        return self.vector_adapter.get_indexed_ids(self, collection_suffix)
-
-    def list_collections(self) -> List[str]:
-        """List all collections in the vectorstore."""
-
-        return self.vector_adapter.list_collections(self)
+        return self.vector_adapter.get_indexed_ids(self, index_name)
+
+    def list_collections(self) -> Any:
+        """List all collections in the vectorstore.
+        Returns a list of collection names, or if no collections exist,
+        returns a dict with an empty list and a message."""
+        raw = self.vector_adapter.list_collections(self)
+        # Normalize raw result to a list of names
+        if not raw:
+            # No collections found
+            return {"collections": [], "message": "No indexed collections"}
+        if isinstance(raw, str):
+            # e.g., Chroma adapter returns comma-separated string
+            cols = [c for c in raw.split(',') if c]
+        else:
+            try:
+                cols = list(raw)
+            except Exception:
+                # Unexpected type, return raw directly
+                return raw
+        if not cols:
+            return {"collections": [], "message": "No indexed collections"}
+        return cols
 
-    def _clean_collection(self, collection_suffix: str = ''):
+    def _clean_collection(self, index_name: str = ''):
         """
         Clean the vectorstore collection by deleting all indexed data.
         """
@@ -225,19 +243,15 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             f"Cleaning collection '{self.dataset}'",
             tool_name="_clean_collection"
         )
-        self.vector_adapter.clean_collection(self, collection_suffix)
+        self.vector_adapter.clean_collection(self, index_name)
         self._log_data(
             f"Collection '{self.dataset}' has been cleaned. ",
             tool_name="_clean_collection"
         )
 
-    def _get_indexed_data(self, collection_name: str):
-        """ Get all indexed data from vectorstore for non-code content """
-        return self.vector_adapter.get_indexed_data(self, collection_name)
-
-    def _get_code_indexed_data(self, collection_suffix: str) -> Dict[str, Dict[str, Any]]:
+    def _get_code_indexed_data(self, index_name: str) -> Dict[str, Dict[str, Any]]:
         """ Get all indexed data from vectorstore for code content """
-        return self.vector_adapter.get_code_indexed_data(self, collection_suffix)
+        return self.vector_adapter.get_code_indexed_data(self, index_name)
 
     def _add_to_collection(self, entry_id, new_collection_value):
         """Add a new collection name to the `collection` key in the `metadata` column."""
@@ -246,7 +260,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
     def _reduce_duplicates(
             self,
             documents: Generator[Any, None, None],
-            collection_suffix: str,
+            index_name: str,
             get_indexed_data: Callable,
             key_fn: Callable,
             compare_fn: Callable,
@@ -255,7 +269,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(collection_suffix)
+        indexed_data = get_indexed_data(index_name)
         indexed_keys = set(indexed_data.keys())
         if not indexed_keys:
             self._log_data("Vectorstore is empty, indexing all incoming documents", tool_name="index_documents")
@@ -266,14 +280,15 @@ class VectorStoreWrapper(BaseToolApiWrapper):
 
         for document in documents:
             key = key_fn(document)
-            if key in indexed_keys and collection_suffix == indexed_data[key]['metadata'].get('collection'):
+            key = key if isinstance(key, str) else str(key)
+            if key in indexed_keys and index_name == 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
+                    #         index_name
                     #     )
                     continue
                 final_docs.append(document)
@@ -290,30 +305,10 @@ class VectorStoreWrapper(BaseToolApiWrapper):
 
         return final_docs
 
-    def _reduce_non_code_duplicates(self, documents: Generator[Any, None, None], collection_suffix: str) -> List[Any]:
-        return self._reduce_duplicates(
-            documents,
-            collection_suffix,
-            self._get_indexed_data,
-            lambda doc: doc.metadata.get('id'),
-            lambda doc, idx: (
-                    doc.metadata.get('updated_on') and
-                    idx['metadata'].get('updated_on') and
-                    doc.metadata.get('updated_on') == idx['metadata'].get('updated_on')
-            ),
-            lambda idx_data, key: (
-                    idx_data[key]['all_chunks'] +
-                    [idx_data[dep_id]['id'] for dep_id in idx_data[key][IndexerKeywords.DEPENDENT_DOCS.value]] +
-                    [chunk_db_id for dep_id in idx_data[key][IndexerKeywords.DEPENDENT_DOCS.value]
-                     for chunk_db_id in idx_data[dep_id]['all_chunks']]
-            ),
-            log_msg="Verification of documents to index started"
-        )
-
-    def _reduce_code_duplicates(self, documents: Generator[Any, None, None], collection_suffix: str) -> List[Any]:
+    def _reduce_code_duplicates(self, documents: Generator[Any, None, None], index_name: str) -> List[Any]:
         return self._reduce_duplicates(
             documents,
-            collection_suffix,
+            index_name,
             self._get_code_indexed_data,
             lambda doc: doc.metadata.get('filename'),
             lambda doc, idx: (
@@ -325,7 +320,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             log_msg="Verification of code documents to index started"
         )
 
-    def index_documents(self, documents: Generator[Document, None, None], collection_suffix: str, progress_step: int = 20, clean_index: bool = True, is_code: bool = False):
+    def index_documents(self, documents: Generator[Document, None, None], index_name: str, progress_step: int = 20, clean_index: bool = True, is_code: bool = True):
         """ Index documents in the vectorstore.
 
         Args:
@@ -336,13 +331,13 @@ class VectorStoreWrapper(BaseToolApiWrapper):
 
         from ..langchain.interfaces.llm_processor import add_documents
 
-        self._log_tool_event(message=f"Starting the indexing... Parameters: {collection_suffix=}, {clean_index=}, {is_code}", tool_name="index_documents")
+        self._log_tool_event(message=f"Starting the indexing... Parameters: {index_name=}, {clean_index=}, {is_code}", tool_name="index_documents")
         # pre-process documents if needed (find duplicates, etc.)
         if clean_index:
             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(collection_suffix)
+                self._clean_collection(index_name)
                 self.vectoradapter.persist()
                 self.vectoradapter.vacuum()
                 self._log_data("Previous index has been removed",
@@ -356,8 +351,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 message="Filter for duplicates",
                 tool_name="index_documents")
             # remove duplicates based on metadata 'id' and 'updated_on' or 'commit_hash' fields
-            documents = self._reduce_code_duplicates(documents, collection_suffix) if is_code \
-                else self._reduce_non_code_duplicates(documents, collection_suffix)
+            documents = self._reduce_code_duplicates(documents, index_name)
             self._log_tool_event(
                 message="All the duplicates were filtered out. Proceeding with indexing.",
                 tool_name="index_documents")
@@ -385,13 +379,13 @@ class VectorStoreWrapper(BaseToolApiWrapper):
         self._log_tool_event(message=f"Documents for indexing were processed. Total documents: {len(documents)}",
                              tool_name="index_documents")
 
-        # if collection_suffix is provided, add it to metadata of each document
-        if collection_suffix:
+        # if index_name is provided, add it to metadata of each document
+        if index_name:
             for doc in documents:
                 if not doc.metadata.get('collection'):
-                    doc.metadata['collection'] = collection_suffix
+                    doc.metadata['collection'] = index_name
                 else:
-                    doc.metadata['collection'] += f";{collection_suffix}"
+                    doc.metadata['collection'] += f";{index_name}"
 
         total_docs = len(documents)
         documents_count = 0
@@ -422,7 +416,8 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 return {"status": "error", "message": f"Error: {format_exc()}"}
         if _documents:
             add_documents(vectorstore=self.vectorstore, documents=_documents)
-        return {"status": "ok", "message": f"successfully indexed {documents_count} documents"}
+        return {"status": "ok", "message": f"successfully indexed {documents_count} documents" if documents_count > 0
+        else "No new documents to index."}
 
     def search_documents(self, query:str, doctype: str = 'code', 
                          filter:dict|str={}, cut_off: float=0.5,
@@ -542,11 +537,18 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             
         # Initialize document map for tracking by ID
         doc_map = {
-            f"{doc.metadata.get('id', f'idx_{i}')}_{doc.metadata['chunk_id']}"
-            if 'chunk_id' in doc.metadata
-            else doc.metadata.get('id', f"idx_{i}"): (doc, score)
+            (
+                f"{doc.metadata.get('id', f'idx_{i}')}_{doc.metadata['chunk_id']}"
+                if 'chunk_id' in doc.metadata
+                else doc.metadata.get('id', f"idx_{i}")
+            ): (doc, 1 - score)
             for i, (doc, score) in enumerate(vector_items)
         }
+
+        # Sort the items by the new score in descending order
+        doc_map = OrderedDict(
+            sorted(doc_map.items(), key=lambda x: x[1][1], reverse=True)
+        )
         
         # Process full-text search if configured
         if full_text_search and full_text_search.get('enabled') and full_text_search.get('fields'):
@@ -597,7 +599,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
         # Apply cutoff filter
         if cut_off:
             # Filter out items above the cutoff score (since the lower the score, the better)
-            combined_items = [item for item in combined_items if abs(item[1]) <= cut_off]
+            combined_items = [item for item in combined_items if abs(item[1]) >= cut_off]
         
         # Sort by score and limit results
         # DISABLED: for chroma we want ascending order (lower score is better), for others descending
@@ -684,8 +686,10 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 ]
             )
         ])
+        # Extract text content safely (handles both string and list content from thinking models)
+        search_query = extract_text_from_completion(result)
         search_results = self.search_documents(
-            result.content, doctype, filter, cut_off, search_top, 
+            search_query, doctype, filter, cut_off, search_top, 
             full_text_search=full_text_search,
             reranking_config=reranking_config,
             extended_search=extended_search
@@ -714,7 +718,8 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 ]
             )
         ])
-        return result.content
+        # Extract text content safely (handles both string and list content from thinking models)
+        return extract_text_from_completion(result)
 
     def _log_data(self, message: str, tool_name: str = "index_data"):
         """Log data and dispatch custom event for indexing progress"""
@@ -758,4 +763,3 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 "args_schema": StepBackSearchDocumentsModel
             }
         ]
-