graphiti-core 0.18.9__py3-none-any.whl → 0.19.0__py3-none-any.whl

graphiti_core/search/search.py

@@ -21,6 +21,7 @@ from time import time
 from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.edges import EntityEdge
+from graphiti_core.embedder.client import EMBEDDING_DIM
 from graphiti_core.errors import SearchRerankerError
 from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import semaphore_gather
@@ -29,6 +30,7 @@ from graphiti_core.search.search_config import (
     DEFAULT_SEARCH_LIMIT,
     CommunityReranker,
     CommunitySearchConfig,
+    CommunitySearchMethod,
     EdgeReranker,
     EdgeSearchConfig,
     EdgeSearchMethod,
@@ -81,11 +83,29 @@ async def search(
 
     if query.strip() == '':
         return SearchResults()
-    query_vector = (
-        query_vector
-        if query_vector is not None
-        else await embedder.create(input_data=[query.replace('\n', ' ')])
-    )
+
+    if (
+        config.edge_config
+        and EdgeSearchMethod.cosine_similarity in config.edge_config.search_methods
+        or config.edge_config
+        and EdgeReranker.mmr == config.edge_config.reranker
+        or config.node_config
+        and NodeSearchMethod.cosine_similarity in config.node_config.search_methods
+        or config.node_config
+        and NodeReranker.mmr == config.node_config.reranker
+        or (
+            config.community_config
+            and CommunitySearchMethod.cosine_similarity in config.community_config.search_methods
+        )
+        or (config.community_config and CommunityReranker.mmr == config.community_config.reranker)
+    ):
+        search_vector = (
+            query_vector
+            if query_vector is not None
+            else await embedder.create(input_data=[query.replace('\n', ' ')])
+        )
+    else:
+        search_vector = [0.0] * EMBEDDING_DIM
 
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids and group_ids != [''] else None
@@ -99,7 +119,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.edge_config,
             search_filter,
@@ -112,7 +132,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.node_config,
             search_filter,
@@ -125,7 +145,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.episode_config,
             search_filter,
@@ -136,7 +156,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.community_config,
             config.limit,
@@ -305,12 +325,20 @@ async def node_search(
     search_tasks = []
     if NodeSearchMethod.bm25 in config.search_methods:
         search_tasks.append(
-            node_fulltext_search(driver, query, search_filter, group_ids, 2 * limit)
+            node_fulltext_search(
+                driver, query, search_filter, group_ids, 2 * limit, config.use_local_indexes
+            )
         )
     if NodeSearchMethod.cosine_similarity in config.search_methods:
         search_tasks.append(
             node_similarity_search(
-                driver, query_vector, search_filter, group_ids, 2 * limit, config.sim_min_score
+                driver,
+                query_vector,
+                search_filter,
+                group_ids,
+                2 * limit,
+                config.sim_min_score,
+                config.use_local_indexes,
             )
         )
     if NodeSearchMethod.bfs in config.search_methods:
@@ -406,7 +434,9 @@ async def episode_search(
     search_results: list[list[EpisodicNode]] = list(
         await semaphore_gather(
             *[
-                episode_fulltext_search(driver, query, search_filter, group_ids, 2 * limit),
+                episode_fulltext_search(
+                    driver, query, search_filter, group_ids, 2 * limit, config.use_local_indexes
+                ),
             ]
         )
     )

graphiti_core/search/search_config.py

@@ -24,6 +24,7 @@ from graphiti_core.search.search_utils import (
     DEFAULT_MIN_SCORE,
     DEFAULT_MMR_LAMBDA,
     MAX_SEARCH_DEPTH,
+    USE_HNSW,
 )
 
 DEFAULT_SEARCH_LIMIT = 10
@@ -91,6 +92,7 @@ class NodeSearchConfig(BaseModel):
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
     bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
+    use_local_indexes: bool = Field(default=USE_HNSW)
 
 
 class EpisodeSearchConfig(BaseModel):
@@ -99,6 +101,7 @@ class EpisodeSearchConfig(BaseModel):
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
     bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
+    use_local_indexes: bool = Field(default=USE_HNSW)
 
 
 class CommunitySearchConfig(BaseModel):
@@ -107,6 +110,7 @@ class CommunitySearchConfig(BaseModel):
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
     bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
+    use_local_indexes: bool = Field(default=USE_HNSW)
 
 
 class SearchConfig(BaseModel):

graphiti_core/search/search_filters.py

@@ -20,6 +20,8 @@ from typing import Any
 
 from pydantic import BaseModel, Field
 
+from graphiti_core.driver.driver import GraphProvider
+
 
 class ComparisonOperator(Enum):
     equals = '='
@@ -54,16 +56,21 @@ class SearchFilters(BaseModel):
 
 def node_search_filter_query_constructor(
     filters: SearchFilters,
-) -> tuple[str, dict[str, Any]]:
-    filter_query: str = ''
+    provider: GraphProvider,
+) -> tuple[list[str], dict[str, Any]]:
+    filter_queries: list[str] = []
     filter_params: dict[str, Any] = {}
 
     if filters.node_labels is not None:
-        node_labels = '|'.join(filters.node_labels)
-        node_label_filter = ' AND n:' + node_labels
-        filter_query += node_label_filter
+        if provider == GraphProvider.KUZU:
+            node_label_filter = 'list_has_all(n.labels, $labels)'
+            filter_params['labels'] = filters.node_labels
+        else:
+            node_labels = '|'.join(filters.node_labels)
+            node_label_filter = 'n:' + node_labels
+        filter_queries.append(node_label_filter)
 
-    return filter_query, filter_params
+    return filter_queries, filter_params
 
 
 def date_filter_query_constructor(
@@ -81,23 +88,29 @@ def date_filter_query_constructor(
 
 def edge_search_filter_query_constructor(
     filters: SearchFilters,
-) -> tuple[str, dict[str, Any]]:
-    filter_query: str = ''
+    provider: GraphProvider,
+) -> tuple[list[str], dict[str, Any]]:
+    filter_queries: list[str] = []
     filter_params: dict[str, Any] = {}
 
     if filters.edge_types is not None:
         edge_types = filters.edge_types
-        edge_types_filter = '\nAND e.name in $edge_types'
-        filter_query += edge_types_filter
+        filter_queries.append('e.name in $edge_types')
         filter_params['edge_types'] = edge_types
 
     if filters.node_labels is not None:
-        node_labels = '|'.join(filters.node_labels)
-        node_label_filter = '\nAND n:' + node_labels + ' AND m:' + node_labels
-        filter_query += node_label_filter
+        if provider == GraphProvider.KUZU:
+            node_label_filter = (
+                'list_has_all(n.labels, $labels) AND list_has_all(m.labels, $labels)'
+            )
+            filter_params['labels'] = filters.node_labels
+        else:
+            node_labels = '|'.join(filters.node_labels)
+            node_label_filter = 'n:' + node_labels + ' AND m:' + node_labels
+        filter_queries.append(node_label_filter)
 
     if filters.valid_at is not None:
-        valid_at_filter = '\nAND ('
+        valid_at_filter = '('
         for i, or_list in enumerate(filters.valid_at):
             for j, date_filter in enumerate(or_list):
                 if date_filter.comparison_operator not in [
@@ -125,10 +138,10 @@ def edge_search_filter_query_constructor(
             else:
                 valid_at_filter += ' OR '
 
-        filter_query += valid_at_filter
+        filter_queries.append(valid_at_filter)
 
     if filters.invalid_at is not None:
-        invalid_at_filter = ' AND ('
+        invalid_at_filter = '('
         for i, or_list in enumerate(filters.invalid_at):
             for j, date_filter in enumerate(or_list):
                 if date_filter.comparison_operator not in [
@@ -156,10 +169,10 @@ def edge_search_filter_query_constructor(
             else:
                 invalid_at_filter += ' OR '
 
-        filter_query += invalid_at_filter
+        filter_queries.append(invalid_at_filter)
 
     if filters.created_at is not None:
-        created_at_filter = ' AND ('
+        created_at_filter = '('
         for i, or_list in enumerate(filters.created_at):
             for j, date_filter in enumerate(or_list):
                 if date_filter.comparison_operator not in [
@@ -187,10 +200,10 @@ def edge_search_filter_query_constructor(
             else:
                 created_at_filter += ' OR '
 
-        filter_query += created_at_filter
+        filter_queries.append(created_at_filter)
 
     if filters.expired_at is not None:
-        expired_at_filter = ' AND ('
+        expired_at_filter = '('
         for i, or_list in enumerate(filters.expired_at):
             for j, date_filter in enumerate(or_list):
                 if date_filter.comparison_operator not in [
@@ -218,6 +231,6 @@ def edge_search_filter_query_constructor(
             else:
                 expired_at_filter += ' OR '
 
-        filter_query += expired_at_filter
+        filter_queries.append(expired_at_filter)
 
-    return filter_query, filter_params
+    return filter_queries, filter_params