graphiti-core 0.3.8__tar.gz → 0.3.9__tar.gz

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.3.8
+Version: 0.3.9
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -17,6 +17,7 @@ Requires-Dist: numpy (>=1.0.0)
 Requires-Dist: openai (>=1.50.2,<2.0.0)
 Requires-Dist: pydantic (>=2.8.2,<3.0.0)
 Requires-Dist: tenacity (<9.0.0)
+Requires-Dist: voyageai (>=0.2.3,<0.3.0)
 Description-Content-Type: text/markdown
 
 <div align="center">

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/edges.py

@@ -188,9 +188,9 @@ class EntityEdge(Edge):
         MATCH (source:Entity {uuid: $source_uuid}) 
         MATCH (target:Entity {uuid: $target_uuid}) 
         MERGE (source)-[r:RELATES_TO {uuid: $uuid}]->(target)
-        SET r = {uuid: $uuid, name: $name, group_id: $group_id, fact: $fact, fact_embedding: $fact_embedding, 
-        episodes: $episodes, created_at: $created_at, expired_at: $expired_at, 
-        valid_at: $valid_at, invalid_at: $invalid_at}
+        SET r = {uuid: $uuid, name: $name, group_id: $group_id, fact: $fact, episodes: $episodes, 
+        created_at: $created_at, expired_at: $expired_at, valid_at: $valid_at, invalid_at: $invalid_at}
+        WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $fact_embedding)
         RETURN r.uuid AS uuid""",
             source_uuid=self.source_node_uuid,
             target_uuid=self.target_node_uuid,

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/embedder/openai.py

@@ -42,7 +42,7 @@ class OpenAIEmbedder(EmbedderClient):
         self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
 
     async def create(
-        self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+            self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         result = await self.client.embeddings.create(input=input, model=self.config.embedding_model)
         return result.data[0].embedding[: self.config.embedding_dim]

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/embedder/voyage.py

@@ -41,7 +41,7 @@ class VoyageAIEmbedder(EmbedderClient):
         self.client = voyageai.AsyncClient(api_key=config.api_key)
 
     async def create(
-        self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+            self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         result = await self.client.embed(input, model=self.config.embedding_model)
         return result.embeddings[0][: self.config.embedding_dim]

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/graphiti.py

@@ -21,11 +21,12 @@ from time import time
 
 from dotenv import load_dotenv
 from neo4j import AsyncGraphDatabase
+from pydantic import BaseModel
 
 from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
 from graphiti_core.llm_client import LLMClient, OpenAIClient
-from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
+from graphiti_core.nodes import CommunityNode, EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.search.search import SearchConfig, search
 from graphiti_core.search.search_config import DEFAULT_SEARCH_LIMIT, SearchResults
 from graphiti_core.search.search_config_recipes import (
@@ -77,6 +78,12 @@ logger = logging.getLogger(__name__)
 load_dotenv()
 
 
+class AddEpisodeResults(BaseModel):
+    episode: EpisodicNode
+    nodes: list[EntityNode]
+    edges: list[EntityEdge]
+
+
 class Graphiti:
     def __init__(
         self,
@@ -245,7 +252,7 @@ class Graphiti:
         group_id: str = '',
         uuid: str | None = None,
         update_communities: bool = False,
-    ):
+    ) -> AddEpisodeResults:
         """
         Process an episode and update the graph.
 
@@ -451,6 +458,8 @@ class Graphiti:
             end = time()
             logger.info(f'Completed add_episode in {(end - start) * 1000} ms')
 
+            return AddEpisodeResults(episode=episode, nodes=nodes, edges=entity_edges)
+
         except Exception as e:
             raise e
 
@@ -567,11 +576,20 @@ class Graphiti:
         except Exception as e:
             raise e
 
-    async def build_communities(self):
+    async def build_communities(self, group_ids: list[str] | None = None) -> list[CommunityNode]:
+        """
+        Use a community clustering algorithm to find communities of nodes. Create community nodes summarising
+        the content of these communities.
+        ----------
+        query : list[str] | None
+            Optional. Create communities only for the listed group_ids. If blank the entire graph will be used.
+        """
         # Clear existing communities
         await remove_communities(self.driver)
 
-        community_nodes, community_edges = await build_communities(self.driver, self.llm_client)
+        community_nodes, community_edges = await build_communities(
+            self.driver, self.llm_client, group_ids
+        )
 
         await asyncio.gather(
             *[node.generate_name_embedding(self.embedder) for node in community_nodes]
@@ -580,6 +598,8 @@ class Graphiti:
         await asyncio.gather(*[node.save(self.driver) for node in community_nodes])
         await asyncio.gather(*[edge.save(self.driver) for edge in community_edges])
 
+        return community_nodes
+
     async def search(
         self,
         query: str,
@@ -700,18 +720,17 @@ class Graphiti:
         ).nodes
         return nodes
 
+    async def get_episode_mentions(self, episode_uuids: list[str]) -> SearchResults:
+        episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids)
 
-async def get_episode_mentions(self, episode_uuids: list[str]) -> SearchResults:
-    episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids)
-
-    edges_list = await asyncio.gather(
-        *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes]
-    )
+        edges_list = await asyncio.gather(
+            *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes]
+        )
 
-    edges: list[EntityEdge] = [edge for lst in edges_list for edge in lst]
+        edges: list[EntityEdge] = [edge for lst in edges_list for edge in lst]
 
-    nodes = await get_mentioned_nodes(self.driver, episodes)
+        nodes = await get_mentioned_nodes(self.driver, episodes)
 
-    communities = await get_communities_by_nodes(self.driver, nodes)
+        communities = await get_communities_by_nodes(self.driver, nodes)
 
-    return SearchResults(edges=edges, nodes=nodes, communities=communities)
+        return SearchResults(edges=edges, nodes=nodes, communities=communities)

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/helpers.py

@@ -16,6 +16,7 @@ limitations under the License.
 
 from datetime import datetime
 
+import numpy as np
 from neo4j import time as neo4j_time
 
 
@@ -25,7 +26,7 @@ def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:
 
 def lucene_sanitize(query: str) -> str:
     # Escape special characters from a query before passing into Lucene
-    # + - && || ! ( ) { } [ ] ^ " ~ * ? : \
+    # + - && || ! ( ) { } [ ] ^ " ~ * ? : \ /
     escape_map = str.maketrans(
         {
             '+': r'\+',
@@ -46,8 +47,21 @@ def lucene_sanitize(query: str) -> str:
             '?': r'\?',
             ':': r'\:',
             '\\': r'\\',
+            '/': r'\/',
         }
     )
 
     sanitized = query.translate(escape_map)
     return sanitized
+
+
+def normalize_l2(embedding: list[float]) -> list[float]:
+    embedding_array = np.array(embedding)
+    if embedding_array.ndim == 1:
+        norm = np.linalg.norm(embedding_array)
+        if norm == 0:
+            return embedding_array.tolist()
+        return (embedding_array / norm).tolist()
+    else:
+        norm = np.linalg.norm(embedding_array, 2, axis=1, keepdims=True)
+        return (np.where(norm == 0, embedding_array, embedding_array / norm)).tolist()

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/nodes.py

@@ -225,7 +225,8 @@ class EntityNode(Node):
         result = await driver.execute_query(
             """
         MERGE (n:Entity {uuid: $uuid})
-        SET n = {uuid: $uuid, name: $name, name_embedding: $name_embedding, group_id: $group_id, summary: $summary, created_at: $created_at}
+        SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
+        WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
         RETURN n.uuid AS uuid""",
             uuid=self.uuid,
             name=self.name,
@@ -308,7 +309,8 @@ class CommunityNode(Node):
         result = await driver.execute_query(
             """
         MERGE (n:Community {uuid: $uuid})
-        SET n = {uuid: $uuid, name: $name, name_embedding: $name_embedding, group_id: $group_id, summary: $summary, created_at: $created_at}
+        SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
+        WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
         RETURN n.uuid AS uuid""",
             uuid=self.uuid,
             name=self.name,

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/prompts/eval.py

@@ -23,11 +23,33 @@ from .models import Message, PromptFunction, PromptVersion
 class Prompt(Protocol):
     qa_prompt: PromptVersion
     eval_prompt: PromptVersion
+    query_expansion: PromptVersion
 
 
 class Versions(TypedDict):
     qa_prompt: PromptFunction
     eval_prompt: PromptFunction
+    query_expansion: PromptFunction
+
+
+def query_expansion(context: dict[str, Any]) -> list[Message]:
+    sys_prompt = """You are an expert at rephrasing questions into queries used in a database retrieval system"""
+
+    user_prompt = f"""
+    Bob is asking Alice a question, are you able to rephrase the question into a simpler one about Alice in the third person
+    that maintains the relevant context?
+    <QUESTION>
+    {json.dumps(context['query'])}
+    </QUESTION>
+    respond with a JSON object in the following format:
+    {{
+        "query": "query optimized for database search"
+    }}
+    """
+    return [
+        Message(role='system', content=sys_prompt),
+        Message(role='user', content=user_prompt),
+    ]
 
 
 def qa_prompt(context: dict[str, Any]) -> list[Message]:
@@ -38,7 +60,7 @@ def qa_prompt(context: dict[str, Any]) -> list[Message]:
     You are given the following entity summaries and facts to help you determine the answer to your question.
     <ENTITY_SUMMARIES>
     {json.dumps(context['entity_summaries'])}
-    </ENTITY_SUMMARIES
+    </ENTITY_SUMMARIES>
     <FACTS>
     {json.dumps(context['facts'])}
     </FACTS>
@@ -87,4 +109,8 @@ def eval_prompt(context: dict[str, Any]) -> list[Message]:
     ]
 
 
-versions: Versions = {'qa_prompt': qa_prompt, 'eval_prompt': eval_prompt}
+versions: Versions = {
+    'qa_prompt': qa_prompt,
+    'eval_prompt': eval_prompt,
+    'query_expansion': query_expansion,
+}

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/prompts/extract_edge_dates.py

@@ -37,7 +37,6 @@ def v1(context: dict[str, Any]) -> list[Message]:
             role='user',
             content=f"""
             Edge:
-            Edge Name: {context['edge_name']}
             Fact: {context['edge_fact']}
 
             Current Episode: {context['current_episode']}
@@ -56,17 +55,17 @@ def v1(context: dict[str, Any]) -> list[Message]:
             Guidelines:
             1. Use ISO 8601 format (YYYY-MM-DDTHH:MM:SSZ) for datetimes.
             2. Use the reference timestamp as the current time when determining the valid_at and invalid_at dates.
-            3. If no temporal information is found that establishes or changes the relationship, leave the fields as null.
-            4. Do not infer dates from related events. Only use dates that are directly stated to establish or change the relationship.
-			5. For relative time mentions directly related to the relationship, calculate the actual datetime based on the reference timestamp.
-            6. If only a date is mentioned without a specific time, use 00:00:00 (midnight) for that date.
-            7. If only a year is mentioned, use January 1st of that year at 00:00:00.
+            3. If the fact is written in the present tense, use the Reference Timestamp for the valid_at date
+            4. If no temporal information is found that establishes or changes the relationship, leave the fields as null.
+            5. Do not infer dates from related events. Only use dates that are directly stated to establish or change the relationship.
+			6. For relative time mentions directly related to the relationship, calculate the actual datetime based on the reference timestamp.
+            7. If only a date is mentioned without a specific time, use 00:00:00 (midnight) for that date.
+            8. If only a year is mentioned, use January 1st of that year at 00:00:00.
             9. Always include the time zone offset (use Z for UTC if no specific time zone is mentioned).
             Respond with a JSON object:
             {{
-                "valid_at": "YYYY-MM-DDTHH:MM:SSZ or null",
-                "invalid_at": "YYYY-MM-DDTHH:MM:SSZ or null",
-                "explanation": "Brief explanation of why these dates were chosen or why they were set to null"
+                "valid_at": "YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null",
+                "invalid_at": "YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null",
             }}
             """,
         ),

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/prompts/extract_edges.py

@@ -113,8 +113,9 @@ def v2(context: dict[str, Any]) -> list[Message]:
         2. Each edge should represent a clear relationship between two DISTINCT nodes.
         3. The relation_type should be a concise, all-caps description of the relationship (e.g., LOVES, IS_FRIENDS_WITH, WORKS_FOR).
         4. Provide a more detailed fact describing the relationship.
-        5. Consider temporal aspects of relationships when relevant.
-        6. Avoid using the same node as the source and target of a relationship
+        5. The fact should include any specific relevant information, including numeric information
+        6. Consider temporal aspects of relationships when relevant.
+        7. Avoid using the same node as the source and target of a relationship
 
         Respond with a JSON object in the following format:
         {{

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/prompts/invalidate_edges.py

@@ -82,7 +82,7 @@ def v2(context: dict[str, Any]) -> list[Message]:
         Message(
             role='user',
             content=f"""
-               Based on the provided Existing Edges and a New Edge, determine which existing edges, if any, should be marked as invalidated due to contradictions with the New Edge.
+               Based on the provided Existing Edges and a New Edge, determine which existing edges, if any, should be marked as invalidated due to invalidations with the New Edge.
 
                 Existing Edges:
                 {context['existing_edges']}

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/search/search.py

@@ -29,13 +29,10 @@ from graphiti_core.search.search_config import (
     DEFAULT_SEARCH_LIMIT,
     CommunityReranker,
     CommunitySearchConfig,
-    CommunitySearchMethod,
     EdgeReranker,
     EdgeSearchConfig,
-    EdgeSearchMethod,
     NodeReranker,
     NodeSearchConfig,
-    NodeSearchMethod,
     SearchConfig,
     SearchResults,
 )
@@ -45,6 +42,7 @@ from graphiti_core.search.search_utils import (
     edge_fulltext_search,
     edge_similarity_search,
     episode_mentions_reranker,
+    maximal_marginal_relevance,
     node_distance_reranker,
     node_fulltext_search,
     node_similarity_search,
@@ -120,22 +118,18 @@ async def edge_search(
     if config is None:
         return []
 
-    search_results: list[list[EntityEdge]] = []
+    query_vector = await embedder.create(input=[query])
 
-    if EdgeSearchMethod.bm25 in config.search_methods:
-        text_search = await edge_fulltext_search(driver, query, None, None, group_ids, 2 * limit)
-        search_results.append(text_search)
-
-    if EdgeSearchMethod.cosine_similarity in config.search_methods:
-        search_vector = await embedder.create(input=[query])
-
-        similarity_search = await edge_similarity_search(
-            driver, search_vector, None, None, group_ids, 2 * limit
+    search_results: list[list[EntityEdge]] = list(
+        await asyncio.gather(
+            *[
+                edge_fulltext_search(driver, query, None, None, group_ids, 2 * limit),
+                edge_similarity_search(
+                    driver, query_vector, None, None, group_ids, 2 * limit, config.sim_min_score
+                ),
+            ]
         )
-        search_results.append(similarity_search)
-
-    if len(search_results) > 1 and config.reranker is None:
-        raise SearchRerankerError('Multiple edge searches enabled without a reranker')
+    )
 
     edge_uuid_map = {edge.uuid: edge for result in search_results for edge in result}
 
@@ -144,6 +138,15 @@ async def edge_search(
         search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
 
         reranked_uuids = rrf(search_result_uuids)
+    elif config.reranker == EdgeReranker.mmr:
+        search_result_uuids_and_vectors = [
+            (edge.uuid, edge.fact_embedding if edge.fact_embedding is not None else [0.0] * 1024)
+            for result in search_results
+            for edge in result
+        ]
+        reranked_uuids = maximal_marginal_relevance(
+            query_vector, search_result_uuids_and_vectors, config.mmr_lambda
+        )
     elif config.reranker == EdgeReranker.node_distance:
         if center_node_uuid is None:
             raise SearchRerankerError('No center node provided for Node Distance reranker')
@@ -184,22 +187,18 @@ async def node_search(
     if config is None:
         return []
 
-    search_results: list[list[EntityNode]] = []
-
-    if NodeSearchMethod.bm25 in config.search_methods:
-        text_search = await node_fulltext_search(driver, query, group_ids, 2 * limit)
-        search_results.append(text_search)
+    query_vector = await embedder.create(input=[query])
 
-    if NodeSearchMethod.cosine_similarity in config.search_methods:
-        search_vector = await embedder.create(input=[query])
-
-        similarity_search = await node_similarity_search(
-            driver, search_vector, group_ids, 2 * limit
+    search_results: list[list[EntityNode]] = list(
+        await asyncio.gather(
+            *[
+                node_fulltext_search(driver, query, group_ids, 2 * limit),
+                node_similarity_search(
+                    driver, query_vector, group_ids, 2 * limit, config.sim_min_score
+                ),
+            ]
         )
-        search_results.append(similarity_search)
-
-    if len(search_results) > 1 and config.reranker is None:
-        raise SearchRerankerError('Multiple node searches enabled without a reranker')
+    )
 
     search_result_uuids = [[node.uuid for node in result] for result in search_results]
     node_uuid_map = {node.uuid: node for result in search_results for node in result}
@@ -207,6 +206,15 @@ async def node_search(
     reranked_uuids: list[str] = []
     if config.reranker == NodeReranker.rrf:
         reranked_uuids = rrf(search_result_uuids)
+    elif config.reranker == NodeReranker.mmr:
+        search_result_uuids_and_vectors = [
+            (node.uuid, node.name_embedding if node.name_embedding is not None else [0.0] * 1024)
+            for result in search_results
+            for node in result
+        ]
+        reranked_uuids = maximal_marginal_relevance(
+            query_vector, search_result_uuids_and_vectors, config.mmr_lambda
+        )
     elif config.reranker == NodeReranker.episode_mentions:
         reranked_uuids = await episode_mentions_reranker(driver, search_result_uuids)
     elif config.reranker == NodeReranker.node_distance:
@@ -232,22 +240,18 @@ async def community_search(
     if config is None:
         return []
 
-    search_results: list[list[CommunityNode]] = []
-
-    if CommunitySearchMethod.bm25 in config.search_methods:
-        text_search = await community_fulltext_search(driver, query, group_ids, 2 * limit)
-        search_results.append(text_search)
-
-    if CommunitySearchMethod.cosine_similarity in config.search_methods:
-        search_vector = await embedder.create(input=[query])
+    query_vector = await embedder.create(input=[query])
 
-        similarity_search = await community_similarity_search(
-            driver, search_vector, group_ids, 2 * limit
+    search_results: list[list[CommunityNode]] = list(
+        await asyncio.gather(
+            *[
+                community_fulltext_search(driver, query, group_ids, 2 * limit),
+                community_similarity_search(
+                    driver, query_vector, group_ids, 2 * limit, config.sim_min_score
+                ),
+            ]
         )
-        search_results.append(similarity_search)
-
-    if len(search_results) > 1 and config.reranker is None:
-        raise SearchRerankerError('Multiple node searches enabled without a reranker')
+    )
 
     search_result_uuids = [[community.uuid for community in result] for result in search_results]
     community_uuid_map = {
@@ -257,6 +261,18 @@ async def community_search(
     reranked_uuids: list[str] = []
     if config.reranker == CommunityReranker.rrf:
         reranked_uuids = rrf(search_result_uuids)
+    elif config.reranker == CommunityReranker.mmr:
+        search_result_uuids_and_vectors = [
+            (
+                community.uuid,
+                community.name_embedding if community.name_embedding is not None else [0.0] * 1024,
+            )
+            for result in search_results
+            for community in result
+        ]
+        reranked_uuids = maximal_marginal_relevance(
+            query_vector, search_result_uuids_and_vectors, config.mmr_lambda
+        )
 
     reranked_communities = [community_uuid_map[uuid] for uuid in reranked_uuids]
 

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/search/search_config.py

@@ -20,6 +20,7 @@ from pydantic import BaseModel, Field
 
 from graphiti_core.edges import EntityEdge
 from graphiti_core.nodes import CommunityNode, EntityNode
+from graphiti_core.search.search_utils import DEFAULT_MIN_SCORE, DEFAULT_MMR_LAMBDA
 
 DEFAULT_SEARCH_LIMIT = 10
 
@@ -43,31 +44,40 @@ class EdgeReranker(Enum):
     rrf = 'reciprocal_rank_fusion'
     node_distance = 'node_distance'
     episode_mentions = 'episode_mentions'
+    mmr = 'mmr'
 
 
 class NodeReranker(Enum):
     rrf = 'reciprocal_rank_fusion'
     node_distance = 'node_distance'
     episode_mentions = 'episode_mentions'
+    mmr = 'mmr'
 
 
 class CommunityReranker(Enum):
     rrf = 'reciprocal_rank_fusion'
+    mmr = 'mmr'
 
 
 class EdgeSearchConfig(BaseModel):
     search_methods: list[EdgeSearchMethod]
-    reranker: EdgeReranker | None
+    reranker: EdgeReranker = Field(default=EdgeReranker.rrf)
+    sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
+    mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
 
 
 class NodeSearchConfig(BaseModel):
     search_methods: list[NodeSearchMethod]
-    reranker: NodeReranker | None
+    reranker: NodeReranker = Field(default=NodeReranker.rrf)
+    sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
+    mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
 
 
 class CommunitySearchConfig(BaseModel):
     search_methods: list[CommunitySearchMethod]
-    reranker: CommunityReranker | None
+    reranker: CommunityReranker = Field(default=CommunityReranker.rrf)
+    sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
+    mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
 
 
 class SearchConfig(BaseModel):

{graphiti_core-0.3.8 → graphiti_core-0.3.9}/graphiti_core/search/search_config_recipes.py

@@ -43,6 +43,22 @@ COMBINED_HYBRID_SEARCH_RRF = SearchConfig(
     ),
 )
 
+# Performs a hybrid search with mmr reranking over edges, nodes, and communities
+COMBINED_HYBRID_SEARCH_MMR = SearchConfig(
+    edge_config=EdgeSearchConfig(
+        search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
+        reranker=EdgeReranker.mmr,
+    ),
+    node_config=NodeSearchConfig(
+        search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
+        reranker=NodeReranker.mmr,
+    ),
+    community_config=CommunitySearchConfig(
+        search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
+        reranker=CommunityReranker.mmr,
+    ),
+)
+
 # performs a hybrid search over edges with rrf reranking
 EDGE_HYBRID_SEARCH_RRF = SearchConfig(
     edge_config=EdgeSearchConfig(
@@ -51,6 +67,14 @@ EDGE_HYBRID_SEARCH_RRF = SearchConfig(
     )
 )
 
+# performs a hybrid search over edges with mmr reranking
+EDGE_HYBRID_SEARCH_mmr = SearchConfig(
+    edge_config=EdgeSearchConfig(
+        search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
+        reranker=EdgeReranker.mmr,
+    )
+)
+
 # performs a hybrid search over edges with node distance reranking
 EDGE_HYBRID_SEARCH_NODE_DISTANCE = SearchConfig(
     edge_config=EdgeSearchConfig(
@@ -75,6 +99,14 @@ NODE_HYBRID_SEARCH_RRF = SearchConfig(
     )
 )
 
+# performs a hybrid search over nodes with mmr reranking
+NODE_HYBRID_SEARCH_MMR = SearchConfig(
+    node_config=NodeSearchConfig(
+        search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
+        reranker=NodeReranker.mmr,
+    )
+)
+
 # performs a hybrid search over nodes with node distance reranking
 NODE_HYBRID_SEARCH_NODE_DISTANCE = SearchConfig(
     node_config=NodeSearchConfig(
@@ -98,3 +130,11 @@ COMMUNITY_HYBRID_SEARCH_RRF = SearchConfig(
         reranker=CommunityReranker.rrf,
     )
 )
+
+# performs a hybrid search over communities with mmr reranking
+COMMUNITY_HYBRID_SEARCH_MMR = SearchConfig(
+    community_config=CommunitySearchConfig(
+        search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
+        reranker=CommunityReranker.mmr,
+    )
+)