diff --git a/src/vanna/chromadb/chromadb_vector.py b/src/vanna/chromadb/chromadb_vector.py index ed320463..668e8c1f 100644 --- a/src/vanna/chromadb/chromadb_vector.py +++ b/src/vanna/chromadb/chromadb_vector.py @@ -154,7 +154,8 @@ def remove_training_data(self, id: str, **kwargs) -> bool: return True else: return False - def reomove_collection(self, collection_name: str) -> bool: + + def remove_collection(self, collection_name: str) -> bool: """ This function can reset the collection to empty state. @@ -184,9 +185,18 @@ def reomove_collection(self, collection_name: str) -> bool: return True else: return False - # Static method to extract the documents from the results of a query + @staticmethod def _extract_documents(query_results) -> list: + """ + Static method to extract the documents from the results of a query. + + Args: + query_results (pd.DataFrame): The dataframe to use. + + Returns: + List[str] or None: The extracted documents, or an empty list or single document if an error occurred. + """ if query_results is None: return []