diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/.gitignore b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/.gitignore new file mode 100644 index 0000000000000..e6bd6fad4afba --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/.gitignore @@ -0,0 +1,29 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class + +# Virtual Environment +venv/ +env/ +.env + +# IDE files +.vscode/ +.idea/ + +# Logs +*.log + +# Build files +build/ +dist/ +*.egg-info/ + +# Test files +.pytest_cache/ +.coverage + +# OS generated files +.DS_Store +Thumbs.db diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/BUILD b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/BUILD new file mode 100644 index 0000000000000..0896ca890d8bf --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/BUILD @@ -0,0 +1,3 @@ +poetry_requirements( + name="poetry", +) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/Makefile b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/Makefile new file mode 100644 index 0000000000000..88e5b05a11e13 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/Makefile @@ -0,0 +1,20 @@ +.PHONY: all format lint test clean + +all: format lint test + +format: + black . + isort . + +lint: + flake8 . + mypy . + +test: + pytest tests/ + +clean: + find . -type f -name "*.pyc" -delete + find . -type d -name "__pycache__" -delete + rm -rf .pytest_cache + rm -rf .mypy_cache diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/README.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/README.md new file mode 100644 index 0000000000000..7b3d1e9d63d83 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/README.md @@ -0,0 +1,144 @@ +# FalkorDB Vector Store Integration for LlamaIndex + +## Overview + +The FalkorDB Vector Store integration for LlamaIndex allows you to use FalkorDB as a vector store backend for your LlamaIndex applications. This integration supports efficient storage and retrieval of vector embeddings, enabling fast similarity searches and other vector operations. + +## Installation + +To use the FalktorDB vector store with LlamaIndex, you need to install the necessary package: + +```bash +pip install llama-index-vector-stores-falkordb +``` + +## Usage + +### Initializing the FalkorDB Vector Store + +To use the FalkorDB vector store, you first need to initialize it with your FalkorDB connection details: + +```python +from llama_index.vector_stores.falkordb import FalkorDBVectorStore + +vector_store = FalkorDBVectorStore( + url="falkor://localhost:7687", + database="your_database_name", + index_name="your_index_name", + node_label="YourNodeLabel", + embedding_node_property="embedding", + text_node_property="text", + distance_strategy="cosine", + embedding_dimension=1536, +) +``` + +### Adding Documents + +You can add documents to the vector store using the `add` method: + +```python +from llama_index.core.schema import Document + +documents = [ + Document("This is the first document."), + Document("This is the second document."), +] + +vector_store.add(documents) +``` + +### Querying the Vector Store + +To perform a similarity search, use the `query` method: + +```python +from llama_index.core.vector_stores.types import VectorStoreQuery + +query_embedding = [0.1, 0.2, 0.3, ...] # Your query embedding +query = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=5) +results = vector_store.query(query) + +for node, score in zip(results.nodes, results.similarities): + print(f"Text: {node.text}, Score: {score}") +``` + +## Advanced Features + +### Creating a New Index + +If you need to create a new vector index, you can use the `create_new_index` method: + +```python +vector_store.create_new_index() +``` + +This method will create a new vector index in FalkorDB based on the parameters you provided when initializing the `FalkorDBVectorStore`. + +### Retrieving Existing Index + +To check if an index already exists and retrieve its information: + +```python +exists = vector_store.retrieve_existing_index() +if exists: + print("Index exists with the following properties:") + print(f"Node Label: {vector_store.node_label}") + print(f"Embedding Property: {vector_store.embedding_node_property}") + print(f"Embedding Dimension: {vector_store.embedding_dimension}") + print(f"Distance Strategy: {vector_store.distance_strategy}") +else: + print("Index does not exist") +``` + +### Deleting Documents + +To delete documents from the vector store: + +```python +ref_doc_id = "your_document_id" +vector_store.delete(ref_doc_id) +``` + +### Using Metadata Filters + +You can use metadata filters when querying the vector store: + +```python +from llama_index.core.vector_stores.types import ( + MetadataFilters, + MetadataFilter, +) + +filters = MetadataFilters( + filters=[ + MetadataFilter( + key="category", value="science", operator=FilterOperator.EQ + ) + ] +) + +query = VectorStoreQuery( + query_embedding=query_embedding, similarity_top_k=5, filters=filters +) + +results = vector_store.query(query) +``` + +## Best Practices + +1. **Connection Management**: Ensure that you properly manage your FalkorDB connections, especially in production environments. +2. **Index Naming**: Use descriptive names for your indexes to easily identify them in your FalkorDB instance. +3. **Error Handling**: Implement proper error handling to manage potential issues with connections or queries. +4. **Performance Tuning**: Adjust the `embedding_dimension` and `distance_strategy` parameters based on your specific use case and performance requirements. + +## Troubleshooting + +If you encounter issues: + +1. Check your FalkorDB connection details (URL, database name). +2. Ensure that the FalkorDB server is running and accessible. +3. Verify that the index exists and has the correct properties. +4. Check the FalkorDB logs for any error messages. + +For more information on FalkorDB and its capabilities, refer to the [official FalkorDB documentation](https://falkordb.com/docs/). diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/falkordb-vectorstore-docs.md b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/falkordb-vectorstore-docs.md new file mode 100644 index 0000000000000..7b3d1e9d63d83 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/falkordb-vectorstore-docs.md @@ -0,0 +1,144 @@ +# FalkorDB Vector Store Integration for LlamaIndex + +## Overview + +The FalkorDB Vector Store integration for LlamaIndex allows you to use FalkorDB as a vector store backend for your LlamaIndex applications. This integration supports efficient storage and retrieval of vector embeddings, enabling fast similarity searches and other vector operations. + +## Installation + +To use the FalktorDB vector store with LlamaIndex, you need to install the necessary package: + +```bash +pip install llama-index-vector-stores-falkordb +``` + +## Usage + +### Initializing the FalkorDB Vector Store + +To use the FalkorDB vector store, you first need to initialize it with your FalkorDB connection details: + +```python +from llama_index.vector_stores.falkordb import FalkorDBVectorStore + +vector_store = FalkorDBVectorStore( + url="falkor://localhost:7687", + database="your_database_name", + index_name="your_index_name", + node_label="YourNodeLabel", + embedding_node_property="embedding", + text_node_property="text", + distance_strategy="cosine", + embedding_dimension=1536, +) +``` + +### Adding Documents + +You can add documents to the vector store using the `add` method: + +```python +from llama_index.core.schema import Document + +documents = [ + Document("This is the first document."), + Document("This is the second document."), +] + +vector_store.add(documents) +``` + +### Querying the Vector Store + +To perform a similarity search, use the `query` method: + +```python +from llama_index.core.vector_stores.types import VectorStoreQuery + +query_embedding = [0.1, 0.2, 0.3, ...] # Your query embedding +query = VectorStoreQuery(query_embedding=query_embedding, similarity_top_k=5) +results = vector_store.query(query) + +for node, score in zip(results.nodes, results.similarities): + print(f"Text: {node.text}, Score: {score}") +``` + +## Advanced Features + +### Creating a New Index + +If you need to create a new vector index, you can use the `create_new_index` method: + +```python +vector_store.create_new_index() +``` + +This method will create a new vector index in FalkorDB based on the parameters you provided when initializing the `FalkorDBVectorStore`. + +### Retrieving Existing Index + +To check if an index already exists and retrieve its information: + +```python +exists = vector_store.retrieve_existing_index() +if exists: + print("Index exists with the following properties:") + print(f"Node Label: {vector_store.node_label}") + print(f"Embedding Property: {vector_store.embedding_node_property}") + print(f"Embedding Dimension: {vector_store.embedding_dimension}") + print(f"Distance Strategy: {vector_store.distance_strategy}") +else: + print("Index does not exist") +``` + +### Deleting Documents + +To delete documents from the vector store: + +```python +ref_doc_id = "your_document_id" +vector_store.delete(ref_doc_id) +``` + +### Using Metadata Filters + +You can use metadata filters when querying the vector store: + +```python +from llama_index.core.vector_stores.types import ( + MetadataFilters, + MetadataFilter, +) + +filters = MetadataFilters( + filters=[ + MetadataFilter( + key="category", value="science", operator=FilterOperator.EQ + ) + ] +) + +query = VectorStoreQuery( + query_embedding=query_embedding, similarity_top_k=5, filters=filters +) + +results = vector_store.query(query) +``` + +## Best Practices + +1. **Connection Management**: Ensure that you properly manage your FalkorDB connections, especially in production environments. +2. **Index Naming**: Use descriptive names for your indexes to easily identify them in your FalkorDB instance. +3. **Error Handling**: Implement proper error handling to manage potential issues with connections or queries. +4. **Performance Tuning**: Adjust the `embedding_dimension` and `distance_strategy` parameters based on your specific use case and performance requirements. + +## Troubleshooting + +If you encounter issues: + +1. Check your FalkorDB connection details (URL, database name). +2. Ensure that the FalkorDB server is running and accessible. +3. Verify that the index exists and has the correct properties. +4. Check the FalkorDB logs for any error messages. + +For more information on FalkorDB and its capabilities, refer to the [official FalkorDB documentation](https://falkordb.com/docs/). diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/BUILD b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/BUILD new file mode 100644 index 0000000000000..db46e8d6c978c --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/BUILD @@ -0,0 +1 @@ +python_sources() diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/__init__.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/__init__.py new file mode 100644 index 0000000000000..9451e28cdc413 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/__init__.py @@ -0,0 +1,3 @@ +from llama_index.vector_stores.falkordb.base import FalkorDBVectorStore + +__all__ = ["FalkorDBVectorStore"] diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/base.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/base.py new file mode 100644 index 0000000000000..f03310bf0a9b0 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/llama_index/vector_stores/falkordb/base.py @@ -0,0 +1,219 @@ +from typing import Any, Dict, List, Optional +import logging + +from falkordb import FalkorDB + +from llama_index.core.bridge.pydantic import PrivateAttr +from llama_index.core.schema import BaseNode, MetadataMode +from llama_index.core.vector_stores.types import ( + BasePydanticVectorStore, + VectorStoreQuery, + VectorStoreQueryResult, + FilterOperator, + MetadataFilters, + FilterCondition, +) +from llama_index.core.vector_stores.utils import ( + metadata_dict_to_node, + node_to_metadata_dict, +) + +_logger = logging.getLogger(__name__) + + +def clean_params(params: List[BaseNode]) -> List[Dict[str, Any]]: + clean_params = [] + for record in params: + text = record.get_content(metadata_mode=MetadataMode.NONE) + embedding = record.get_embedding() + id = record.node_id + metadata = node_to_metadata_dict(record, remove_text=True, flat_metadata=False) + for k in ["document_id", "doc_id"]: + if k in metadata: + del metadata[k] + clean_params.append( + {"text": text, "embedding": embedding, "id": id, "metadata": metadata} + ) + return clean_params + + +def _to_falkordb_operator(operator: FilterOperator) -> str: + operator_map = { + FilterOperator.EQ: "=", + FilterOperator.GT: ">", + FilterOperator.LT: "<", + FilterOperator.NE: "<>", + FilterOperator.GTE: ">=", + FilterOperator.LTE: "<=", + FilterOperator.IN: "IN", + FilterOperator.NIN: "NOT IN", + FilterOperator.CONTAINS: "CONTAINS", + } + return operator_map.get(operator, "=") + + +def construct_metadata_filter(filters: MetadataFilters): + cypher_snippets = [] + params = {} + for index, filter in enumerate(filters.filters): + cypher_snippets.append( + f"n.`{filter.key}` {_to_falkordb_operator(filter.operator)} $param_{index}" + ) + params[f"param_{index}"] = filter.value + + condition = " OR " if filters.condition == FilterCondition.OR else " AND " + return condition.join(cypher_snippets), params + + +class FalkorDBVectorStore(BasePydanticVectorStore): + stores_text: bool = True + flat_metadata: bool = True + + distance_strategy: str + index_name: str + node_label: str + embedding_node_property: str + text_node_property: str + embedding_dimension: int + + _client: FalkorDB = PrivateAttr() + _database: str = PrivateAttr() + + def __init__( + self, + 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__( + distance_strategy=distance_strategy, + index_name=index_name, + node_label=node_label, + embedding_node_property=embedding_node_property, + text_node_property=text_node_property, + embedding_dimension=embedding_dimension, + ) + + if distance_strategy not in ["cosine", "euclidean"]: + raise ValueError("distance_strategy must be either 'euclidean' or 'cosine'") + + self._client = driver or FalkorDB.from_url(url).select_graph(database) + self._database = database + + # Inline check_if_not_null function + for prop, value in zip( + [ + "index_name", + "node_label", + "embedding_node_property", + "text_node_property", + ], + [index_name, node_label, embedding_node_property, text_node_property], + ): + if not value: + raise ValueError(f"Parameter `{prop}` must not be None or empty string") + + if not self.retrieve_existing_index(): + self.create_new_index() + + @property + def client(self) -> FalkorDB: + return self._client + + def create_new_index(self) -> None: + index_query = ( + f"CREATE VECTOR INDEX {self.index_name} " + f"FOR (n:`{self.node_label}`) " + f"ON (n.`{self.embedding_node_property}`) " + f"OPTIONS {{dimension: {self.embedding_dimension}, metric: '{self.distance_strategy}'}}" + ) + self._client.query(index_query) + + def retrieve_existing_index(self) -> bool: + index_information = self._client.query( + "CALL db.indexes() " + "YIELD label, properties, types, options, entitytype " + "WHERE types = ['VECTOR'] AND label = $index_name", + params={"index_name": self.index_name}, + ) + if index_information.result_set: + index = index_information.result_set[0] + self.node_label = index["entitytype"] + self.embedding_node_property = index["properties"][0] + self.embedding_dimension = index["options"]["dimension"] + self.distance_strategy = index["options"]["metric"] + return True + return False + + def add(self, nodes: List[BaseNode], **add_kwargs: Any) -> List[str]: + ids = [r.node_id for r in nodes] + import_query = ( + "UNWIND $data AS row " + f"MERGE (c:`{self.node_label}` {{id: row.id}}) " + f"SET c.`{self.embedding_node_property}` = row.embedding, " + f"c.`{self.text_node_property}` = row.text, " + "c += row.metadata" + ) + + self._client.query(import_query, params={"data": clean_params(nodes)}) + return ids + + def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult: + base_query = ( + f"MATCH (n:`{self.node_label}`) " + f"WHERE n.`{self.embedding_node_property}` IS NOT NULL " + ) + + if query.filters: + filter_snippets, filter_params = construct_metadata_filter(query.filters) + base_query += f"AND {filter_snippets} " + else: + filter_params = {} + + similarity_query = ( + f"WITH n, vector.similarity.{self.distance_strategy}(" + f"n.`{self.embedding_node_property}`, $embedding) AS score " + "ORDER BY score DESC LIMIT toInteger($k) " + ) + + return_query = ( + f"RETURN n.`{self.text_node_property}` AS text, score, " + "n.id AS id, " + f"n {{.*, `{self.text_node_property}`: NULL, " + f"`{self.embedding_node_property}`: NULL, id: NULL}} AS metadata" + ) + + full_query = base_query + similarity_query + return_query + + parameters = { + "k": query.similarity_top_k, + "embedding": query.query_embedding, + **filter_params, + } + + results = self._client.query(full_query, params=parameters) + + nodes = [] + similarities = [] + ids = [] + for record in results.result_set: + node = metadata_dict_to_node(record["metadata"]) + node.set_content(str(record["text"])) + nodes.append(node) + similarities.append(record["score"]) + ids.append(record["id"]) + + return VectorStoreQueryResult(nodes=nodes, similarities=similarities, ids=ids) + + def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: + self._client.query( + f"MATCH (n:`{self.node_label}`) WHERE n.ref_doc_id = $id DETACH DELETE n", + params={"id": ref_doc_id}, + ) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/pyproject.toml new file mode 100644 index 0000000000000..bcc84ec45a311 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/pyproject.toml @@ -0,0 +1,63 @@ +[build-system] +build-backend = "poetry.core.masonry.api" +requires = ["poetry-core"] + +[tool.codespell] +check-filenames = true +check-hidden = true +skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb" + +[tool.llamahub] +contains_example = false +import_path = "llama_index.vector_stores.falkordb" + +[tool.llamahub.class_authors] +FalkorDBVectorStore = "llama-index" + +[tool.mypy] +disallow_untyped_defs = true +exclude = ["_static", "build", "examples", "notebooks", "venv"] +ignore_missing_imports = true +python_version = "3.8" + +[tool.poetry] +authors = ["Your Name "] +description = "llama-index vector stores falkordb integration" +exclude = ["**/BUILD"] +license = "MIT" +name = "llama-index-vector-stores-falkordb" +readme = "README.md" +version = "0.1.0" + +[tool.poetry.dependencies] +python = ">=3.8.1,<4.0" +falkordb = "^1.0.8" +llama-index-core = "^0.11.0" + +[tool.poetry.group.dev.dependencies] +ipython = "8.10.0" +jupyter = "^1.0.0" +mypy = "0.991" +pre-commit = "3.2.0" +pylint = "2.15.10" +pytest = "7.2.1" +pytest-mock = "3.11.1" +ruff = "0.0.292" +tree-sitter-languages = "^1.8.0" +types-Deprecated = ">=0.1.0" +types-PyYAML = "^6.0.12.12" +types-protobuf = "^4.24.0.4" +types-redis = "4.5.5.0" +types-requests = "2.28.11.8" +types-setuptools = "67.1.0.0" + +[tool.poetry.group.dev.dependencies.black] +extras = ["jupyter"] +version = "<=23.9.1,>=23.7.0" + +[tool.poetry.group.dev.dependencies.codespell] +extras = ["toml"] +version = ">=v2.2.6" + +[[tool.poetry.packages]] +include = "llama_index/" diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/tests/BUILD b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/tests/BUILD new file mode 100644 index 0000000000000..dabf212d7e716 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/tests/BUILD @@ -0,0 +1 @@ +python_tests() diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/tests/__init__.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/tests/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/tests/test_falkordb.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/tests/test_falkordb.py new file mode 100644 index 0000000000000..e0371254a8f1b --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-falkordb/tests/test_falkordb.py @@ -0,0 +1,214 @@ +import pytest +from unittest.mock import MagicMock, patch + +from llama_index.core.schema import TextNode +from llama_index.vector_stores.falkordb import FalkorDBVectorStore +from llama_index.core.vector_stores.types import ( + VectorStoreQuery, + MetadataFilters, + ExactMatchFilter, +) + + +# Mock FalkorDB client +class MockFalkorDBClient: + def __init__(self) -> None: + self.nodes = {} + self.query_results = [] + self.index_exists = False + + def query(self, query, params=None): + if "CREATE VECTOR INDEX" in query: + self.index_exists = True + return MagicMock() + elif "SHOW INDEXES" in query: + if self.index_exists: + return MagicMock(result_set=[{"name": "test_index"}]) + return MagicMock(result_set=[]) + elif "MATCH (n:Chunk)" in query and "RETURN" in query: + # Query operation + return MagicMock(result_set=self.query_results) + elif "MERGE (c:Chunk" in query: + # Add operation + if params and "data" in params: + for node in params["data"]: + self.nodes[node["id"]] = node + return MagicMock() + elif "MATCH (n:Chunk) WHERE n.id" in query and "DELETE" in query: + # Delete operation + node_id = params["id"] + if node_id in self.nodes: + del self.nodes[node_id] + return MagicMock() + elif "MATCH (n:Chunk) WHERE n.id" in query: + # Get operation + node_id = params["id"] + if node_id in self.nodes: + return MagicMock(result_set=[self.nodes[node_id]]) + return MagicMock(result_set=[]) + return MagicMock() + + def set_query_results(self, results): + self.query_results = results + + +@pytest.fixture() +def mock_falkordb(): + with patch("falkordb.FalkorDB") as mock: + client = MockFalkorDBClient() + mock.return_value = client + mock.from_url.return_value.select_graph.return_value = client + yield client + + +@pytest.fixture() +def falkordb_store(mock_falkordb): + store = FalkorDBVectorStore( + url="redis://localhost:6379", + database="testdb", + index_name="test_index", + node_label="Chunk", + embedding_node_property="embedding", + text_node_property="text", + ) + # Ensure store uses the mock client + store._client = mock_falkordb + return store + + +def test_falkordb_add(falkordb_store): + nodes = [ + TextNode( + text="Hello world", + id_="1", + embedding=[1.0, 0.0, 0.0], + metadata={"key": "value"}, + ), + TextNode( + text="Hello world 2", + id_="2", + embedding=[0.0, 1.0, 0.0], + metadata={"key2": "value2"}, + ), + ] + ids = falkordb_store.add(nodes) + assert ids == ["1", "2"] + assert len(falkordb_store._client.nodes) == 2 + assert falkordb_store._client.nodes["1"]["text"] == "Hello world" + assert falkordb_store._client.nodes["2"]["text"] == "Hello world 2" + + +def test_falkordb_delete(falkordb_store): + node = TextNode( + text="Hello world", + id_="test_node", + embedding=[1.0, 0.0, 0.0], + ) + falkordb_store.add([node]) + assert "test_node" in falkordb_store._client.nodes + + falkordb_store.delete("test_node") + assert "test_node" not in falkordb_store._client.nodes + + +def test_falkordb_query(falkordb_store): + mock_results = [ + { + "n": { + "text": "Hello world", + "id": "1", + "embedding": [1.0, 0.0, 0.0], + "metadata": {"key": "value"}, + }, + "score": 0.9, + }, + { + "n": { + "text": "Hello world 2", + "id": "2", + "embedding": [0.0, 1.0, 0.0], + "metadata": {"key2": "value2"}, + }, + "score": 0.7, + }, + ] + falkordb_store._client.set_query_results(mock_results) + + query = VectorStoreQuery( + query_embedding=[1.0, 0.0, 0.0], + similarity_top_k=2, + ) + results = falkordb_store.query(query) + + assert len(results.nodes) == 2 + assert results.nodes[0].text == "Hello world" + assert results.nodes[1].text == "Hello world 2" + assert results.similarities == [0.9, 0.7] + + +def test_falkordb_query_with_filters(falkordb_store): + mock_results = [ + { + "n": { + "text": "Hello world", + "id": "1", + "embedding": [1.0, 0.0, 0.0], + "metadata": {"key": "value"}, + }, + "score": 0.9, + }, + ] + falkordb_store._client.set_query_results(mock_results) + + query = VectorStoreQuery( + query_embedding=[1.0, 0.0, 0.0], + similarity_top_k=2, + filters=MetadataFilters(filters=[ExactMatchFilter(key="key", value="value")]), + ) + results = falkordb_store.query(query) + + assert len(results.nodes) == 1 + assert results.nodes[0].text == "Hello world" + assert results.similarities == [0.9] + + +def test_falkordb_update(falkordb_store): + node = TextNode( + text="Original text", + id_="update_node", + embedding=[1.0, 0.0, 0.0], + ) + falkordb_store.add([node]) + + updated_node = TextNode( + text="Updated text", + id_="update_node", + embedding=[0.0, 1.0, 0.0], + ) + falkordb_store.update(updated_node) + + assert falkordb_store._client.nodes["update_node"]["text"] == "Updated text" + assert falkordb_store._client.nodes["update_node"]["embedding"] == [0.0, 1.0, 0.0] + + +def test_falkordb_get(falkordb_store): + node = TextNode( + text="Get test", + id_="get_node", + embedding=[1.0, 1.0, 1.0], + ) + falkordb_store.add([node]) + + retrieved_node = falkordb_store.get("get_node") + assert retrieved_node is not None + assert retrieved_node.text == "Get test" + assert retrieved_node.embedding == [1.0, 1.0, 1.0] + + +def test_falkordb_nonexistent_get(falkordb_store): + retrieved_node = falkordb_store.get("nonexistent_node") + assert retrieved_node is None + + +if __name__ == "__main__": + pytest.main()