fix: add tests and adjustments to Chroma component (#5571)

* Refactor: Update EditNodeComponent to hide table options and block hide

* 🐛 (model.py): fix issue with search_documents method not returning empty list when search_query is empty
♻️ (model.py): refactor search_documents method to handle search_query logic more efficiently
📝 (chroma.py): add typing_extensions override import for build_vector_store method to improve code readability

* ✨ (tests): add unit tests for ChromaVectorStoreComponent

- Introduced new test suite for ChromaVectorStoreComponent, covering various functionalities including database creation, collection management, similarity and MMR searches, and duplicate handling.
- Implemented tests for creating collections with and without data, ensuring proper functionality and data integrity.
- Verified search capabilities with different query types and result limits, enhancing overall test coverage for the component.

* fix: remove unnecessary whitespace in model.py and add missing import in chroma.py

* fix: mypy error module has no attribute "timeout"

* ♻️ (async_helpers.py): Remove unnecessary type hint ignore comment from timeout_context function

* 📝 (async_helpers.py): add a comment with issue reference PGH003 to document the reason for ignoring type checking in timeout_context function

* ♻️ (async_helpers.py): Remove unnecessary type hint comment to improve code readability and maintainability

* ♻️ (async_helpers.py): Add type ignore comment to suppress miscellaneous type error for timeout_context function

* ♻️ (async_helpers.py): refactor timeout_context function to remove unnecessary type ignore comments and improve code readability

* [autofix.ci] apply automated fixes

* 📝 (async_helpers.py): add a blank line for better code readability and consistency

* fix: mypy error: incompatible redefinition

---------

Co-authored-by: anovazzi1 <otavio2204@gmail.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Co-authored-by: italojohnny <italojohnnydosanjos@gmail.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

src/backend/base/langflow/base/vectorstores/model.py

@@ -122,17 +122,17 @@ class LCVectorStoreComponent(Component):
 
     def search_documents(self) -> list[Data]:
         """Search for documents in the vector store."""
-        search_query: str = self.search_query
-        if not search_query:
-            self.status = ""
-            return []
-
         if self._cached_vector_store is not None:
             vector_store = self._cached_vector_store
         else:
             vector_store = self.build_vector_store()
             self._cached_vector_store = vector_store
 
+        search_query: str = self.search_query
+        if not search_query:
+            self.status = ""
+            return []
+
         self.log(f"Search input: {search_query}")
         self.log(f"Search type: {self.search_type}")
         self.log(f"Number of results: {self.number_of_results}")

src/backend/base/langflow/components/tools/mcp_sse.py

@@ -1,5 +1,4 @@
 # from langflow.field_typing import Data
-import asyncio
 from contextlib import AsyncExitStack
 
 import httpx
@@ -11,6 +10,7 @@ from langflow.components.tools.mcp_stdio import create_input_schema_from_json_sc
 from langflow.custom import Component
 from langflow.field_typing import Tool
 from langflow.io import MessageTextInput, Output
+from langflow.utils.async_helpers import timeout_context
 
 # Define constant for status code
 HTTP_TEMPORARY_REDIRECT = 307
@@ -39,7 +39,7 @@ class MCPSseClient:
             headers = {}
         url = await self.pre_check_redirect(url)
 
-        async with asyncio.timeout(timeout_seconds):
+        async with timeout_context(timeout_seconds):
             sse_transport = await self.exit_stack.enter_async_context(
                 sse_client(url, headers, timeout_seconds, sse_read_timeout_seconds)
             )

src/backend/base/langflow/components/vectorstores/chroma.py

@@ -2,6 +2,7 @@ from copy import deepcopy
 
 from chromadb.config import Settings
 from langchain_chroma import Chroma
+from typing_extensions import override
 
 from langflow.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store
 from langflow.base.vectorstores.utils import chroma_collection_to_data
@@ -82,6 +83,7 @@ class ChromaVectorStoreComponent(LCVectorStoreComponent):
         ),
     ]
 
+    @override
     @check_cached_vector_store
     def build_vector_store(self) -> Chroma:
         """Builds the Chroma object."""

src/backend/base/langflow/utils/async_helpers.py

@@ -1,4 +1,21 @@
 import asyncio
+from contextlib import asynccontextmanager
+
+if hasattr(asyncio, "timeout"):
+
+    @asynccontextmanager
+    async def timeout_context(timeout_seconds):
+        with asyncio.timeout(timeout_seconds) as ctx:
+            yield ctx
+else:
+
+    @asynccontextmanager
+    async def timeout_context(timeout_seconds):
+        try:
+            yield await asyncio.wait_for(asyncio.Future(), timeout=timeout_seconds)
+        except asyncio.TimeoutError as e:
+            msg = f"Operation timed out after {timeout_seconds} seconds"
+            raise TimeoutError(msg) from e
 
 
 def run_until_complete(coro):