graphiti-core 0.1.0__py3-none-any.whl

graphiti_core/search/search_utils.py

@@ -0,0 +1,454 @@
+import asyncio
+import logging
+import re
+import typing
+from collections import defaultdict
+from time import time
+
+from neo4j import AsyncDriver
+
+from graphiti_core.edges import EntityEdge
+from graphiti_core.helpers import parse_db_date
+from graphiti_core.nodes import EntityNode, EpisodicNode
+
+logger = logging.getLogger(__name__)
+
+RELEVANT_SCHEMA_LIMIT = 3
+
+
+async def get_mentioned_nodes(driver: AsyncDriver, episodes: list[EpisodicNode]):
+    episode_uuids = [episode.uuid for episode in episodes]
+    records, _, _ = await driver.execute_query(
+        """
+        MATCH (episode:Episodic)-[:MENTIONS]->(n:Entity) WHERE episode.uuid IN $uuids
+        RETURN DISTINCT
+            n.uuid As uuid, 
+            n.name AS name, 
+            n.created_at AS created_at, 
+            n.summary AS summary
+        """,
+        uuids=episode_uuids,
+    )
+
+    nodes: list[EntityNode] = []
+
+    for record in records:
+        nodes.append(
+            EntityNode(
+                uuid=record['uuid'],
+                name=record['name'],
+                labels=['Entity'],
+                created_at=record['created_at'].to_native(),
+                summary=record['summary'],
+            )
+        )
+
+    return nodes
+
+
+async def bfs(node_ids: list[str], driver: AsyncDriver):
+    records, _, _ = await driver.execute_query(
+        """
+        MATCH (n WHERE n.uuid in $node_ids)-[r]->(m)
+        RETURN DISTINCT
+            n.uuid AS source_node_uuid,
+            n.name AS source_name, 
+            n.summary AS source_summary,
+            m.uuid AS target_node_uuid,
+            m.name AS target_name, 
+            m.summary AS target_summary,
+            r.uuid AS uuid,
+            r.created_at AS created_at,
+            r.name AS name,
+            r.fact AS fact,
+            r.fact_embedding AS fact_embedding,
+            r.episodes AS episodes,
+            r.expired_at AS expired_at,
+            r.valid_at AS valid_at,
+            r.invalid_at AS invalid_at
+        
+    """,
+        node_ids=node_ids,
+    )
+
+    context: dict[str, typing.Any] = {}
+
+    for record in records:
+        n_uuid = record['source_node_uuid']
+        if n_uuid in context:
+            context[n_uuid]['facts'].append(record['fact'])
+        else:
+            context[n_uuid] = {
+                'name': record['source_name'],
+                'summary': record['source_summary'],
+                'facts': [record['fact']],
+            }
+
+        m_uuid = record['target_node_uuid']
+        if m_uuid not in context:
+            context[m_uuid] = {
+                'name': record['target_name'],
+                'summary': record['target_summary'],
+                'facts': [],
+            }
+    logger.info(f'bfs search returned context: {context}')
+    return context
+
+
+async def edge_similarity_search(
+    search_vector: list[float], driver: AsyncDriver, limit=RELEVANT_SCHEMA_LIMIT
+) -> list[EntityEdge]:
+    # vector similarity search over embedded facts
+    records, _, _ = await driver.execute_query(
+        """
+                CALL db.index.vector.queryRelationships("fact_embedding", $limit, $search_vector)
+                YIELD relationship AS r, score
+                MATCH (n)-[r:RELATES_TO]->(m)
+                RETURN
+                    r.uuid AS uuid,
+                    n.uuid AS source_node_uuid,
+                    m.uuid AS target_node_uuid,
+                    r.created_at AS created_at,
+                    r.name AS name,
+                    r.fact AS fact,
+                    r.fact_embedding AS fact_embedding,
+                    r.episodes AS episodes,
+                    r.expired_at AS expired_at,
+                    r.valid_at AS valid_at,
+                    r.invalid_at AS invalid_at
+                ORDER BY score DESC
+                """,
+        search_vector=search_vector,
+        limit=limit,
+    )
+
+    edges: list[EntityEdge] = []
+
+    for record in records:
+        edge = EntityEdge(
+            uuid=record['uuid'],
+            source_node_uuid=record['source_node_uuid'],
+            target_node_uuid=record['target_node_uuid'],
+            fact=record['fact'],
+            name=record['name'],
+            episodes=record['episodes'],
+            fact_embedding=record['fact_embedding'],
+            created_at=record['created_at'].to_native(),
+            expired_at=parse_db_date(record['expired_at']),
+            valid_at=parse_db_date(record['valid_at']),
+            invalid_at=parse_db_date(record['invalid_at']),
+        )
+
+        edges.append(edge)
+
+    return edges
+
+
+async def entity_similarity_search(
+    search_vector: list[float], driver: AsyncDriver, limit=RELEVANT_SCHEMA_LIMIT
+) -> list[EntityNode]:
+    # vector similarity search over entity names
+    records, _, _ = await driver.execute_query(
+        """
+                CALL db.index.vector.queryNodes("name_embedding", $limit, $search_vector)
+                YIELD node AS n, score
+                RETURN
+                    n.uuid As uuid, 
+                    n.name AS name, 
+                    n.created_at AS created_at, 
+                    n.summary AS summary
+                ORDER BY score DESC
+                """,
+        search_vector=search_vector,
+        limit=limit,
+    )
+    nodes: list[EntityNode] = []
+
+    for record in records:
+        nodes.append(
+            EntityNode(
+                uuid=record['uuid'],
+                name=record['name'],
+                labels=['Entity'],
+                created_at=record['created_at'].to_native(),
+                summary=record['summary'],
+            )
+        )
+
+    return nodes
+
+
+async def entity_fulltext_search(
+    query: str, driver: AsyncDriver, limit=RELEVANT_SCHEMA_LIMIT
+) -> list[EntityNode]:
+    # BM25 search to get top nodes
+    fuzzy_query = re.sub(r'[^\w\s]', '', query) + '~'
+    records, _, _ = await driver.execute_query(
+        """
+    CALL db.index.fulltext.queryNodes("name_and_summary", $query) YIELD node, score
+    RETURN
+        node.uuid As uuid, 
+        node.name AS name, 
+        node.created_at AS created_at, 
+        node.summary AS summary
+    ORDER BY score DESC
+    LIMIT $limit
+    """,
+        query=fuzzy_query,
+        limit=limit,
+    )
+    nodes: list[EntityNode] = []
+
+    for record in records:
+        nodes.append(
+            EntityNode(
+                uuid=record['uuid'],
+                name=record['name'],
+                labels=['Entity'],
+                created_at=record['created_at'].to_native(),
+                summary=record['summary'],
+            )
+        )
+
+    return nodes
+
+
+async def edge_fulltext_search(
+    query: str, driver: AsyncDriver, limit=RELEVANT_SCHEMA_LIMIT
+) -> list[EntityEdge]:
+    # fulltext search over facts
+    fuzzy_query = re.sub(r'[^\w\s]', '', query) + '~'
+
+    records, _, _ = await driver.execute_query(
+        """
+                CALL db.index.fulltext.queryRelationships("name_and_fact", $query) 
+                YIELD relationship AS r, score
+                MATCH (n:Entity)-[r]->(m:Entity)
+                RETURN 
+                    r.uuid AS uuid,
+                    n.uuid AS source_node_uuid,
+                    m.uuid AS target_node_uuid,
+                    r.created_at AS created_at,
+                    r.name AS name,
+                    r.fact AS fact,
+                    r.fact_embedding AS fact_embedding,
+                    r.episodes AS episodes,
+                    r.expired_at AS expired_at,
+                    r.valid_at AS valid_at,
+                    r.invalid_at AS invalid_at
+                ORDER BY score DESC LIMIT $limit
+                """,
+        query=fuzzy_query,
+        limit=limit,
+    )
+
+    edges: list[EntityEdge] = []
+
+    for record in records:
+        edge = EntityEdge(
+            uuid=record['uuid'],
+            source_node_uuid=record['source_node_uuid'],
+            target_node_uuid=record['target_node_uuid'],
+            fact=record['fact'],
+            name=record['name'],
+            episodes=record['episodes'],
+            fact_embedding=record['fact_embedding'],
+            created_at=record['created_at'].to_native(),
+            expired_at=parse_db_date(record['expired_at']),
+            valid_at=parse_db_date(record['valid_at']),
+            invalid_at=parse_db_date(record['invalid_at']),
+        )
+
+        edges.append(edge)
+
+    return edges
+
+
+async def hybrid_node_search(
+    queries: list[str],
+    embeddings: list[list[float]],
+    driver: AsyncDriver,
+    limit: int | None = None,
+) -> list[EntityNode]:
+    """
+    Perform a hybrid search for nodes using both text queries and embeddings.
+
+    This method combines fulltext search and vector similarity search to find
+    relevant nodes in the graph database.
+
+    Parameters
+    ----------
+    queries : list[str]
+        A list of text queries to search for.
+    embeddings : list[list[float]]
+        A list of embedding vectors corresponding to the queries. If empty only fulltext search is performed.
+    driver : AsyncDriver
+        The Neo4j driver instance for database operations.
+    limit : int | None, optional
+        The maximum number of results to return per search method. If None, a default limit will be applied.
+
+    Returns
+    -------
+    list[EntityNode]
+        A list of unique EntityNode objects that match the search criteria.
+
+    Notes
+    -----
+    This method performs the following steps:
+    1. Executes fulltext searches for each query.
+    2. Executes vector similarity searches for each embedding.
+    3. Combines and deduplicates the results from both search types.
+    4. Logs the performance metrics of the search operation.
+
+    The search results are deduplicated based on the node UUIDs to ensure
+    uniqueness in the returned list. The 'limit' parameter is applied to each
+    individual search method before deduplication. If not specified, a default
+    limit (defined in the individual search functions) will be used.
+    """
+
+    start = time()
+    relevant_nodes: list[EntityNode] = []
+    relevant_node_uuids = set()
+
+    results = await asyncio.gather(
+        *[entity_fulltext_search(q, driver, 2 * (limit or RELEVANT_SCHEMA_LIMIT)) for q in queries],
+        *[
+            entity_similarity_search(e, driver, 2 * (limit or RELEVANT_SCHEMA_LIMIT))
+            for e in embeddings
+        ],
+    )
+
+    for result in results:
+        for node in result:
+            if node.uuid in relevant_node_uuids:
+                continue
+
+            relevant_node_uuids.add(node.uuid)
+            relevant_nodes.append(node)
+
+    end = time()
+    logger.info(f'Found relevant nodes: {relevant_node_uuids} in {(end - start) * 1000} ms')
+    return relevant_nodes
+
+
+async def get_relevant_nodes(
+    nodes: list[EntityNode],
+    driver: AsyncDriver,
+) -> list[EntityNode]:
+    """
+    Retrieve relevant nodes based on the provided list of EntityNodes.
+
+    This method performs a hybrid search using both the names and embeddings
+    of the input nodes to find relevant nodes in the graph database.
+
+    Parameters
+    ----------
+    nodes : list[EntityNode]
+        A list of EntityNode objects to use as the basis for the search.
+    driver : AsyncDriver
+        The Neo4j driver instance for database operations.
+
+    Returns
+    -------
+    list[EntityNode]
+        A list of EntityNode objects that are deemed relevant based on the input nodes.
+
+    Notes
+    -----
+    This method uses the hybrid_node_search function to perform the search,
+    which combines fulltext search and vector similarity search.
+    It extracts the names and name embeddings (if available) from the input nodes
+    to use as search criteria.
+    """
+    relevant_nodes = await hybrid_node_search(
+        [node.name for node in nodes],
+        [node.name_embedding for node in nodes if node.name_embedding is not None],
+        driver,
+    )
+    return relevant_nodes
+
+
+async def get_relevant_edges(
+    edges: list[EntityEdge],
+    driver: AsyncDriver,
+) -> list[EntityEdge]:
+    start = time()
+    relevant_edges: list[EntityEdge] = []
+    relevant_edge_uuids = set()
+
+    results = await asyncio.gather(
+        *[
+            edge_similarity_search(edge.fact_embedding, driver)
+            for edge in edges
+            if edge.fact_embedding is not None
+        ],
+        *[edge_fulltext_search(edge.fact, driver) for edge in edges],
+    )
+
+    for result in results:
+        for edge in result:
+            if edge.uuid in relevant_edge_uuids:
+                continue
+
+            relevant_edge_uuids.add(edge.uuid)
+            relevant_edges.append(edge)
+
+    end = time()
+    logger.info(f'Found relevant edges: {relevant_edge_uuids} in {(end - start) * 1000} ms')
+
+    return relevant_edges
+
+
+# takes in a list of rankings of uuids
+def rrf(results: list[list[str]], rank_const=1) -> list[str]:
+    scores: dict[str, float] = defaultdict(float)
+    for result in results:
+        for i, uuid in enumerate(result):
+            scores[uuid] += 1 / (i + rank_const)
+
+    scored_uuids = [term for term in scores.items()]
+    scored_uuids.sort(reverse=True, key=lambda term: term[1])
+
+    sorted_uuids = [term[0] for term in scored_uuids]
+
+    return sorted_uuids
+
+
+async def node_distance_reranker(
+    driver: AsyncDriver, results: list[list[str]], center_node_uuid: str
+) -> list[str]:
+    # use rrf as a preliminary ranker
+    sorted_uuids = rrf(results)
+    scores: dict[str, float] = {}
+
+    for uuid in sorted_uuids:
+        # Find shortest path to center node
+        records, _, _ = await driver.execute_query(
+            """  
+        MATCH (source:Entity)-[r:RELATES_TO {uuid: $edge_uuid}]->(target:Entity)
+        MATCH p = SHORTEST 1 (center:Entity)-[:RELATES_TO]-+(n:Entity)
+        WHERE center.uuid = $center_uuid AND n.uuid IN [source.uuid, target.uuid]
+        RETURN min(length(p)) AS score, source.uuid AS source_uuid, target.uuid AS target_uuid
+        """,
+            edge_uuid=uuid,
+            center_uuid=center_node_uuid,
+        )
+        distance = 0.01
+
+        for record in records:
+            if (
+                record['source_uuid'] == center_node_uuid
+                or record['target_uuid'] == center_node_uuid
+            ):
+                continue
+            distance = record['score']
+
+        if uuid in scores:
+            scores[uuid] = min(1 / distance, scores[uuid])
+        else:
+            scores[uuid] = 1 / distance
+
+    # rerank on shortest distance
+    sorted_uuids.sort(reverse=True, key=lambda cur_uuid: scores[cur_uuid])
+
+    return sorted_uuids

graphiti_core/utils/__init__.py

@@ -0,0 +1,15 @@
+from .maintenance import (
+    build_episodic_edges,
+    clear_data,
+    extract_edges,
+    extract_nodes,
+    retrieve_episodes,
+)
+
+__all__ = [
+    'extract_edges',
+    'build_episodic_edges',
+    'extract_nodes',
+    'clear_data',
+    'retrieve_episodes',
+]

graphiti_core/utils/bulk_utils.py

@@ -0,0 +1,227 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import asyncio
+import typing
+from datetime import datetime
+
+from neo4j import AsyncDriver
+from numpy import dot
+from pydantic import BaseModel
+
+from graphiti_core.edges import Edge, EntityEdge, EpisodicEdge
+from graphiti_core.llm_client import LLMClient
+from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
+from graphiti_core.search.search_utils import get_relevant_edges, get_relevant_nodes
+from graphiti_core.utils import retrieve_episodes
+from graphiti_core.utils.maintenance.edge_operations import (
+    build_episodic_edges,
+    dedupe_edge_list,
+    dedupe_extracted_edges,
+    extract_edges,
+)
+from graphiti_core.utils.maintenance.graph_data_operations import EPISODE_WINDOW_LEN
+from graphiti_core.utils.maintenance.node_operations import (
+    dedupe_extracted_nodes,
+    dedupe_node_list,
+    extract_nodes,
+)
+
+CHUNK_SIZE = 15
+
+
+class RawEpisode(BaseModel):
+    name: str
+    content: str
+    source_description: str
+    source: EpisodeType
+    reference_time: datetime
+
+
+async def retrieve_previous_episodes_bulk(
+    driver: AsyncDriver, episodes: list[EpisodicNode]
+) -> list[tuple[EpisodicNode, list[EpisodicNode]]]:
+    previous_episodes_list = await asyncio.gather(
+        *[
+            retrieve_episodes(driver, episode.valid_at, last_n=EPISODE_WINDOW_LEN)
+            for episode in episodes
+        ]
+    )
+    episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]] = [
+        (episode, previous_episodes_list[i]) for i, episode in enumerate(episodes)
+    ]
+
+    return episode_tuples
+
+
+async def extract_nodes_and_edges_bulk(
+    llm_client: LLMClient, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]]
+) -> tuple[list[EntityNode], list[EntityEdge], list[EpisodicEdge]]:
+    extracted_nodes_bulk = await asyncio.gather(
+        *[
+            extract_nodes(llm_client, episode, previous_episodes)
+            for episode, previous_episodes in episode_tuples
+        ]
+    )
+
+    episodes, previous_episodes_list = (
+        [episode[0] for episode in episode_tuples],
+        [episode[1] for episode in episode_tuples],
+    )
+
+    extracted_edges_bulk = await asyncio.gather(
+        *[
+            extract_edges(llm_client, episode, extracted_nodes_bulk[i], previous_episodes_list[i])
+            for i, episode in enumerate(episodes)
+        ]
+    )
+
+    episodic_edges: list[EpisodicEdge] = []
+    for i, episode in enumerate(episodes):
+        episodic_edges += build_episodic_edges(extracted_nodes_bulk[i], episode, episode.created_at)
+
+    nodes: list[EntityNode] = []
+    for extracted_nodes in extracted_nodes_bulk:
+        nodes += extracted_nodes
+
+    edges: list[EntityEdge] = []
+    for extracted_edges in extracted_edges_bulk:
+        edges += extracted_edges
+
+    return nodes, edges, episodic_edges
+
+
+async def dedupe_nodes_bulk(
+    driver: AsyncDriver,
+    llm_client: LLMClient,
+    extracted_nodes: list[EntityNode],
+) -> tuple[list[EntityNode], dict[str, str]]:
+    # Compress nodes
+    nodes, uuid_map = node_name_match(extracted_nodes)
+
+    compressed_nodes, compressed_map = await compress_nodes(llm_client, nodes, uuid_map)
+
+    existing_nodes = await get_relevant_nodes(compressed_nodes, driver)
+
+    nodes, partial_uuid_map, _ = await dedupe_extracted_nodes(
+        llm_client, compressed_nodes, existing_nodes
+    )
+
+    compressed_map.update(partial_uuid_map)
+
+    return nodes, compressed_map
+
+
+async def dedupe_edges_bulk(
+    driver: AsyncDriver, llm_client: LLMClient, extracted_edges: list[EntityEdge]
+) -> list[EntityEdge]:
+    # Compress edges
+    compressed_edges = await compress_edges(llm_client, extracted_edges)
+
+    existing_edges = await get_relevant_edges(compressed_edges, driver)
+
+    edges = await dedupe_extracted_edges(llm_client, compressed_edges, existing_edges)
+
+    return edges
+
+
+def node_name_match(nodes: list[EntityNode]) -> tuple[list[EntityNode], dict[str, str]]:
+    uuid_map: dict[str, str] = {}
+    name_map: dict[str, EntityNode] = {}
+    for node in nodes:
+        if node.name in name_map:
+            uuid_map[node.uuid] = name_map[node.name].uuid
+            continue
+
+        name_map[node.name] = node
+
+    return [node for node in name_map.values()], uuid_map
+
+
+async def compress_nodes(
+    llm_client: LLMClient, nodes: list[EntityNode], uuid_map: dict[str, str]
+) -> tuple[list[EntityNode], dict[str, str]]:
+    if len(nodes) == 0:
+        return nodes, uuid_map
+
+    anchor = nodes[0]
+    nodes.sort(key=lambda node: dot(anchor.name_embedding or [], node.name_embedding or []))
+
+    node_chunks = [nodes[i : i + CHUNK_SIZE] for i in range(0, len(nodes), CHUNK_SIZE)]
+
+    results = await asyncio.gather(*[dedupe_node_list(llm_client, chunk) for chunk in node_chunks])
+
+    extended_map = dict(uuid_map)
+    compressed_nodes: list[EntityNode] = []
+    for node_chunk, uuid_map_chunk in results:
+        compressed_nodes += node_chunk
+        extended_map.update(uuid_map_chunk)
+
+    # Check if we have removed all duplicates
+    if len(compressed_nodes) == len(nodes):
+        compressed_uuid_map = compress_uuid_map(extended_map)
+        return compressed_nodes, compressed_uuid_map
+
+    return await compress_nodes(llm_client, compressed_nodes, extended_map)
+
+
+async def compress_edges(llm_client: LLMClient, edges: list[EntityEdge]) -> list[EntityEdge]:
+    if len(edges) == 0:
+        return edges
+
+    anchor = edges[0]
+    edges.sort(
+        key=lambda embedding: dot(anchor.fact_embedding or [], embedding.fact_embedding or [])
+    )
+
+    edge_chunks = [edges[i : i + CHUNK_SIZE] for i in range(0, len(edges), CHUNK_SIZE)]
+
+    results = await asyncio.gather(*[dedupe_edge_list(llm_client, chunk) for chunk in edge_chunks])
+
+    compressed_edges: list[EntityEdge] = []
+    for edge_chunk in results:
+        compressed_edges += edge_chunk
+
+    # Check if we have removed all duplicates
+    if len(compressed_edges) == len(edges):
+        return compressed_edges
+
+    return await compress_edges(llm_client, compressed_edges)
+
+
+def compress_uuid_map(uuid_map: dict[str, str]) -> dict[str, str]:
+    # make sure all uuid values aren't mapped to other uuids
+    compressed_map = {}
+    for key, uuid in uuid_map.items():
+        curr_value = uuid
+        while curr_value in uuid_map:
+            curr_value = uuid_map[curr_value]
+
+        compressed_map[key] = curr_value
+    return compressed_map
+
+
+E = typing.TypeVar('E', bound=Edge)
+
+
+def resolve_edge_pointers(edges: list[E], uuid_map: dict[str, str]):
+    for edge in edges:
+        source_uuid = edge.source_node_uuid
+        target_uuid = edge.target_node_uuid
+        edge.source_node_uuid = uuid_map.get(source_uuid, source_uuid)
+        edge.target_node_uuid = uuid_map.get(target_uuid, target_uuid)
+
+    return edges

graphiti_core/utils/maintenance/__init__.py

@@ -0,0 +1,16 @@
+from .edge_operations import build_episodic_edges, extract_edges
+from .graph_data_operations import (
+    clear_data,
+    retrieve_episodes,
+)
+from .node_operations import extract_nodes
+from .temporal_operations import invalidate_edges
+
+__all__ = [
+    'extract_edges',
+    'build_episodic_edges',
+    'extract_nodes',
+    'clear_data',
+    'retrieve_episodes',
+    'invalidate_edges',
+]