Skip to content

Falkordb

FalkorDBGraphStore #

Bases: GraphStore

FalkorDB Graph Store.

In this graph store, triplets are stored within FalkorDB.

Parameters:

Name Type Description Default
url str

The URL for the FalkorDB database.

 required
database str

The name of the graph to connect to. Defaults to "falkor".

 'falkor'
node_label str

The label used for every entity node. Defaults to "Entity".

 'Entity'
**kwargs Any

Additional keyword arguments forwarded to the FalkorDB client (e.g. username, password, ssl).

 {}
Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/base.py 
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
class FalkorDBGraphStore(GraphStore):
    """
    FalkorDB Graph Store.

    In this graph store, triplets are stored within FalkorDB.

    Args:
        url (str): The URL for the FalkorDB database.
        database (str): The name of the graph to connect to. Defaults to "falkor".
        node_label (str): The label used for every entity node. Defaults to "Entity".
        **kwargs (Any): Additional keyword arguments forwarded to the FalkorDB
            client (e.g. ``username``, ``password``, ``ssl``).

    """

    def __init__(
        self,
        url: str,
        database: str = "falkor",
        node_label: str = "Entity",
        **kwargs: Any,
    ) -> None:
        """Initialize params."""
        self._node_label = node_label

        self._driver = FalkorDB.from_url(url, **kwargs)
        self._graph = self._driver.select_graph(database)
        self._create_index()

        self._database = database

        self.schema = ""
        self.get_query = f"""
            MATCH (n1:`{self._node_label}`)-[r]->(n2:`{self._node_label}`)
            WHERE n1.id = $subj RETURN type(r), n2.id
        """

    def _create_index(self) -> None:
        """Create the index backing every `id` lookup, if it does not exist."""
        try:
            self._graph.query(f"CREATE INDEX FOR (n:`{self._node_label}`) ON (n.id)")
        except redis.ResponseError as e:
            if "already indexed" not in str(e).lower():
                logger.warning("Create index failed: %s", e)

    @property
    def client(self) -> None:
        return self._graph

    def get(self, subj: str) -> List[List[str]]:
        """Get triplets."""
        result = self._graph.query(self.get_query, params={"subj": subj})
        return result.result_set

    def get_rel_map(
        self, subjs: Optional[List[str]] = None, depth: int = 2, limit: int = 30
    ) -> Dict[str, List[List[str]]]:
        """Get flat rel map."""
        # The flat means for multi-hop relation path, we could get
        # knowledge like: subj -> rel -> obj -> rel -> obj -> rel -> obj.
        # This type of knowledge is useful for some tasks.
        # +-------------+------------------------------------+
        # | subj        | flattened_rels                     |
        # +-------------+------------------------------------+
        # | "player101" | [95, "player125", 2002, "team204"] |
        # | "player100" | [1997, "team204"]                  |
        # ...
        # +-------------+------------------------------------+

        rel_map: Dict[Any, List[Any]] = {}
        if subjs is None or len(subjs) == 0:
            # unlike simple graph_store, we don't do get_all here
            return rel_map

        query = f"""
            MATCH (n1:{self._node_label})
            WHERE n1.id IN $subjs
            WITH n1
            MATCH p=(n1)-[e*1..{depth}]->(z)
            RETURN p LIMIT {limit}
        """

        data = self.query(query, params={"subjs": subjs})
        if not data:
            return rel_map

        for record in data:
            nodes = record[0].nodes()
            edges = record[0].edges()

            subj_id = nodes[0].properties["id"]
            path = []
            for i, edge in enumerate(edges):
                dest = nodes[i + 1]
                dest_id = dest.properties["id"]
                path.append(edge.relation)
                path.append(dest_id)

            paths = rel_map[subj_id] if subj_id in rel_map else []
            paths.append(path)
            rel_map[subj_id] = paths

        return rel_map

    def upsert_triplet(self, subj: str, rel: str, obj: str) -> None:
        """Add triplet."""
        query = """
            MERGE (n1:`%s` {id:$subj})
            MERGE (n2:`%s` {id:$obj})
            MERGE (n1)-[:`%s`]->(n2)
        """

        prepared_statement = query % (
            self._node_label,
            self._node_label,
            rel.replace(" ", "_").upper(),
        )

        # Call FalkorDB with prepared statement
        self._graph.query(prepared_statement, params={"subj": subj, "obj": obj})

    def delete(self, subj: str, rel: str, obj: str) -> None:
        """Delete triplet."""

        def delete_rel(subj: str, obj: str, rel: str) -> None:
            rel = rel.replace(" ", "_").upper()
            query = f"""
                MATCH (n1:`{self._node_label}`)-[r:`{rel}`]->(n2:`{self._node_label}`)
                WHERE n1.id = $subj AND n2.id = $obj DELETE r
            """

            # Call FalkorDB with prepared statement
            self._graph.query(query, params={"subj": subj, "obj": obj})

        def delete_entity(entity: str) -> None:
            query = f"MATCH (n:`{self._node_label}`) WHERE n.id = $entity DELETE n"

            # Call FalkorDB with prepared statement
            self._graph.query(query, params={"entity": entity})

        def check_edges(entity: str) -> bool:
            query = f"""
                MATCH (n1:`{self._node_label}`)--()
                WHERE n1.id = $entity RETURN count(*)
            """

            # Call FalkorDB with prepared statement
            result = self._graph.query(query, params={"entity": entity})
            # `RETURN count(*)` always yields a single row, so the row count
            # itself carries no information - the counter value does.
            return bool(result.result_set) and result.result_set[0][0] > 0

        delete_rel(subj, obj, rel)
        if not check_edges(subj):
            delete_entity(subj)
        if not check_edges(obj):
            delete_entity(obj)

    def refresh_schema(self) -> None:
        """
        Refreshes the FalkorDB graph schema information.
        """
        node_properties = self.query("CALL DB.PROPERTYKEYS()")
        relationships = self.query("CALL DB.RELATIONSHIPTYPES()")

        self.schema = f"""
        Properties: {node_properties}
        Relationships: {relationships}
        """

    def get_schema(self, refresh: bool = False) -> str:
        """Get the schema of the FalkorDBGraph store."""
        if self.schema and not refresh:
            return self.schema
        self.refresh_schema()
        logger.debug(f"get_schema() schema:\n{self.schema}")
        return self.schema

    def query(self, query: str, params: Optional[Dict[str, Any]] = None) -> Any:
        result = self._graph.query(query, params=params)
        return result.result_set

    def switch_graph(self, graph_name: str) -> None:
        """
        Switch to the given graph name (`graph_name`).

        This method allows users to change the active graph within the same
        database connection.

        Args:
            graph_name (str): The name of the graph to switch to.

        """
        self._graph = self._driver.select_graph(graph_name)
        self._database = graph_name
        self._create_index()

        try:
            self.refresh_schema()
        except Exception as e:
            raise ValueError(f"Could not refresh schema. Error: {e}")

    def close(self) -> None:
        """Explicitly close the FalkorDB connection."""
        if hasattr(self, "_driver"):
            try:
                self._driver.connection.close()
            finally:
                delattr(self, "_driver")

    def __enter__(self) -> "FalkorDBGraphStore":
        """Enter the runtime context, enabling `with FalkorDBGraphStore(...)`."""
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        """Close the connection when leaving the runtime context."""
        self.close()

    def __del__(self) -> None:
        """Best-effort cleanup; prefer `close()` or the context manager."""
        try:
            self.close()
        except Exception:
            # Suppress any exceptions during garbage collection
            pass

get #

get(subj: str) -> List[List[str]]

Get triplets.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/base.py 
63
64
65
66
def get(self, subj: str) -> List[List[str]]:
    """Get triplets."""
    result = self._graph.query(self.get_query, params={"subj": subj})
    return result.result_set

get_rel_map #

get_rel_map(
    subjs: Optional[List[str]] = None,
    depth: int = 2,
    limit: int = 30,
) -> Dict[str, List[List[str]]]

Get flat rel map.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/base.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def get_rel_map(
    self, subjs: Optional[List[str]] = None, depth: int = 2, limit: int = 30
) -> Dict[str, List[List[str]]]:
    """Get flat rel map."""
    # The flat means for multi-hop relation path, we could get
    # knowledge like: subj -> rel -> obj -> rel -> obj -> rel -> obj.
    # This type of knowledge is useful for some tasks.
    # +-------------+------------------------------------+
    # | subj        | flattened_rels                     |
    # +-------------+------------------------------------+
    # | "player101" | [95, "player125", 2002, "team204"] |
    # | "player100" | [1997, "team204"]                  |
    # ...
    # +-------------+------------------------------------+

    rel_map: Dict[Any, List[Any]] = {}
    if subjs is None or len(subjs) == 0:
        # unlike simple graph_store, we don't do get_all here
        return rel_map

    query = f"""
        MATCH (n1:{self._node_label})
        WHERE n1.id IN $subjs
        WITH n1
        MATCH p=(n1)-[e*1..{depth}]->(z)
        RETURN p LIMIT {limit}
    """

    data = self.query(query, params={"subjs": subjs})
    if not data:
        return rel_map

    for record in data:
        nodes = record[0].nodes()
        edges = record[0].edges()

        subj_id = nodes[0].properties["id"]
        path = []
        for i, edge in enumerate(edges):
            dest = nodes[i + 1]
            dest_id = dest.properties["id"]
            path.append(edge.relation)
            path.append(dest_id)

        paths = rel_map[subj_id] if subj_id in rel_map else []
        paths.append(path)
        rel_map[subj_id] = paths

    return rel_map

upsert_triplet #

upsert_triplet(subj: str, rel: str, obj: str) -> None

Add triplet.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/base.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def upsert_triplet(self, subj: str, rel: str, obj: str) -> None:
    """Add triplet."""
    query = """
        MERGE (n1:`%s` {id:$subj})
        MERGE (n2:`%s` {id:$obj})
        MERGE (n1)-[:`%s`]->(n2)
    """

    prepared_statement = query % (
        self._node_label,
        self._node_label,
        rel.replace(" ", "_").upper(),
    )

    # Call FalkorDB with prepared statement
    self._graph.query(prepared_statement, params={"subj": subj, "obj": obj})

delete #

delete(subj: str, rel: str, obj: str) -> None

Delete triplet.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/base.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def delete(self, subj: str, rel: str, obj: str) -> None:
    """Delete triplet."""

    def delete_rel(subj: str, obj: str, rel: str) -> None:
        rel = rel.replace(" ", "_").upper()
        query = f"""
            MATCH (n1:`{self._node_label}`)-[r:`{rel}`]->(n2:`{self._node_label}`)
            WHERE n1.id = $subj AND n2.id = $obj DELETE r
        """

        # Call FalkorDB with prepared statement
        self._graph.query(query, params={"subj": subj, "obj": obj})

    def delete_entity(entity: str) -> None:
        query = f"MATCH (n:`{self._node_label}`) WHERE n.id = $entity DELETE n"

        # Call FalkorDB with prepared statement
        self._graph.query(query, params={"entity": entity})

    def check_edges(entity: str) -> bool:
        query = f"""
            MATCH (n1:`{self._node_label}`)--()
            WHERE n1.id = $entity RETURN count(*)
        """

        # Call FalkorDB with prepared statement
        result = self._graph.query(query, params={"entity": entity})
        # `RETURN count(*)` always yields a single row, so the row count
        # itself carries no information - the counter value does.
        return bool(result.result_set) and result.result_set[0][0] > 0

    delete_rel(subj, obj, rel)
    if not check_edges(subj):
        delete_entity(subj)
    if not check_edges(obj):
        delete_entity(obj)

refresh_schema #

refresh_schema() -> None

Refreshes the FalkorDB graph schema information.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/base.py 
172
173
174
175
176
177
178
179
180
181
182
def refresh_schema(self) -> None:
    """
    Refreshes the FalkorDB graph schema information.
    """
    node_properties = self.query("CALL DB.PROPERTYKEYS()")
    relationships = self.query("CALL DB.RELATIONSHIPTYPES()")

    self.schema = f"""
    Properties: {node_properties}
    Relationships: {relationships}
    """

get_schema #

get_schema(refresh: bool = False) -> str

Get the schema of the FalkorDBGraph store.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/base.py 
184
185
186
187
188
189
190
def get_schema(self, refresh: bool = False) -> str:
    """Get the schema of the FalkorDBGraph store."""
    if self.schema and not refresh:
        return self.schema
    self.refresh_schema()
    logger.debug(f"get_schema() schema:\n{self.schema}")
    return self.schema

switch_graph #

switch_graph(graph_name: str) -> None

Switch to the given graph name (graph_name).

This method allows users to change the active graph within the same database connection.

Parameters:

Name Type Description Default
graph_name str

The name of the graph to switch to.

 required
Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/base.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def switch_graph(self, graph_name: str) -> None:
    """
    Switch to the given graph name (`graph_name`).

    This method allows users to change the active graph within the same
    database connection.

    Args:
        graph_name (str): The name of the graph to switch to.

    """
    self._graph = self._driver.select_graph(graph_name)
    self._database = graph_name
    self._create_index()

    try:
        self.refresh_schema()
    except Exception as e:
        raise ValueError(f"Could not refresh schema. Error: {e}")

close #

close() -> None

Explicitly close the FalkorDB connection.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/base.py 
216
217
218
219
220
221
222
def close(self) -> None:
    """Explicitly close the FalkorDB connection."""
    if hasattr(self, "_driver"):
        try:
            self._driver.connection.close()
        finally:
            delattr(self, "_driver")

FalkorDBPropertyGraphStore #

Bases: PropertyGraphStore

FalkorDB Property Graph Store.

This class implements a FalkorDB property graph store.

If you are using local FalkorDB instead of FalkorDB Cloud, here's a helpful command for launching the docker container:

docker run \
    -p 3000:3000 -p 6379:6379 \
    -v $PWD/data:/data \
    falkordb/falkordb:latest

Parameters:

Name Type Description Default
url str

The URL for the FalkorDB database.

 required
database Optional[str]

The name of the graph to connect to. Defaults to "falkor".

 'falkor'
refresh_schema bool

Whether to read the graph schema on startup. Defaults to True.

 True
sanitize_query_output bool

Whether to strip oversized values (such as embeddings) from query results. Defaults to True.

 True
create_indexes bool

Whether to create the indexes used for fast lookups and vector search. Defaults to True.

 True
timeout Optional[int]

Query timeout in milliseconds. Defaults to None.

 None
**falkordb_kwargs Any

Additional keyword arguments forwarded to the FalkorDB client (e.g. username, password, ssl).

 {}

Examples:

pip install llama-index-graph-stores-falkordb

from llama_index.core.indices.property_graph import PropertyGraphIndex
from llama_index.graph_stores.falkordb import FalkorDBPropertyGraphStore

# Create a FalkorDBPropertyGraphStore instance
graph_store = FalkorDBPropertyGraphStore(
    url="falkordb://localhost:6379",
    database="falkor"
)

# create the index
index = PropertyGraphIndex.from_documents(
    documents,
    property_graph_store=graph_store,
)
Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/falkordb_property_graph.py 
 136
 137
 138
 139
 140
 141
 142
 143
 144
 145
 146
 147
 148
 149
 150
 151
 152
 153
 154
 155
 156
 157
 158
 159
 160
 161
 162
 163
 164
 165
 166
 167
 168
 169
 170
 171
 172
 173
 174
 175
 176
 177
 178
 179
 180
 181
 182
 183
 184
 185
 186
 187
 188
 189
 190
 191
 192
 193
 194
 195
 196
 197
 198
 199
 200
 201
 202
 203
 204
 205
 206
 207
 208
 209
 210
 211
 212
 213
 214
 215
 216
 217
 218
 219
 220
 221
 222
 223
 224
 225
 226
 227
 228
 229
 230
 231
 232
 233
 234
 235
 236
 237
 238
 239
 240
 241
 242
 243
 244
 245
 246
 247
 248
 249
 250
 251
 252
 253
 254
 255
 256
 257
 258
 259
 260
 261
 262
 263
 264
 265
 266
 267
 268
 269
 270
 271
 272
 273
 274
 275
 276
 277
 278
 279
 280
 281
 282
 283
 284
 285
 286
 287
 288
 289
 290
 291
 292
 293
 294
 295
 296
 297
 298
 299
 300
 301
 302
 303
 304
 305
 306
 307
 308
 309
 310
 311
 312
 313
 314
 315
 316
 317
 318
 319
 320
 321
 322
 323
 324
 325
 326
 327
 328
 329
 330
 331
 332
 333
 334
 335
 336
 337
 338
 339
 340
 341
 342
 343
 344
 345
 346
 347
 348
 349
 350
 351
 352
 353
 354
 355
 356
 357
 358
 359
 360
 361
 362
 363
 364
 365
 366
 367
 368
 369
 370
 371
 372
 373
 374
 375
 376
 377
 378
 379
 380
 381
 382
 383
 384
 385
 386
 387
 388
 389
 390
 391
 392
 393
 394
 395
 396
 397
 398
 399
 400
 401
 402
 403
 404
 405
 406
 407
 408
 409
 410
 411
 412
 413
 414
 415
 416
 417
 418
 419
 420
 421
 422
 423
 424
 425
 426
 427
 428
 429
 430
 431
 432
 433
 434
 435
 436
 437
 438
 439
 440
 441
 442
 443
 444
 445
 446
 447
 448
 449
 450
 451
 452
 453
 454
 455
 456
 457
 458
 459
 460
 461
 462
 463
 464
 465
 466
 467
 468
 469
 470
 471
 472
 473
 474
 475
 476
 477
 478
 479
 480
 481
 482
 483
 484
 485
 486
 487
 488
 489
 490
 491
 492
 493
 494
 495
 496
 497
 498
 499
 500
 501
 502
 503
 504
 505
 506
 507
 508
 509
 510
 511
 512
 513
 514
 515
 516
 517
 518
 519
 520
 521
 522
 523
 524
 525
 526
 527
 528
 529
 530
 531
 532
 533
 534
 535
 536
 537
 538
 539
 540
 541
 542
 543
 544
 545
 546
 547
 548
 549
 550
 551
 552
 553
 554
 555
 556
 557
 558
 559
 560
 561
 562
 563
 564
 565
 566
 567
 568
 569
 570
 571
 572
 573
 574
 575
 576
 577
 578
 579
 580
 581
 582
 583
 584
 585
 586
 587
 588
 589
 590
 591
 592
 593
 594
 595
 596
 597
 598
 599
 600
 601
 602
 603
 604
 605
 606
 607
 608
 609
 610
 611
 612
 613
 614
 615
 616
 617
 618
 619
 620
 621
 622
 623
 624
 625
 626
 627
 628
 629
 630
 631
 632
 633
 634
 635
 636
 637
 638
 639
 640
 641
 642
 643
 644
 645
 646
 647
 648
 649
 650
 651
 652
 653
 654
 655
 656
 657
 658
 659
 660
 661
 662
 663
 664
 665
 666
 667
 668
 669
 670
 671
 672
 673
 674
 675
 676
 677
 678
 679
 680
 681
 682
 683
 684
 685
 686
 687
 688
 689
 690
 691
 692
 693
 694
 695
 696
 697
 698
 699
 700
 701
 702
 703
 704
 705
 706
 707
 708
 709
 710
 711
 712
 713
 714
 715
 716
 717
 718
 719
 720
 721
 722
 723
 724
 725
 726
 727
 728
 729
 730
 731
 732
 733
 734
 735
 736
 737
 738
 739
 740
 741
 742
 743
 744
 745
 746
 747
 748
 749
 750
 751
 752
 753
 754
 755
 756
 757
 758
 759
 760
 761
 762
 763
 764
 765
 766
 767
 768
 769
 770
 771
 772
 773
 774
 775
 776
 777
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
class FalkorDBPropertyGraphStore(PropertyGraphStore):
    r"""
    FalkorDB Property Graph Store.

    This class implements a FalkorDB property graph store.

    If you are using local FalkorDB instead of FalkorDB Cloud, here's a helpful
    command for launching the docker container:

    ```bash
    docker run \
        -p 3000:3000 -p 6379:6379 \
        -v $PWD/data:/data \
        falkordb/falkordb:latest
    ```

    Args:
        url (str): The URL for the FalkorDB database.
        database (Optional[str]): The name of the graph to connect to. Defaults to "falkor".
        refresh_schema (bool): Whether to read the graph schema on startup. Defaults to True.
        sanitize_query_output (bool): Whether to strip oversized values (such as
            embeddings) from query results. Defaults to True.
        create_indexes (bool): Whether to create the indexes used for fast lookups
            and vector search. Defaults to True.
        timeout (Optional[int]): Query timeout in milliseconds. Defaults to None.
        **falkordb_kwargs (Any): Additional keyword arguments forwarded to the
            FalkorDB client (e.g. ``username``, ``password``, ``ssl``).

    Examples:
        `pip install llama-index-graph-stores-falkordb`

        ```python
        from llama_index.core.indices.property_graph import PropertyGraphIndex
        from llama_index.graph_stores.falkordb import FalkorDBPropertyGraphStore

        # Create a FalkorDBPropertyGraphStore instance
        graph_store = FalkorDBPropertyGraphStore(
            url="falkordb://localhost:6379",
            database="falkor"
        )

        # create the index
        index = PropertyGraphIndex.from_documents(
            documents,
            property_graph_store=graph_store,
        )
        ```

    """

    supports_structured_queries: bool = True
    supports_vector_queries: bool = True
    text_to_cypher_template: PromptTemplate = DEFAULT_CYPHER_TEMPALTE

    def __init__(
        self,
        url: str,
        database: str = "falkor",
        refresh_schema: bool = True,
        sanitize_query_output: bool = True,
        create_indexes: bool = True,
        timeout: Optional[int] = None,
        **falkordb_kwargs: Any,
    ) -> None:
        self.sanitize_query_output = sanitize_query_output
        self._create_indexes = create_indexes
        self._timeout = timeout
        self._driver = FalkorDB.from_url(url, **falkordb_kwargs)
        self._graph = self._driver.select_graph(database)
        self._database = database
        self.structured_schema = {}
        # label -> dimension of the vector index defined on `embedding`
        self._vector_index_dimensions: Dict[str, int] = {}

        if self._create_indexes:
            self._create_range_indexes()
        self._refresh_vector_index_info()

        if refresh_schema:
            self.refresh_schema()

    @property
    def client(self):
        return self._graph

    ### ----- Index management ----- ###

    def _run_index_statement(self, statement: str) -> bool:
        """Run an index/constraint statement, tolerating "already exists" errors."""
        try:
            self.structured_query(statement)
        except redis.exceptions.ResponseError as e:
            message = str(e).lower()
            if "already indexed" in message or "already exists" in message:
                return True
            logger.warning("Could not create index with `%s`: %s", statement, e)
            return False
        return True

    def _create_range_indexes(self) -> None:
        """Create the range indexes backing every `id` lookup and MERGE."""
        for label in (BASE_ENTITY_LABEL, CHUNK_LABEL):
            self._run_index_statement(
                f"CREATE INDEX FOR (n:{escape_identifier(label)}) ON (n.id)"
            )

    def _refresh_vector_index_info(self) -> None:
        """Record the dimension of every existing vector index on `embedding`."""
        self._vector_index_dimensions = {}
        try:
            indexes = self.structured_query("CALL db.indexes()")
        except redis.exceptions.ResponseError as e:
            logger.debug("Could not list indexes: %s", e)
            return

        for index in indexes or []:
            types = index.get("types") or {}
            if "VECTOR" not in (types.get(EMBEDDING_KEY) or []):
                continue
            options = (index.get("options") or {}).get(EMBEDDING_KEY) or {}
            dimension = options.get("dimension")
            if index.get("label") is not None and dimension is not None:
                self._vector_index_dimensions[index["label"]] = int(dimension)

    def _ensure_vector_index(self, label: str, dimension: int) -> None:
        """
        Create the vector index for ``label`` if it does not exist yet.

        FalkorDB requires the embedding dimension up front, so the index can
        only be created once the first embedding is known.
        """
        if not self._create_indexes:
            return

        existing = self._vector_index_dimensions.get(label)
        if existing == dimension:
            return
        if existing is not None:
            logger.warning(
                "A vector index on `%s`.embedding already exists with dimension %s, "
                "but an embedding of dimension %s was provided. Vector search will "
                "fall back to a full scan.",
                label,
                existing,
                dimension,
            )
            return

        created = self._run_index_statement(
            f"CREATE VECTOR INDEX FOR (n:{escape_identifier(label)}) ON (n.{EMBEDDING_KEY}) "
            f"OPTIONS {{dimension: {int(dimension)}, "
            f"similarityFunction: '{VECTOR_SIMILARITY_FUNCTION}'}}"
        )
        if created:
            self._refresh_vector_index_info()

    ### ----- Schema ----- ###

    def _sample_node_properties(self) -> List[dict]:
        """
        Derive the property types of every node label in the graph.

        Labels are enumerated exhaustively and only the property *values* are
        sampled, so a label that happens to sit outside the first
        ``SCHEMA_SAMPLE_SIZE`` scanned nodes still appears in the schema.
        """
        excluded = {*EXCLUDED_LABELS, BASE_ENTITY_LABEL}
        labels = [
            row["label"]
            for row in self.structured_query("CALL db.labels()") or []
            if row["label"] not in excluded
        ]

        outputs = []
        for label in labels:
            rows = self.structured_query(
                f"""
                MATCH (n:{escape_identifier(label)})
                WITH n LIMIT $SAMPLE_SIZE
                UNWIND [k IN keys(n) WHERE k <> '{EMBEDDING_KEY}'] AS key
                WITH key, collect(n[key])[0] AS sample
                RETURN collect({{property: key, sample: sample}}) AS keys
                """,
                param_map={"SAMPLE_SIZE": SCHEMA_SAMPLE_SIZE},
            )
            outputs.append({"label": label, "keys": rows[0]["keys"] if rows else []})
        return outputs

    def _sample_rel_properties(self) -> List[dict]:
        """Derive the property types of every relationship type in the graph."""
        rel_types = [
            row["relationshipType"]
            for row in self.structured_query("CALL db.relationshipTypes()") or []
            if row["relationshipType"] not in EXCLUDED_RELS
        ]

        outputs = []
        for rel_type in rel_types:
            rows = self.structured_query(
                f"""
                MATCH ()-[r:{escape_identifier(rel_type)}]->()
                WITH r LIMIT $SAMPLE_SIZE
                UNWIND keys(r) AS key
                WITH key, collect(r[key])[0] AS sample
                RETURN collect({{property: key, sample: sample}}) AS keys
                """,
                param_map={"SAMPLE_SIZE": SCHEMA_SAMPLE_SIZE},
            )
            outputs.append({"type": rel_type, "keys": rows[0]["keys"] if rows else []})
        return outputs

    def _sample_rel_triples(self) -> List[dict]:
        """
        Derive the (start label, type, end label) triples present in the graph.

        Relationship types are enumerated exhaustively and their endpoint labels
        are sampled per type. The obvious `MATCH (n)-[r]->(m)` instead traverses
        every relationship in the graph, and `PropertyGraphIndex` refreshes the
        schema after *every* ingestion batch, so that cost is paid over and over
        and grows with the graph.
        """
        triples: List[dict] = []
        seen: Set[Tuple[str, str, str]] = set()

        for row in self.structured_query("CALL db.relationshipTypes()") or []:
            rel_type = row["relationshipType"]
            if rel_type in EXCLUDED_RELS:
                continue

            rows = self.structured_query(
                f"""
                MATCH (n)-[r:{escape_identifier(rel_type)}]->(m)
                WITH n, m LIMIT $SAMPLE_SIZE
                UNWIND labels(n) AS start_label
                UNWIND labels(m) AS end_label
                RETURN DISTINCT start_label, end_label
                """,
                param_map={"SAMPLE_SIZE": SCHEMA_SAMPLE_SIZE},
            )
            for record in rows or []:
                key = (record["start_label"], rel_type, record["end_label"])
                if key in seen:
                    continue
                seen.add(key)
                triples.append(
                    {
                        "start": record["start_label"],
                        "type": rel_type,
                        "end": record["end_label"],
                    }
                )
        return triples

    def refresh_schema(self) -> None:
        """Refresh the schema."""
        node_properties = self._sample_node_properties()
        rel_properties = self._sample_rel_properties()
        relationships = self._sample_rel_triples()

        # Get constraints & indexes
        try:
            constraint = self.structured_query("CALL db.constraints()")
            index = self.structured_query(
                "CALL db.indexes() YIELD label, properties, entitytype RETURN *"
            )
        except (
            redis.exceptions.ResponseError
        ):  # Read-only user might not have access to schema information
            constraint = []
            index = []

        self.structured_schema = {
            "node_props": {
                el["label"]: self._format_properties(el["keys"])
                for el in node_properties
            },
            "rel_props": {
                el["type"]: self._format_properties(el["keys"]) for el in rel_properties
            },
            "relationships": relationships,
            "metadata": {"constraint": constraint, "index": index},
        }

    @staticmethod
    def _format_properties(keys: Optional[List[Any]]) -> List[Dict[str, str]]:
        """Turn sampled property values into ``{property, type}`` entries."""
        formatted = []
        for key in keys or []:
            if isinstance(key, dict):
                formatted.append(
                    {
                        "property": key.get("property"),
                        "type": sample_value_type(key.get("sample")),
                    }
                )
            else:
                # Defensive: older schemas stored bare property names
                formatted.append({"property": key, "type": "UNKNOWN"})
        return formatted

    def get_schema(self, refresh: bool = False) -> Any:
        if refresh:
            self.refresh_schema()

        return self.structured_schema

    def get_schema_str(self, refresh: bool = False) -> str:
        schema = self.get_schema(refresh=refresh)

        formatted_node_props = []
        formatted_rel_props = []

        # Format node properties
        for label, props in schema["node_props"].items():
            props_str = ", ".join(
                [f"{prop['property']}: {prop['type']}" for prop in props]
            )
            formatted_node_props.append(f"{label} {{{props_str}}}")

        # Format relationship properties using structured_schema
        for type, props in schema["rel_props"].items():
            props_str = ", ".join(
                [f"{prop['property']}: {prop['type']}" for prop in props]
            )
            formatted_rel_props.append(f"{type} {{{props_str}}}")

        # Format relationships
        formatted_rels = [
            f"(:{el['start']})-[:{el['type']}]->(:{el['end']})"
            for el in schema["relationships"]
        ]

        return "\n".join(
            [
                "Node properties:",
                "\n".join(formatted_node_props),
                "Relationship properties:",
                "\n".join(formatted_rel_props),
                "The relationships:",
                "\n".join(formatted_rels),
            ]
        )

    ### ----- Writes ----- ###

    def upsert_nodes(self, nodes: Sequence[LabelledNode]) -> None:
        # Lists to hold separated types
        entity_dicts: List[dict] = []
        chunk_dicts: List[dict] = []

        # Sort by type
        for item in nodes:
            if isinstance(item, EntityNode):
                entity_dicts.append({**item.model_dump(), "id": item.id})
            elif isinstance(item, ChunkNode):
                chunk_dicts.append({**item.model_dump(), "id": item.id})
            else:
                logger.warning(
                    "Unsupported node type `%s`, skipping.", type(item).__name__
                )

        if chunk_dicts:
            for batch in _batched(chunk_dicts):
                self.structured_query(
                    f"""
                    UNWIND $data AS row
                    MERGE (c:{escape_identifier(CHUNK_LABEL)} {{id: row.id}})
                    SET c.text = row.text
                    SET c += row.properties
                    SET c.{EMBEDDING_KEY} = CASE WHEN row.embedding IS NULL
                        THEN c.{EMBEDDING_KEY} ELSE vecf32(row.embedding) END
                    RETURN count(*)
                    """,
                    param_map={"data": batch},
                )

        if not entity_dicts:
            return

        for dimension in {
            len(entity["embedding"])
            for entity in entity_dicts
            if entity.get("embedding")
        }:
            self._ensure_vector_index(BASE_ENTITY_LABEL, dimension)

        # FalkorDB has no APOC, so labels cannot be parameterized. Group the
        # entities by label to keep a single statement per label instead of one
        # statement per node.
        entities_by_label: Dict[str, List[dict]] = defaultdict(list)
        for entity in entity_dicts:
            entities_by_label[entity["label"]].append(entity)

        for label, rows in entities_by_label.items():
            for batch in _batched(rows):
                self.structured_query(
                    f"""
                    UNWIND $data AS row
                    MERGE (e:{escape_identifier(BASE_ENTITY_LABEL)} {{id: row.id}})
                    SET e += row.properties
                    SET e.name = row.name
                    SET e:{escape_identifier(label)}
                    SET e.{EMBEDDING_KEY} = CASE WHEN row.embedding IS NULL
                        THEN e.{EMBEDDING_KEY} ELSE vecf32(row.embedding) END
                    RETURN count(*)
                    """,
                    param_map={"data": batch},
                )

        # Create MENTIONS relationships for entities with a triplet_source_id
        mentions = [
            {
                "entity_id": entity["id"],
                "chunk_id": entity["properties"]["triplet_source_id"],
            }
            for entity in entity_dicts
            if (entity.get("properties") or {}).get("triplet_source_id")
        ]
        for batch in _batched(mentions):
            self.structured_query(
                f"""
                UNWIND $data AS row
                MATCH (e:{escape_identifier(BASE_ENTITY_LABEL)} {{id: row.entity_id}})
                MERGE (c:{escape_identifier(CHUNK_LABEL)} {{id: row.chunk_id}})
                MERGE (e)<-[:MENTIONS]-(c)
                RETURN count(*)
                """,
                param_map={"data": batch},
            )

    def _existing_chunk_ids(self, ids: List[str]) -> Set[str]:
        """Return the subset of ``ids`` that already exist as chunk nodes."""
        found: Set[str] = set()
        for batch in _batched([{"id": node_id} for node_id in ids]):
            rows = self.structured_query(
                f"""
                UNWIND $data AS row
                MATCH (c:{escape_identifier(CHUNK_LABEL)} {{id: row.id}})
                RETURN c.id AS id
                """,
                param_map={"data": batch},
            )
            found.update(row["id"] for row in rows or [])
        return found

    def upsert_relations(self, relations: List[Relation]) -> None:
        """Add relations."""
        if not relations:
            return

        # An endpoint that does not exist yet is created with BASE_ENTITY_LABEL
        # so that a later `upsert_nodes` MERGE matches this very node instead of
        # inserting a second one with the same id. Endpoints that already exist
        # as chunks keep their label for the same reason.
        endpoint_ids = sorted(
            {relation.source_id for relation in relations}
            | {relation.target_id for relation in relations}
        )
        chunk_ids = self._existing_chunk_ids(endpoint_ids)

        def label_for(node_id: str) -> str:
            return CHUNK_LABEL if node_id in chunk_ids else BASE_ENTITY_LABEL

        # Neither relationship types nor labels can be parameterized, so group
        # the rows by every identifier the statement has to interpolate. Merging
        # on a labelled pattern is also what lets FalkorDB use the `id` indexes
        # instead of scanning every node in the graph for each row.
        grouped: Dict[Tuple[str, str, str], List[dict]] = defaultdict(list)
        for relation in relations:
            key = (
                relation.label,
                label_for(relation.source_id),
                label_for(relation.target_id),
            )
            grouped[key].append(relation.model_dump())

        for (label, source_label, target_label), rows in grouped.items():
            for batch in _batched(rows):
                self.structured_query(
                    f"""
                    UNWIND $data AS row
                    MERGE (source:{escape_identifier(source_label)} {{id: row.source_id}})
                    MERGE (target:{escape_identifier(target_label)} {{id: row.target_id}})
                    MERGE (source)-[r:{escape_identifier(label)}]->(target)
                    SET r += row.properties
                    RETURN count(*)
                    """,
                    param_map={"data": batch},
                )

    ### ----- Reads ----- ###

    def _match_store_nodes(
        self, conditions: List[str], params: Dict[str, Any]
    ) -> List[dict]:
        """
        Run a node lookup once per store-managed label and merge the results.

        Every node this store writes carries either ``BASE_ENTITY_LABEL`` or
        ``CHUNK_LABEL``. Scoping the match to those labels is what lets FalkorDB
        use the `id` range indexes: its indexes are label-scoped, so an
        unlabelled `MATCH (e)` degrades to an `All Node Scan` of the whole
        graph. That matters because `PropertyGraphIndex` calls `get()` twice per
        ingestion batch to deduplicate, making ingestion quadratic.
        """
        where = f"WHERE {' AND '.join(conditions)} " if conditions else ""
        return_statement = f"""
        WITH e
        RETURN e.id AS name,
               [l in labels(e) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS type,
               e{{.* , embedding: Null, id: Null}} AS properties
        """

        records: List[dict] = []
        seen: Set[str] = set()
        for label in STORE_NODE_LABELS:
            rows = self.structured_query(
                f"MATCH (e:{escape_identifier(label)}) {where}{return_statement}",
                param_map=params,
            )
            for record in rows or []:
                # A node carrying both labels would otherwise be returned twice.
                if record["name"] in seen:
                    continue
                seen.add(record["name"])
                records.append(record)
        return records

    def get(
        self,
        properties: Optional[dict] = None,
        ids: Optional[List[str]] = None,
    ) -> List[LabelledNode]:
        """Get nodes."""
        params: Dict[str, Any] = {}
        conditions: List[str] = []

        if ids:
            conditions.append("e.id IN $ids")
            params["ids"] = ids

        if properties:
            for i, prop in enumerate(properties):
                conditions.append(f"e.{escape_identifier(prop)} = $property_{i}")
                params[f"property_{i}"] = properties[prop]

        response = self._match_store_nodes(conditions, params)

        nodes = []
        for record in response:
            # text indicates a chunk node
            # none on the type indicates an implicit node, likely a chunk node
            if "text" in record["properties"] or record["type"] is None:
                text = record["properties"].pop("text", "")
                nodes.append(
                    ChunkNode(
                        id_=record["name"],
                        text=text,
                        properties=remove_empty_values(record["properties"]),
                    )
                )
            else:
                nodes.append(
                    EntityNode(
                        name=record["name"],
                        label=record["type"],
                        properties=remove_empty_values(record["properties"]),
                    )
                )

        return nodes

    def get_triplets(
        self,
        entity_names: Optional[List[str]] = None,
        relation_names: Optional[List[str]] = None,
        properties: Optional[dict] = None,
        ids: Optional[List[str]] = None,
    ) -> List[Triplet]:
        # TODO: handle ids of chunk nodes
        cypher_statement = f"MATCH (e:{escape_identifier(BASE_ENTITY_LABEL)}) "

        params: Dict[str, Any] = {}
        conditions: List[str] = []

        if entity_names:
            conditions.append("e.name IN $entity_names")
            params["entity_names"] = entity_names

        if ids:
            conditions.append("e.id IN $ids")
            params["ids"] = ids

        if properties:
            for i, prop in enumerate(properties):
                conditions.append(f"e.{escape_identifier(prop)} = $property_{i}")
                params[f"property_{i}"] = properties[prop]

        if conditions:
            cypher_statement += "WHERE " + " AND ".join(conditions) + " "

        relation_filter = (
            ":" + "|".join(escape_identifier(rel) for rel in relation_names)
            if relation_names
            else ""
        )

        base = escape_identifier(BASE_ENTITY_LABEL)
        return_statement = f"""
        WITH e
        CALL {{
            WITH e
            MATCH (e)-[r{relation_filter}]->(t:{base})
            RETURN e.name AS source_id, [l in labels(e) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS source_type,
                   e{{.* , embedding: Null, name: Null}} AS source_properties,
                   type(r) AS type, r{{.*}} AS rel_properties,
                   t.name AS target_id, [l in labels(t) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS target_type,
                   t{{.* , embedding: Null, name: Null}} AS target_properties
            UNION ALL
            WITH e
            MATCH (e)<-[r{relation_filter}]-(t:{base})
            RETURN t.name AS source_id, [l in labels(t) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS source_type,
                   t{{.* , embedding: Null, name: Null}} AS source_properties,
                   type(r) AS type, r{{.*}} AS rel_properties,
                   e.name AS target_id, [l in labels(e) WHERE l <> '{BASE_ENTITY_LABEL}' | l][0] AS target_type,
                   e{{.* , embedding: Null, name: Null}} AS target_properties
        }}
        RETURN source_id, source_type, type, rel_properties, target_id, target_type, source_properties, target_properties"""
        cypher_statement += return_statement

        data = self.structured_query(cypher_statement, param_map=params)
        data = data if data else []

        triples = []
        for record in data:
            source = EntityNode(
                name=record["source_id"],
                label=record["source_type"],
                properties=remove_empty_values(record["source_properties"]),
            )
            target = EntityNode(
                name=record["target_id"],
                label=record["target_type"],
                properties=remove_empty_values(record["target_properties"]),
            )
            rel = Relation(
                source_id=record["source_id"],
                target_id=record["target_id"],
                label=record["type"],
                properties=remove_empty_values(record["rel_properties"] or {}),
            )
            triples.append([source, rel, target])
        return triples

    def get_rel_map(
        self,
        graph_nodes: List[LabelledNode],
        depth: int = 2,
        limit: int = 30,
        ignore_rels: Optional[List[str]] = None,
    ) -> List[Triplet]:
        """Get depth-aware rel map."""
        triples = []

        ids = [node.id for node in graph_nodes]
        if not ids:
            return triples

        # Filter the ignored relationships server side so that they do not eat
        # into the `limit` budget.
        ignored = list({*(ignore_rels or [])})

        response = self.structured_query(
            f"""
            WITH $ids AS id_list
            UNWIND range(0, size(id_list) - 1) AS idx
            MATCH (e:{escape_identifier(BASE_ENTITY_LABEL)})
            WHERE e.id = id_list[idx]
            MATCH p=(e)-[r*1..{int(depth)}]-(other)
            WHERE ALL(rel in relationships(p) WHERE type(rel) <> 'MENTIONS')
            UNWIND relationships(p) AS rel
            WITH distinct rel, idx
            WHERE NOT type(rel) IN $ignore_rels
            WITH startNode(rel) AS source,
                type(rel) AS type,
                rel{{.*}} AS rel_properties,
                endNode(rel) AS endNode,
                idx
            LIMIT $limit
            RETURN source.id AS source_id, [l in labels(source) WHERE l <> '__Entity__' | l][0] AS source_type,
                source{{.* , embedding: Null, id: Null}} AS source_properties,
                type,
                rel_properties,
                endNode.id AS target_id, [l in labels(endNode) WHERE l <> '__Entity__' | l][0] AS target_type,
                endNode{{.* , embedding: Null, id: Null}} AS target_properties,
                idx
            ORDER BY idx
            LIMIT $limit
            """,
            param_map={"ids": ids, "limit": int(limit), "ignore_rels": ignored},
        )
        response = response if response else []

        for record in response:
            source = EntityNode(
                name=record["source_id"],
                label=record["source_type"],
                properties=remove_empty_values(record["source_properties"]),
            )
            target = EntityNode(
                name=record["target_id"],
                label=record["target_type"],
                properties=remove_empty_values(record["target_properties"]),
            )
            rel = Relation(
                source_id=record["source_id"],
                target_id=record["target_id"],
                label=record["type"],
                properties=remove_empty_values(record["rel_properties"] or {}),
            )
            triples.append([source, rel, target])

        return triples

    def structured_query(
        self, query: str, param_map: Optional[Dict[str, Any]] = None
    ) -> Any:
        param_map = param_map or {}

        result = self._graph.query(query, param_map, timeout=self._timeout)
        full_result = [
            {
                header[1]: to_plain_value(value)
                for header, value in zip(result.header, row)
            }
            for row in result.result_set
        ]

        if self.sanitize_query_output:
            return [value_sanitize(el) for el in full_result]
        return full_result

    def vector_query(
        self, query: VectorStoreQuery, **kwargs: Any
    ) -> Tuple[List[LabelledNode], List[float]]:
        """Query the graph store with a vector store query."""
        conditions = []
        filter_params: Dict[str, Any] = {}
        if query.filters:
            for index, filter_ in enumerate(query.filters.filters):
                conditions.append(
                    f"{'NOT ' if filter_.operator.value == 'nin' else ''}"
                    f"e.{escape_identifier(filter_.key)} "
                    f"{convert_operator(filter_.operator.value)} $param_{index}"
                )
                filter_params[f"param_{index}"] = filter_.value
        filters = (
            f" {query.filters.condition.value} ".join(conditions)
            if conditions
            else "1 = 1"
        )

        dimension = len(query.query_embedding or [])
        use_vector_index = (
            not conditions
            and self._vector_index_dimensions.get(BASE_ENTITY_LABEL) == dimension
        )

        if use_vector_index:
            data = self.structured_query(
                f"""CALL db.idx.vector.queryNodes(
                        '{BASE_ENTITY_LABEL}', '{EMBEDDING_KEY}',
                        $limit, vecf32($embedding))
                YIELD node AS e, score
                RETURN e.id AS name,
                    [l in labels(e) WHERE l <> '__Entity__' | l][0] AS type,
                    e{{.* , embedding: Null, name: Null, id: Null}} AS properties,
                    1 - score AS score
                ORDER BY score DESC""",
                param_map={
                    "embedding": query.query_embedding,
                    "limit": int(query.similarity_top_k),
                },
            )
        else:
            try:
                data = self.structured_query(
                    f"""MATCH (e:{escape_identifier(BASE_ENTITY_LABEL)})
                    WHERE e.{EMBEDDING_KEY} IS NOT NULL AND ({filters})
                    WITH e, 1 - vec.cosineDistance(e.{EMBEDDING_KEY}, vecf32($embedding)) AS score
                    ORDER BY score DESC LIMIT $limit
                    RETURN e.id AS name,
                    [l in labels(e) WHERE l <> '__Entity__' | l][0] AS type,
                    e{{.* , embedding: Null, name: Null, id: Null}} AS properties,
                    score""",
                    param_map={
                        "embedding": query.query_embedding,
                        "dimension": dimension,
                        "limit": int(query.similarity_top_k),
                        **filter_params,
                    },
                )
            except redis.exceptions.ResponseError as e:
                if "dimension mismatch" in str(e).lower():
                    raise ValueError(
                        "The graph contains embeddings whose dimension differs from "
                        f"the query embedding ({dimension}). Make sure a single "
                        "embedding model is used for the whole graph."
                    ) from e
                raise

        data = data if data else []

        nodes = []
        scores = []
        for record in data:
            node = EntityNode(
                name=record["name"],
                label=record["type"],
                properties=remove_empty_values(record["properties"]),
            )
            nodes.append(node)
            scores.append(record["score"])

        return (nodes, scores)

    def delete(
        self,
        entity_names: Optional[List[str]] = None,
        relation_names: Optional[List[str]] = None,
        properties: Optional[dict] = None,
        ids: Optional[List[str]] = None,
    ) -> None:
        """Delete matching data."""

        # As in `get()`, the match has to name a label for FalkorDB to be able
        # to use its (label-scoped) indexes instead of scanning the graph.
        def delete_by(condition: str, param_map: Dict[str, Any]) -> None:
            for label in STORE_NODE_LABELS:
                self.structured_query(
                    f"MATCH (e:{escape_identifier(label)}) WHERE {condition} "
                    "DETACH DELETE e",
                    param_map=param_map,
                )

        if entity_names:
            delete_by("e.name IN $entity_names", {"entity_names": entity_names})

        if ids:
            delete_by("e.id IN $ids", {"ids": ids})

        if relation_names:
            for rel in relation_names:
                self.structured_query(
                    f"MATCH ()-[r:{escape_identifier(rel)}]->() DELETE r"
                )

        if properties:
            prop_list = []
            params: Dict[str, Any] = {}
            for i, prop in enumerate(properties):
                prop_list.append(f"e.{escape_identifier(prop)} = $property_{i}")
                params[f"property_{i}"] = properties[prop]
            delete_by(" AND ".join(prop_list), params)

    ### ----- Connection management ----- ###

    def switch_graph(self, graph_name: str) -> None:
        """
        Switch to the given graph name (`graph_name`).

        This method allows users to change the active graph within the same
        database connection.

        Args:
            graph_name (str): The name of the graph to switch to.

        """
        self._graph = self._driver.select_graph(graph_name)
        self._database = graph_name

        if self._create_indexes:
            self._create_range_indexes()
        self._refresh_vector_index_info()

        try:
            self.refresh_schema()
        except Exception as e:
            raise ValueError(f"Could not refresh schema. Error: {e}")

    def close(self) -> None:
        """Explicitly close the FalkorDB connection."""
        if hasattr(self, "_driver"):
            try:
                self._driver.connection.close()
            finally:
                delattr(self, "_driver")

    def __enter__(self) -> "FalkorDBPropertyGraphStore":
        """Enter the runtime context, enabling `with FalkorDBPropertyGraphStore(...)`."""
        return self

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        """Close the connection when leaving the runtime context."""
        self.close()

    def __del__(self) -> None:
        """Best-effort cleanup; prefer `close()` or the context manager."""
        try:
            self.close()
        except Exception:
            # Suppress any exceptions during garbage collection
            pass

refresh_schema #

refresh_schema() -> None

Refresh the schema.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/falkordb_property_graph.py 
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
def refresh_schema(self) -> None:
    """Refresh the schema."""
    node_properties = self._sample_node_properties()
    rel_properties = self._sample_rel_properties()
    relationships = self._sample_rel_triples()

    # Get constraints & indexes
    try:
        constraint = self.structured_query("CALL db.constraints()")
        index = self.structured_query(
            "CALL db.indexes() YIELD label, properties, entitytype RETURN *"
        )
    except (
        redis.exceptions.ResponseError
    ):  # Read-only user might not have access to schema information
        constraint = []
        index = []

    self.structured_schema = {
        "node_props": {
            el["label"]: self._format_properties(el["keys"])
            for el in node_properties
        },
        "rel_props": {
            el["type"]: self._format_properties(el["keys"]) for el in rel_properties
        },
        "relationships": relationships,
        "metadata": {"constraint": constraint, "index": index},
    }

upsert_relations #

upsert_relations(relations: List[Relation]) -> None

Add relations.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/falkordb_property_graph.py 
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
def upsert_relations(self, relations: List[Relation]) -> None:
    """Add relations."""
    if not relations:
        return

    # An endpoint that does not exist yet is created with BASE_ENTITY_LABEL
    # so that a later `upsert_nodes` MERGE matches this very node instead of
    # inserting a second one with the same id. Endpoints that already exist
    # as chunks keep their label for the same reason.
    endpoint_ids = sorted(
        {relation.source_id for relation in relations}
        | {relation.target_id for relation in relations}
    )
    chunk_ids = self._existing_chunk_ids(endpoint_ids)

    def label_for(node_id: str) -> str:
        return CHUNK_LABEL if node_id in chunk_ids else BASE_ENTITY_LABEL

    # Neither relationship types nor labels can be parameterized, so group
    # the rows by every identifier the statement has to interpolate. Merging
    # on a labelled pattern is also what lets FalkorDB use the `id` indexes
    # instead of scanning every node in the graph for each row.
    grouped: Dict[Tuple[str, str, str], List[dict]] = defaultdict(list)
    for relation in relations:
        key = (
            relation.label,
            label_for(relation.source_id),
            label_for(relation.target_id),
        )
        grouped[key].append(relation.model_dump())

    for (label, source_label, target_label), rows in grouped.items():
        for batch in _batched(rows):
            self.structured_query(
                f"""
                UNWIND $data AS row
                MERGE (source:{escape_identifier(source_label)} {{id: row.source_id}})
                MERGE (target:{escape_identifier(target_label)} {{id: row.target_id}})
                MERGE (source)-[r:{escape_identifier(label)}]->(target)
                SET r += row.properties
                RETURN count(*)
                """,
                param_map={"data": batch},
            )

get #

get(
    properties: Optional[dict] = None,
    ids: Optional[List[str]] = None,
) -> List[LabelledNode]

Get nodes.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/falkordb_property_graph.py 
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
def get(
    self,
    properties: Optional[dict] = None,
    ids: Optional[List[str]] = None,
) -> List[LabelledNode]:
    """Get nodes."""
    params: Dict[str, Any] = {}
    conditions: List[str] = []

    if ids:
        conditions.append("e.id IN $ids")
        params["ids"] = ids

    if properties:
        for i, prop in enumerate(properties):
            conditions.append(f"e.{escape_identifier(prop)} = $property_{i}")
            params[f"property_{i}"] = properties[prop]

    response = self._match_store_nodes(conditions, params)

    nodes = []
    for record in response:
        # text indicates a chunk node
        # none on the type indicates an implicit node, likely a chunk node
        if "text" in record["properties"] or record["type"] is None:
            text = record["properties"].pop("text", "")
            nodes.append(
                ChunkNode(
                    id_=record["name"],
                    text=text,
                    properties=remove_empty_values(record["properties"]),
                )
            )
        else:
            nodes.append(
                EntityNode(
                    name=record["name"],
                    label=record["type"],
                    properties=remove_empty_values(record["properties"]),
                )
            )

    return nodes

get_rel_map #

get_rel_map(
    graph_nodes: List[LabelledNode],
    depth: int = 2,
    limit: int = 30,
    ignore_rels: Optional[List[str]] = None,
) -> List[Triplet]

Get depth-aware rel map.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/falkordb_property_graph.py 
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
def get_rel_map(
    self,
    graph_nodes: List[LabelledNode],
    depth: int = 2,
    limit: int = 30,
    ignore_rels: Optional[List[str]] = None,
) -> List[Triplet]:
    """Get depth-aware rel map."""
    triples = []

    ids = [node.id for node in graph_nodes]
    if not ids:
        return triples

    # Filter the ignored relationships server side so that they do not eat
    # into the `limit` budget.
    ignored = list({*(ignore_rels or [])})

    response = self.structured_query(
        f"""
        WITH $ids AS id_list
        UNWIND range(0, size(id_list) - 1) AS idx
        MATCH (e:{escape_identifier(BASE_ENTITY_LABEL)})
        WHERE e.id = id_list[idx]
        MATCH p=(e)-[r*1..{int(depth)}]-(other)
        WHERE ALL(rel in relationships(p) WHERE type(rel) <> 'MENTIONS')
        UNWIND relationships(p) AS rel
        WITH distinct rel, idx
        WHERE NOT type(rel) IN $ignore_rels
        WITH startNode(rel) AS source,
            type(rel) AS type,
            rel{{.*}} AS rel_properties,
            endNode(rel) AS endNode,
            idx
        LIMIT $limit
        RETURN source.id AS source_id, [l in labels(source) WHERE l <> '__Entity__' | l][0] AS source_type,
            source{{.* , embedding: Null, id: Null}} AS source_properties,
            type,
            rel_properties,
            endNode.id AS target_id, [l in labels(endNode) WHERE l <> '__Entity__' | l][0] AS target_type,
            endNode{{.* , embedding: Null, id: Null}} AS target_properties,
            idx
        ORDER BY idx
        LIMIT $limit
        """,
        param_map={"ids": ids, "limit": int(limit), "ignore_rels": ignored},
    )
    response = response if response else []

    for record in response:
        source = EntityNode(
            name=record["source_id"],
            label=record["source_type"],
            properties=remove_empty_values(record["source_properties"]),
        )
        target = EntityNode(
            name=record["target_id"],
            label=record["target_type"],
            properties=remove_empty_values(record["target_properties"]),
        )
        rel = Relation(
            source_id=record["source_id"],
            target_id=record["target_id"],
            label=record["type"],
            properties=remove_empty_values(record["rel_properties"] or {}),
        )
        triples.append([source, rel, target])

    return triples

vector_query #

vector_query(
    query: VectorStoreQuery, **kwargs: Any
) -> Tuple[List[LabelledNode], List[float]]

Query the graph store with a vector store query.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/falkordb_property_graph.py 
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
def vector_query(
    self, query: VectorStoreQuery, **kwargs: Any
) -> Tuple[List[LabelledNode], List[float]]:
    """Query the graph store with a vector store query."""
    conditions = []
    filter_params: Dict[str, Any] = {}
    if query.filters:
        for index, filter_ in enumerate(query.filters.filters):
            conditions.append(
                f"{'NOT ' if filter_.operator.value == 'nin' else ''}"
                f"e.{escape_identifier(filter_.key)} "
                f"{convert_operator(filter_.operator.value)} $param_{index}"
            )
            filter_params[f"param_{index}"] = filter_.value
    filters = (
        f" {query.filters.condition.value} ".join(conditions)
        if conditions
        else "1 = 1"
    )

    dimension = len(query.query_embedding or [])
    use_vector_index = (
        not conditions
        and self._vector_index_dimensions.get(BASE_ENTITY_LABEL) == dimension
    )

    if use_vector_index:
        data = self.structured_query(
            f"""CALL db.idx.vector.queryNodes(
                    '{BASE_ENTITY_LABEL}', '{EMBEDDING_KEY}',
                    $limit, vecf32($embedding))
            YIELD node AS e, score
            RETURN e.id AS name,
                [l in labels(e) WHERE l <> '__Entity__' | l][0] AS type,
                e{{.* , embedding: Null, name: Null, id: Null}} AS properties,
                1 - score AS score
            ORDER BY score DESC""",
            param_map={
                "embedding": query.query_embedding,
                "limit": int(query.similarity_top_k),
            },
        )
    else:
        try:
            data = self.structured_query(
                f"""MATCH (e:{escape_identifier(BASE_ENTITY_LABEL)})
                WHERE e.{EMBEDDING_KEY} IS NOT NULL AND ({filters})
                WITH e, 1 - vec.cosineDistance(e.{EMBEDDING_KEY}, vecf32($embedding)) AS score
                ORDER BY score DESC LIMIT $limit
                RETURN e.id AS name,
                [l in labels(e) WHERE l <> '__Entity__' | l][0] AS type,
                e{{.* , embedding: Null, name: Null, id: Null}} AS properties,
                score""",
                param_map={
                    "embedding": query.query_embedding,
                    "dimension": dimension,
                    "limit": int(query.similarity_top_k),
                    **filter_params,
                },
            )
        except redis.exceptions.ResponseError as e:
            if "dimension mismatch" in str(e).lower():
                raise ValueError(
                    "The graph contains embeddings whose dimension differs from "
                    f"the query embedding ({dimension}). Make sure a single "
                    "embedding model is used for the whole graph."
                ) from e
            raise

    data = data if data else []

    nodes = []
    scores = []
    for record in data:
        node = EntityNode(
            name=record["name"],
            label=record["type"],
            properties=remove_empty_values(record["properties"]),
        )
        nodes.append(node)
        scores.append(record["score"])

    return (nodes, scores)

delete #

delete(
    entity_names: Optional[List[str]] = None,
    relation_names: Optional[List[str]] = None,
    properties: Optional[dict] = None,
    ids: Optional[List[str]] = None,
) -> None

Delete matching data.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/falkordb_property_graph.py 
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
def delete(
    self,
    entity_names: Optional[List[str]] = None,
    relation_names: Optional[List[str]] = None,
    properties: Optional[dict] = None,
    ids: Optional[List[str]] = None,
) -> None:
    """Delete matching data."""

    # As in `get()`, the match has to name a label for FalkorDB to be able
    # to use its (label-scoped) indexes instead of scanning the graph.
    def delete_by(condition: str, param_map: Dict[str, Any]) -> None:
        for label in STORE_NODE_LABELS:
            self.structured_query(
                f"MATCH (e:{escape_identifier(label)}) WHERE {condition} "
                "DETACH DELETE e",
                param_map=param_map,
            )

    if entity_names:
        delete_by("e.name IN $entity_names", {"entity_names": entity_names})

    if ids:
        delete_by("e.id IN $ids", {"ids": ids})

    if relation_names:
        for rel in relation_names:
            self.structured_query(
                f"MATCH ()-[r:{escape_identifier(rel)}]->() DELETE r"
            )

    if properties:
        prop_list = []
        params: Dict[str, Any] = {}
        for i, prop in enumerate(properties):
            prop_list.append(f"e.{escape_identifier(prop)} = $property_{i}")
            params[f"property_{i}"] = properties[prop]
        delete_by(" AND ".join(prop_list), params)

switch_graph #

switch_graph(graph_name: str) -> None

Switch to the given graph name (graph_name).

This method allows users to change the active graph within the same database connection.

Parameters:

Name Type Description Default
graph_name str

The name of the graph to switch to.

 required
Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/falkordb_property_graph.py 
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
def switch_graph(self, graph_name: str) -> None:
    """
    Switch to the given graph name (`graph_name`).

    This method allows users to change the active graph within the same
    database connection.

    Args:
        graph_name (str): The name of the graph to switch to.

    """
    self._graph = self._driver.select_graph(graph_name)
    self._database = graph_name

    if self._create_indexes:
        self._create_range_indexes()
    self._refresh_vector_index_info()

    try:
        self.refresh_schema()
    except Exception as e:
        raise ValueError(f"Could not refresh schema. Error: {e}")

close #

close() -> None

Explicitly close the FalkorDB connection.

Source code in llama-index-integrations/graph_stores/llama-index-graph-stores-falkordb/llama_index/graph_stores/falkordb/falkordb_property_graph.py 
1026
1027
1028
1029
1030
1031
1032
def close(self) -> None:
    """Explicitly close the FalkorDB connection."""
    if hasattr(self, "_driver"):
        try:
            self._driver.connection.close()
        finally:
            delattr(self, "_driver")

options: members: - FalkorDBGraphStore - FalkorDBPropertyGraphStore