graphiti-core 0.1.0__py3-none-any.whl

graphiti_core/prompts/invalidate_edges.py

@@ -0,0 +1,74 @@
+"""
+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.
+"""
+
+from typing import Any, Protocol, TypedDict
+
+from .models import Message, PromptFunction, PromptVersion
+
+
+class Prompt(Protocol):
+    v1: PromptVersion
+
+
+class Versions(TypedDict):
+    v1: PromptFunction
+
+
+def v1(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are an AI assistant that helps determine which relationships in a knowledge graph should be invalidated based solely on explicit contradictions in newer information.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+               Based on the provided existing edges and new edges with their timestamps, determine which relationships, if any, should be marked as expired due to contradictions or updates in the newer edges.
+               Use the start and end dates of the edges to determine which edges are to be marked expired.
+                Only mark a relationship as invalid if there is clear evidence from other edges that the relationship is no longer true.
+                Do not invalidate relationships merely because they weren't mentioned in the episodes. You may use the current episode and previous episodes as well as the facts of each edge to understand the context of the relationships.
+
+                Previous Episodes:
+                {context['previous_episodes']}
+
+                Current Episode:
+                {context['current_episode']}
+
+                Existing Edges (sorted by timestamp, newest first):
+                {context['existing_edges']}
+
+                New Edges:
+                {context['new_edges']}
+
+                Each edge is formatted as: "UUID | SOURCE_NODE - EDGE_NAME - TARGET_NODE (fact: EDGE_FACT), START_DATE (END_DATE, optional))"
+
+                For each existing edge that should be invalidated, respond with a JSON object in the following format:
+                {{
+                    "invalidated_edges": [
+                        {{
+                            "edge_uuid": "The UUID of the edge to be invalidated (the part before the | character)",
+                            "fact": "Updated fact of the edge"
+                        }}
+                    ]
+                }}
+
+                If no relationships need to be invalidated based on these strict criteria, return an empty list for "invalidated_edges".
+            """,
+        ),
+    ]
+
+
+versions: Versions = {'v1': v1}

graphiti_core/prompts/lib.py

@@ -0,0 +1,122 @@
+"""
+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.
+"""
+
+from typing import Any, Protocol, TypedDict
+
+from .dedupe_edges import (
+    Prompt as DedupeEdgesPrompt,
+)
+from .dedupe_edges import (
+    Versions as DedupeEdgesVersions,
+)
+from .dedupe_edges import (
+    versions as dedupe_edges_versions,
+)
+from .dedupe_nodes import (
+    Prompt as DedupeNodesPrompt,
+)
+from .dedupe_nodes import (
+    Versions as DedupeNodesVersions,
+)
+from .dedupe_nodes import (
+    versions as dedupe_nodes_versions,
+)
+from .extract_edge_dates import (
+    Prompt as ExtractEdgeDatesPrompt,
+)
+from .extract_edge_dates import (
+    Versions as ExtractEdgeDatesVersions,
+)
+from .extract_edge_dates import (
+    versions as extract_edge_dates_versions,
+)
+from .extract_edges import (
+    Prompt as ExtractEdgesPrompt,
+)
+from .extract_edges import (
+    Versions as ExtractEdgesVersions,
+)
+from .extract_edges import (
+    versions as extract_edges_versions,
+)
+from .extract_nodes import (
+    Prompt as ExtractNodesPrompt,
+)
+from .extract_nodes import (
+    Versions as ExtractNodesVersions,
+)
+from .extract_nodes import (
+    versions as extract_nodes_versions,
+)
+from .invalidate_edges import (
+    Prompt as InvalidateEdgesPrompt,
+)
+from .invalidate_edges import (
+    Versions as InvalidateEdgesVersions,
+)
+from .invalidate_edges import (
+    versions as invalidate_edges_versions,
+)
+from .models import Message, PromptFunction
+
+
+class PromptLibrary(Protocol):
+    extract_nodes: ExtractNodesPrompt
+    dedupe_nodes: DedupeNodesPrompt
+    extract_edges: ExtractEdgesPrompt
+    dedupe_edges: DedupeEdgesPrompt
+    invalidate_edges: InvalidateEdgesPrompt
+    extract_edge_dates: ExtractEdgeDatesPrompt
+
+
+class PromptLibraryImpl(TypedDict):
+    extract_nodes: ExtractNodesVersions
+    dedupe_nodes: DedupeNodesVersions
+    extract_edges: ExtractEdgesVersions
+    dedupe_edges: DedupeEdgesVersions
+    invalidate_edges: InvalidateEdgesVersions
+    extract_edge_dates: ExtractEdgeDatesVersions
+
+
+class VersionWrapper:
+    def __init__(self, func: PromptFunction):
+        self.func = func
+
+    def __call__(self, context: dict[str, Any]) -> list[Message]:
+        return self.func(context)
+
+
+class PromptTypeWrapper:
+    def __init__(self, versions: dict[str, PromptFunction]):
+        for version, func in versions.items():
+            setattr(self, version, VersionWrapper(func))
+
+
+class PromptLibraryWrapper:
+    def __init__(self, library: PromptLibraryImpl):
+        for prompt_type, versions in library.items():
+            setattr(self, prompt_type, PromptTypeWrapper(versions))  # type: ignore[arg-type]
+
+
+PROMPT_LIBRARY_IMPL: PromptLibraryImpl = {
+    'extract_nodes': extract_nodes_versions,
+    'dedupe_nodes': dedupe_nodes_versions,
+    'extract_edges': extract_edges_versions,
+    'dedupe_edges': dedupe_edges_versions,
+    'invalidate_edges': invalidate_edges_versions,
+    'extract_edge_dates': extract_edge_dates_versions,
+}
+prompt_library: PromptLibrary = PromptLibraryWrapper(PROMPT_LIBRARY_IMPL)  # type: ignore[assignment]

graphiti_core/prompts/models.py

@@ -0,0 +1,31 @@
+"""
+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.
+"""
+
+from typing import Any, Callable, Protocol
+
+from pydantic import BaseModel
+
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+
+class PromptVersion(Protocol):
+    def __call__(self, context: dict[str, Any]) -> list[Message]: ...
+
+
+PromptFunction = Callable[[dict[str, Any]], list[Message]]

graphiti_core/search/search.py

@@ -0,0 +1,142 @@
+"""
+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 logging
+from datetime import datetime
+from enum import Enum
+from time import time
+
+from neo4j import AsyncDriver
+from pydantic import BaseModel, Field
+
+from graphiti_core.edges import EntityEdge
+from graphiti_core.llm_client.config import EMBEDDING_DIM
+from graphiti_core.nodes import EntityNode, EpisodicNode
+from graphiti_core.search.search_utils import (
+    edge_fulltext_search,
+    edge_similarity_search,
+    get_mentioned_nodes,
+    node_distance_reranker,
+    rrf,
+)
+from graphiti_core.utils import retrieve_episodes
+from graphiti_core.utils.maintenance.graph_data_operations import EPISODE_WINDOW_LEN
+
+logger = logging.getLogger(__name__)
+
+
+class SearchMethod(Enum):
+    cosine_similarity = 'cosine_similarity'
+    bm25 = 'bm25'
+
+
+class Reranker(Enum):
+    rrf = 'reciprocal_rank_fusion'
+    node_distance = 'node_distance'
+
+
+class SearchConfig(BaseModel):
+    num_edges: int = Field(default=10)
+    num_nodes: int = Field(default=10)
+    num_episodes: int = EPISODE_WINDOW_LEN
+    search_methods: list[SearchMethod]
+    reranker: Reranker | None
+
+
+class SearchResults(BaseModel):
+    episodes: list[EpisodicNode]
+    nodes: list[EntityNode]
+    edges: list[EntityEdge]
+
+
+async def hybrid_search(
+        driver: AsyncDriver,
+        embedder,
+        query: str,
+        timestamp: datetime,
+        config: SearchConfig,
+        center_node_uuid: str | None = None,
+) -> SearchResults:
+    start = time()
+
+    episodes = []
+    nodes = []
+    edges = []
+
+    search_results = []
+
+    if config.num_episodes > 0:
+        episodes.extend(await retrieve_episodes(driver, timestamp, config.num_episodes))
+        nodes.extend(await get_mentioned_nodes(driver, episodes))
+
+    if SearchMethod.bm25 in config.search_methods:
+        text_search = await edge_fulltext_search(query, driver, 2 * config.num_edges)
+        search_results.append(text_search)
+
+    if SearchMethod.cosine_similarity in config.search_methods:
+        query_text = query.replace('\n', ' ')
+        search_vector = (
+            (await embedder.create(input=[query_text], model='text-embedding-3-small'))
+            .data[0]
+            .embedding[:EMBEDDING_DIM]
+        )
+
+        similarity_search = await edge_similarity_search(
+            search_vector, driver, 2 * config.num_edges
+        )
+        search_results.append(similarity_search)
+
+    if len(search_results) > 1 and config.reranker is None:
+        logger.exception('Multiple searches enabled without a reranker')
+        raise Exception('Multiple searches enabled without a reranker')
+
+    else:
+        edge_uuid_map = {}
+        search_result_uuids = []
+
+        for result in search_results:
+            result_uuids = []
+            for edge in result:
+                result_uuids.append(edge.uuid)
+                edge_uuid_map[edge.uuid] = edge
+
+            search_result_uuids.append(result_uuids)
+
+        search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
+
+        reranked_uuids: list[str] = []
+        if config.reranker == Reranker.rrf:
+            reranked_uuids = rrf(search_result_uuids)
+        elif config.reranker == Reranker.node_distance:
+            if center_node_uuid is None:
+                logger.exception('No center node provided for Node Distance reranker')
+                raise Exception('No center node provided for Node Distance reranker')
+            reranked_uuids = await node_distance_reranker(
+                driver, search_result_uuids, center_node_uuid
+            )
+
+        reranked_edges = [edge_uuid_map[uuid] for uuid in reranked_uuids]
+        edges.extend(reranked_edges)
+
+    context = SearchResults(
+        episodes=episodes, nodes=nodes[: config.num_nodes], edges=edges[: config.num_edges]
+    )
+
+    end = time()
+
+    logger.info(f'search returned context for query {query} in {(end - start) * 1000} ms')
+
+    return context