Fix some issues

run-llama · Oct 4, 2024 · d064c47 · d064c47 · A5jadAli · Oct 4, 2024
1 parent 046c0cf
commit d064c47
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/..._stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/__init__.py b/..._stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/__init__.py
@@ -1,3 +1,3 @@
-from llama_index.vector_stores.falkordb import FalkorDBVectorStore
+from llama_index.vector_stores.falkordb.base import FalkorDBVectorStore
 
 __all__ = ["FalkorDBVectorStore"]
diff --git a/...ctor_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/base.py b/...ctor_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/base.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 import logging
 
 from falkordb import FalkorDB
@@ -81,14 +81,15 @@ class FalkorDBVectorStore(BasePydanticVectorStore):
 
     def __init__(
         self,
-        url: str,
+        url: Optional[str] = None,
         database: str = "falkor",
         index_name: str = "vector",
         node_label: str = "Chunk",
         embedding_node_property: str = "embedding",
         text_node_property: str = "text",
         distance_strategy: str = "cosine",
         embedding_dimension: int = 1536,
+        driver: Optional[FalkorDB] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -103,7 +104,7 @@ def __init__(
         if distance_strategy not in ["cosine", "euclidean"]:
             raise ValueError("distance_strategy must be either 'euclidean' or 'cosine'")
 
-        self._driver = FalkorDB.from_url(url).select_graph(database)
+        self._driver = driver or FalkorDB.from_url(url).select_graph(database)
         self._database = database
 
         # Inline check_if_not_null function

diff --git a/...dex-vector-stores-falkordb/tests/tests.py → ...or-stores-falkordb/tests/test_falkordb.py b/...dex-vector-stores-falkordb/tests/tests.py → ...or-stores-falkordb/tests/test_falkordb.py
@@ -4,7 +4,7 @@
 
 from llama_index.core.schema import TextNode
 from llama_index.vector_stores.falkordb import FalkorDBVectorStore
-from llama_index.vector_stores.types import (
+from llama_index.core.vector_stores.types import (
     VectorStoreQuery,
     MetadataFilters,
     ExactMatchFilter,
@@ -50,7 +50,7 @@ def mock_falkordb():
 @pytest.fixture()
 def falkordb_store(mock_falkordb):
     return FalkorDBVectorStore(
-        url="bolt://localhost:7687",
+        driver=mock_falkordb,
         database="testdb",
         index_name="test_index",
         node_label="Chunk",