graphiti-core 0.3.2__tar.gz → 0.3.4__tar.gz

{graphiti_core-0.3.2 → graphiti_core-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.3.2
+Version: 0.3.4
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -21,7 +21,7 @@ Description-Content-Type: text/markdown
 
 <div align="center">
 
-# Graphiti
+<img width="350" alt="Graphiti-ts-small" src="https://github.com/user-attachments/assets/bbd02947-e435-4a05-b25a-bbbac36d52c8">
 
 ## Temporal Knowledge Graphs for Agentic Applications
 
@@ -37,7 +37,9 @@ Description-Content-Type: text/markdown
 
 </div>
 
-Graphiti builds dynamic, temporally aware Knowledge Graphs that represent complex, evolving relationships between entities over time. Graphiti ingests both unstructured and structured data, and the resulting graph may be queried using a fusion of time, full-text, semantic, and graph algorithm approaches.
+Graphiti builds dynamic, temporally aware Knowledge Graphs that represent complex, evolving relationships between
+entities over time. Graphiti ingests both unstructured and structured data, and the resulting graph may be queried using
+a fusion of time, full-text, semantic, and graph algorithm approaches.
 
 <br />
 
@@ -47,25 +49,39 @@ Graphiti builds dynamic, temporally aware Knowledge Graphs that represent comple
 
 <br />
 
-Graphiti helps you create and query Knowledge Graphs that evolve over time. A knowledge graph is a network of interconnected facts, such as _“Kendra loves Adidas shoes.”_ Each fact is a “triplet” represented by two entities, or nodes (_”Kendra”_, _“Adidas shoes”_), and their relationship, or edge (_”loves”_). Knowledge Graphs have been explored extensively for information retrieval. What makes Graphiti unique is its ability to autonomously build a knowledge graph while handling changing relationships and maintaining historical context.
+Graphiti helps you create and query Knowledge Graphs that evolve over time. A knowledge graph is a network of
+interconnected facts, such as _“Kendra loves Adidas shoes.”_ Each fact is a “triplet” represented by two entities, or
+nodes (_”Kendra”_, _“Adidas shoes”_), and their relationship, or edge (_”loves”_). Knowledge Graphs have been explored
+extensively for information retrieval. What makes Graphiti unique is its ability to autonomously build a knowledge graph
+while handling changing relationships and maintaining historical context.
 
 With Graphiti, you can build LLM applications such as:
 
-- Assistants that learn from user interactions, fusing personal knowledge with dynamic data from business systems like CRMs and billing platforms.
+- Assistants that learn from user interactions, fusing personal knowledge with dynamic data from business systems like
+  CRMs and billing platforms.
 - Agents that autonomously execute complex tasks, reasoning with state changes from multiple dynamic sources.
 
-Graphiti supports a wide range of applications in sales, customer service, health, finance, and more, enabling long-term recall and state-based reasoning for both assistants and agents.
+Graphiti supports a wide range of applications in sales, customer service, health, finance, and more, enabling long-term
+recall and state-based reasoning for both assistants and agents.
 
 ## Why Graphiti?
 
-We were intrigued by Microsoft’s GraphRAG, which expanded on RAG text chunking by using a graph to better model a document corpus and making this representation available via semantic and graph search techniques. However, GraphRAG did not address our core problem: It's primarily designed for static documents and doesn't inherently handle temporal aspects of data.
-
-Graphiti is designed from the ground up to handle constantly changing information, hybrid semantic and graph search, and scale:
-
-- **Temporal Awareness:** Tracks changes in facts and relationships over time, enabling point-in-time queries. Graph edges include temporal metadata to record relationship lifecycles.
-- **Episodic Processing:** Ingests data as discrete episodes, maintaining data provenance and allowing incremental entity and relationship extraction.
-- **Hybrid Search:** Combines semantic and BM25 full-text search, with the ability to rerank results by distance from a central node e.g. “Kendra”.
-- **Scalable:** Designed for processing large datasets, with parallelization of LLM calls for bulk processing while preserving the chronology of events.
+We were intrigued by Microsoft’s GraphRAG, which expanded on RAG text chunking by using a graph to better model a
+document corpus and making this representation available via semantic and graph search techniques. However, GraphRAG did
+not address our core problem: It's primarily designed for static documents and doesn't inherently handle temporal
+aspects of data.
+
+Graphiti is designed from the ground up to handle constantly changing information, hybrid semantic and graph search, and
+scale:
+
+- **Temporal Awareness:** Tracks changes in facts and relationships over time, enabling point-in-time queries. Graph
+  edges include temporal metadata to record relationship lifecycles.
+- **Episodic Processing:** Ingests data as discrete episodes, maintaining data provenance and allowing incremental
+  entity and relationship extraction.
+- **Hybrid Search:** Combines semantic and BM25 full-text search, with the ability to rerank results by distance from a
+  central node e.g. “Kendra”.
+- **Scalable:** Designed for processing large datasets, with parallelization of LLM calls for bulk processing while
+  preserving the chronology of events.
 - **Supports Varied Sources:** Can ingest both unstructured text and structured JSON data.
 
 <p align="center">
@@ -91,7 +107,8 @@ Optional:
 - Anthropic or Groq API key (for alternative LLM providers)
 
 > [!TIP]
-> The simplest way to install Neo4j is via [Neo4j Desktop](https://neo4j.com/download/). It provides a user-friendly interface to manage Neo4j instances and databases.
+> The simplest way to install Neo4j is via [Neo4j Desktop](https://neo4j.com/download/). It provides a user-friendly
+> interface to manage Neo4j instances and databases.
 
 ```bash
 pip install graphiti-core
@@ -106,7 +123,8 @@ poetry add graphiti-core
 ## Quick Start
 
 > [!IMPORTANT]
-> Graphiti uses OpenAI for LLM inference and embedding. Ensure that an `OPENAI_API_KEY` is set in your environment. Support for Anthropic and Groq LLM inferences is available, too.
+> Graphiti uses OpenAI for LLM inference and embedding. Ensure that an `OPENAI_API_KEY` is set in your environment.
+> Support for Anthropic and Groq LLM inferences is available, too.
 
 ```python
 from graphiti_core import Graphiti
@@ -140,25 +158,25 @@ for i, episode in enumerate(episodes):
 results = await graphiti.search('Who was the California Attorney General?')
 [
     EntityEdge(
-    │   uuid='3133258f738e487383f07b04e15d4ac0',
-    │   source_node_uuid='2a85789b318d4e418050506879906e62',
-    │   target_node_uuid='baf7781f445945989d6e4f927f881556',
-    │   created_at=datetime.datetime(2024, 8, 26, 13, 13, 24, 861097),
-    │   name='HELD_POSITION',
-        # the fact reflects the updated state that Harris is
-        # no longer the AG of California
-    │   fact='Kamala Harris was the Attorney General of California',
-    │   fact_embedding=[
-    │   │   -0.009955154731869698,
-    │       ...
-    │   │   0.00784289836883545
-    │   ],
-    │   episodes=['b43e98ad0a904088a76c67985caecc22'],
-    │   expired_at=datetime.datetime(2024, 8, 26, 20, 18, 1, 53812),
-        # These dates represent the date this edge was true.
-    │   valid_at=datetime.datetime(2011, 1, 3, 0, 0, tzinfo=<UTC>),
-    │   invalid_at=datetime.datetime(2017, 1, 3, 0, 0, tzinfo=<UTC>)
-    )
+│   uuid = '3133258f738e487383f07b04e15d4ac0',
+│   source_node_uuid = '2a85789b318d4e418050506879906e62',
+│   target_node_uuid = 'baf7781f445945989d6e4f927f881556',
+│   created_at = datetime.datetime(2024, 8, 26, 13, 13, 24, 861097),
+│   name = 'HELD_POSITION',
+# the fact reflects the updated state that Harris is
+# no longer the AG of California
+│   fact = 'Kamala Harris was the Attorney General of California',
+│   fact_embedding = [
+│   │   -0.009955154731869698,
+│       ...
+│   │   0.00784289836883545
+│],
+│   episodes = ['b43e98ad0a904088a76c67985caecc22'],
+│   expired_at = datetime.datetime(2024, 8, 26, 20, 18, 1, 53812),
+# These dates represent the date this edge was true.
+│   valid_at = datetime.datetime(2011, 1, 3, 0, 0, tzinfo= < UTC >),
+│   invalid_at = datetime.datetime(2017, 1, 3, 0, 0, tzinfo= < UTC >)
+)
 ]
 
 # Rerank search results based on graph distance
@@ -191,14 +209,16 @@ Graphiti is under active development. We aim to maintain API stability while wor
 - [ ] Achieving good performance with different LLM and embedding models
 - [ ] Creating a dedicated embedder interface
 - [ ] Supporting custom graph schemas:
-  - Allow developers to provide their own defined node and edge classes when ingesting episodes
-  - Enable more flexible knowledge representation tailored to specific use cases
+    - Allow developers to provide their own defined node and edge classes when ingesting episodes
+    - Enable more flexible knowledge representation tailored to specific use cases
 - [ ] Enhancing retrieval capabilities with more robust and configurable options
 - [ ] Expanding test coverage to ensure reliability and catch edge cases
 
 ## Contributing
 
-We encourage and appreciate all forms of contributions, whether it's code, documentation, addressing GitHub Issues, or answering questions in the Graphiti Discord channel. For detailed guidelines on code contributions, please refer to [CONTRIBUTING](CONTRIBUTING.md).
+We encourage and appreciate all forms of contributions, whether it's code, documentation, addressing GitHub Issues, or
+answering questions in the Graphiti Discord channel. For detailed guidelines on code contributions, please refer
+to [CONTRIBUTING](CONTRIBUTING.md).
 
 ## Support
 

{graphiti_core-0.3.2 → graphiti_core-0.3.4}/README.md

@@ -1,6 +1,6 @@
 <div align="center">
 
-# Graphiti
+<img width="350" alt="Graphiti-ts-small" src="https://github.com/user-attachments/assets/bbd02947-e435-4a05-b25a-bbbac36d52c8">
 
 ## Temporal Knowledge Graphs for Agentic Applications
 
@@ -16,7 +16,9 @@
 
 </div>
 
-Graphiti builds dynamic, temporally aware Knowledge Graphs that represent complex, evolving relationships between entities over time. Graphiti ingests both unstructured and structured data, and the resulting graph may be queried using a fusion of time, full-text, semantic, and graph algorithm approaches.
+Graphiti builds dynamic, temporally aware Knowledge Graphs that represent complex, evolving relationships between
+entities over time. Graphiti ingests both unstructured and structured data, and the resulting graph may be queried using
+a fusion of time, full-text, semantic, and graph algorithm approaches.
 
 <br />
 
@@ -26,25 +28,39 @@ Graphiti builds dynamic, temporally aware Knowledge Graphs that represent comple
 
 <br />
 
-Graphiti helps you create and query Knowledge Graphs that evolve over time. A knowledge graph is a network of interconnected facts, such as _“Kendra loves Adidas shoes.”_ Each fact is a “triplet” represented by two entities, or nodes (_”Kendra”_, _“Adidas shoes”_), and their relationship, or edge (_”loves”_). Knowledge Graphs have been explored extensively for information retrieval. What makes Graphiti unique is its ability to autonomously build a knowledge graph while handling changing relationships and maintaining historical context.
+Graphiti helps you create and query Knowledge Graphs that evolve over time. A knowledge graph is a network of
+interconnected facts, such as _“Kendra loves Adidas shoes.”_ Each fact is a “triplet” represented by two entities, or
+nodes (_”Kendra”_, _“Adidas shoes”_), and their relationship, or edge (_”loves”_). Knowledge Graphs have been explored
+extensively for information retrieval. What makes Graphiti unique is its ability to autonomously build a knowledge graph
+while handling changing relationships and maintaining historical context.
 
 With Graphiti, you can build LLM applications such as:
 
-- Assistants that learn from user interactions, fusing personal knowledge with dynamic data from business systems like CRMs and billing platforms.
+- Assistants that learn from user interactions, fusing personal knowledge with dynamic data from business systems like
+  CRMs and billing platforms.
 - Agents that autonomously execute complex tasks, reasoning with state changes from multiple dynamic sources.
 
-Graphiti supports a wide range of applications in sales, customer service, health, finance, and more, enabling long-term recall and state-based reasoning for both assistants and agents.
+Graphiti supports a wide range of applications in sales, customer service, health, finance, and more, enabling long-term
+recall and state-based reasoning for both assistants and agents.
 
 ## Why Graphiti?
 
-We were intrigued by Microsoft’s GraphRAG, which expanded on RAG text chunking by using a graph to better model a document corpus and making this representation available via semantic and graph search techniques. However, GraphRAG did not address our core problem: It's primarily designed for static documents and doesn't inherently handle temporal aspects of data.
-
-Graphiti is designed from the ground up to handle constantly changing information, hybrid semantic and graph search, and scale:
-
-- **Temporal Awareness:** Tracks changes in facts and relationships over time, enabling point-in-time queries. Graph edges include temporal metadata to record relationship lifecycles.
-- **Episodic Processing:** Ingests data as discrete episodes, maintaining data provenance and allowing incremental entity and relationship extraction.
-- **Hybrid Search:** Combines semantic and BM25 full-text search, with the ability to rerank results by distance from a central node e.g. “Kendra”.
-- **Scalable:** Designed for processing large datasets, with parallelization of LLM calls for bulk processing while preserving the chronology of events.
+We were intrigued by Microsoft’s GraphRAG, which expanded on RAG text chunking by using a graph to better model a
+document corpus and making this representation available via semantic and graph search techniques. However, GraphRAG did
+not address our core problem: It's primarily designed for static documents and doesn't inherently handle temporal
+aspects of data.
+
+Graphiti is designed from the ground up to handle constantly changing information, hybrid semantic and graph search, and
+scale:
+
+- **Temporal Awareness:** Tracks changes in facts and relationships over time, enabling point-in-time queries. Graph
+  edges include temporal metadata to record relationship lifecycles.
+- **Episodic Processing:** Ingests data as discrete episodes, maintaining data provenance and allowing incremental
+  entity and relationship extraction.
+- **Hybrid Search:** Combines semantic and BM25 full-text search, with the ability to rerank results by distance from a
+  central node e.g. “Kendra”.
+- **Scalable:** Designed for processing large datasets, with parallelization of LLM calls for bulk processing while
+  preserving the chronology of events.
 - **Supports Varied Sources:** Can ingest both unstructured text and structured JSON data.
 
 <p align="center">
@@ -70,7 +86,8 @@ Optional:
 - Anthropic or Groq API key (for alternative LLM providers)
 
 > [!TIP]
-> The simplest way to install Neo4j is via [Neo4j Desktop](https://neo4j.com/download/). It provides a user-friendly interface to manage Neo4j instances and databases.
+> The simplest way to install Neo4j is via [Neo4j Desktop](https://neo4j.com/download/). It provides a user-friendly
+> interface to manage Neo4j instances and databases.
 
 ```bash
 pip install graphiti-core
@@ -85,7 +102,8 @@ poetry add graphiti-core
 ## Quick Start
 
 > [!IMPORTANT]
-> Graphiti uses OpenAI for LLM inference and embedding. Ensure that an `OPENAI_API_KEY` is set in your environment. Support for Anthropic and Groq LLM inferences is available, too.
+> Graphiti uses OpenAI for LLM inference and embedding. Ensure that an `OPENAI_API_KEY` is set in your environment.
+> Support for Anthropic and Groq LLM inferences is available, too.
 
 ```python
 from graphiti_core import Graphiti
@@ -119,25 +137,25 @@ for i, episode in enumerate(episodes):
 results = await graphiti.search('Who was the California Attorney General?')
 [
     EntityEdge(
-    │   uuid='3133258f738e487383f07b04e15d4ac0',
-    │   source_node_uuid='2a85789b318d4e418050506879906e62',
-    │   target_node_uuid='baf7781f445945989d6e4f927f881556',
-    │   created_at=datetime.datetime(2024, 8, 26, 13, 13, 24, 861097),
-    │   name='HELD_POSITION',
-        # the fact reflects the updated state that Harris is
-        # no longer the AG of California
-    │   fact='Kamala Harris was the Attorney General of California',
-    │   fact_embedding=[
-    │   │   -0.009955154731869698,
-    │       ...
-    │   │   0.00784289836883545
-    │   ],
-    │   episodes=['b43e98ad0a904088a76c67985caecc22'],
-    │   expired_at=datetime.datetime(2024, 8, 26, 20, 18, 1, 53812),
-        # These dates represent the date this edge was true.
-    │   valid_at=datetime.datetime(2011, 1, 3, 0, 0, tzinfo=<UTC>),
-    │   invalid_at=datetime.datetime(2017, 1, 3, 0, 0, tzinfo=<UTC>)
-    )
+│   uuid = '3133258f738e487383f07b04e15d4ac0',
+│   source_node_uuid = '2a85789b318d4e418050506879906e62',
+│   target_node_uuid = 'baf7781f445945989d6e4f927f881556',
+│   created_at = datetime.datetime(2024, 8, 26, 13, 13, 24, 861097),
+│   name = 'HELD_POSITION',
+# the fact reflects the updated state that Harris is
+# no longer the AG of California
+│   fact = 'Kamala Harris was the Attorney General of California',
+│   fact_embedding = [
+│   │   -0.009955154731869698,
+│       ...
+│   │   0.00784289836883545
+│],
+│   episodes = ['b43e98ad0a904088a76c67985caecc22'],
+│   expired_at = datetime.datetime(2024, 8, 26, 20, 18, 1, 53812),
+# These dates represent the date this edge was true.
+│   valid_at = datetime.datetime(2011, 1, 3, 0, 0, tzinfo= < UTC >),
+│   invalid_at = datetime.datetime(2017, 1, 3, 0, 0, tzinfo= < UTC >)
+)
 ]
 
 # Rerank search results based on graph distance
@@ -170,14 +188,16 @@ Graphiti is under active development. We aim to maintain API stability while wor
 - [ ] Achieving good performance with different LLM and embedding models
 - [ ] Creating a dedicated embedder interface
 - [ ] Supporting custom graph schemas:
-  - Allow developers to provide their own defined node and edge classes when ingesting episodes
-  - Enable more flexible knowledge representation tailored to specific use cases
+    - Allow developers to provide their own defined node and edge classes when ingesting episodes
+    - Enable more flexible knowledge representation tailored to specific use cases
 - [ ] Enhancing retrieval capabilities with more robust and configurable options
 - [ ] Expanding test coverage to ensure reliability and catch edge cases
 
 ## Contributing
 
-We encourage and appreciate all forms of contributions, whether it's code, documentation, addressing GitHub Issues, or answering questions in the Graphiti Discord channel. For detailed guidelines on code contributions, please refer to [CONTRIBUTING](CONTRIBUTING.md).
+We encourage and appreciate all forms of contributions, whether it's code, documentation, addressing GitHub Issues, or
+answering questions in the Graphiti Discord channel. For detailed guidelines on code contributions, please refer
+to [CONTRIBUTING](CONTRIBUTING.md).
 
 ## Support
 

{graphiti_core-0.3.2 → graphiti_core-0.3.4}/graphiti_core/edges.py

@@ -33,7 +33,7 @@ logger = logging.getLogger(__name__)
 
 
 class Edge(BaseModel, ABC):
-    uuid: str = Field(default_factory=lambda: uuid4().hex)
+    uuid: str = Field(default_factory=lambda: str(uuid4()))
     group_id: str | None = Field(description='partition of the graph')
     source_node_uuid: str
     target_node_uuid: str
@@ -104,18 +104,62 @@ class EpisodicEdge(Edge):
 
         edges = [get_episodic_edge_from_record(record) for record in records]
 
-        logger.info(f'Found Edge: {uuid}')
         if len(edges) == 0:
             raise EdgeNotFoundError(uuid)
         return edges[0]
 
+    @classmethod
+    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
+        WHERE e.uuid IN $uuids
+        RETURN
+            e.uuid As uuid,
+            e.group_id AS group_id,
+            n.uuid AS source_node_uuid, 
+            m.uuid AS target_node_uuid, 
+            e.created_at AS created_at
+        """,
+            uuids=uuids,
+        )
+
+        edges = [get_episodic_edge_from_record(record) for record in records]
+
+        if len(edges) == 0:
+            raise EdgeNotFoundError(uuids[0])
+        return edges
+
+    @classmethod
+    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str | None]):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
+        WHERE e.group_id IN $group_ids
+        RETURN
+            e.uuid As uuid,
+            e.group_id AS group_id,
+            n.uuid AS source_node_uuid, 
+            m.uuid AS target_node_uuid, 
+            e.created_at AS created_at
+        """,
+            group_ids=group_ids,
+        )
+
+        edges = [get_episodic_edge_from_record(record) for record in records]
+        uuids = [edge.uuid for edge in edges]
+
+        if len(edges) == 0:
+            raise EdgeNotFoundError(uuids[0])
+        return edges
+
 
 class EntityEdge(Edge):
     name: str = Field(description='name of the edge, relation name')
     fact: str = Field(description='fact representing the edge and nodes that it connects')
     fact_embedding: list[float] | None = Field(default=None, description='embedding of the fact')
-    episodes: list[str] | None = Field(
-        default=None,
+    episodes: list[str] = Field(
+        default=[],
         description='list of episode ids that reference these entity edges',
     )
     expired_at: datetime | None = Field(
@@ -192,11 +236,69 @@ class EntityEdge(Edge):
 
         edges = [get_entity_edge_from_record(record) for record in records]
 
-        logger.info(f'Found Edge: {uuid}')
         if len(edges) == 0:
             raise EdgeNotFoundError(uuid)
         return edges[0]
 
+    @classmethod
+    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
+        WHERE e.uuid IN $uuids
+        RETURN
+            e.uuid AS uuid,
+            n.uuid AS source_node_uuid,
+            m.uuid AS target_node_uuid,
+            e.created_at AS created_at,
+            e.name AS name,
+            e.group_id AS group_id,
+            e.fact AS fact,
+            e.fact_embedding AS fact_embedding,
+            e.episodes AS episodes,
+            e.expired_at AS expired_at,
+            e.valid_at AS valid_at,
+            e.invalid_at AS invalid_at
+        """,
+            uuids=uuids,
+        )
+
+        edges = [get_entity_edge_from_record(record) for record in records]
+
+        if len(edges) == 0:
+            raise EdgeNotFoundError(uuids[0])
+        return edges
+
+    @classmethod
+    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str | None]):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
+        WHERE e.group_id IN $group_ids
+        RETURN
+            e.uuid AS uuid,
+            n.uuid AS source_node_uuid,
+            m.uuid AS target_node_uuid,
+            e.created_at AS created_at,
+            e.name AS name,
+            e.group_id AS group_id,
+            e.fact AS fact,
+            e.fact_embedding AS fact_embedding,
+            e.episodes AS episodes,
+            e.expired_at AS expired_at,
+            e.valid_at AS valid_at,
+            e.invalid_at AS invalid_at
+        """,
+            group_ids=group_ids,
+        )
+
+        edges = [get_entity_edge_from_record(record) for record in records]
+        uuids = [edge.uuid for edge in edges]
+
+        if len(edges) == 0:
+            raise EdgeNotFoundError(uuids[0])
+        return edges
+
 
 class CommunityEdge(Edge):
     async def save(self, driver: AsyncDriver):
@@ -235,10 +337,48 @@ class CommunityEdge(Edge):
 
         edges = [get_community_edge_from_record(record) for record in records]
 
-        logger.info(f'Found Edge: {uuid}')
-
         return edges[0]
 
+    @classmethod
+    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Community)-[e:HAS_MEMBER]->(m:Entity | Community)
+        WHERE e.uuid IN $uuids
+        RETURN
+            e.uuid As uuid,
+            e.group_id AS group_id,
+            n.uuid AS source_node_uuid, 
+            m.uuid AS target_node_uuid, 
+            e.created_at AS created_at
+        """,
+            uuids=uuids,
+        )
+
+        edges = [get_community_edge_from_record(record) for record in records]
+
+        return edges
+
+    @classmethod
+    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str | None]):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Community)-[e:HAS_MEMBER]->(m:Entity | Community)
+        WHERE e.group_id IN $group_ids
+        RETURN
+            e.uuid As uuid,
+            e.group_id AS group_id,
+            n.uuid AS source_node_uuid, 
+            m.uuid AS target_node_uuid, 
+            e.created_at AS created_at
+        """,
+            group_ids=group_ids,
+        )
+
+        edges = [get_community_edge_from_record(record) for record in records]
+
+        return edges
+
 
 # Edge helpers
 def get_episodic_edge_from_record(record: Any) -> EpisodicEdge:

{graphiti_core-0.3.2 → graphiti_core-0.3.4}/graphiti_core/graphiti.py

@@ -35,6 +35,8 @@ from graphiti_core.search.search_config_recipes import (
 )
 from graphiti_core.search.search_utils import (
     RELEVANT_SCHEMA_LIMIT,
+    get_communities_by_nodes,
+    get_mentioned_nodes,
     get_relevant_edges,
     get_relevant_nodes,
 )
@@ -54,6 +56,7 @@ from graphiti_core.utils.bulk_utils import (
 from graphiti_core.utils.maintenance.community_operations import (
     build_communities,
     remove_communities,
+    update_community,
 )
 from graphiti_core.utils.maintenance.edge_operations import (
     extract_edges,
@@ -74,7 +77,14 @@ load_dotenv()
 
 
 class Graphiti:
-    def __init__(self, uri: str, user: str, password: str, llm_client: LLMClient | None = None):
+    def __init__(
+        self,
+        uri: str,
+        user: str,
+        password: str,
+        llm_client: LLMClient | None = None,
+        store_raw_episode_content: bool = True,
+    ):
         """
         Initialize a Graphiti instance.
 
@@ -113,6 +123,7 @@ class Graphiti:
         """
         self.driver = AsyncGraphDatabase.driver(uri, auth=(user, password))
         self.database = 'neo4j'
+        self.store_raw_episode_content = store_raw_episode_content
         if llm_client:
             self.llm_client = llm_client
         else:
@@ -147,8 +158,8 @@ class Graphiti:
                 # Use graphiti...
             finally:
                 graphiti.close()
-        self.driver.close()
         """
+        self.driver.close()
 
     async def build_indices_and_constraints(self):
         """
@@ -224,6 +235,7 @@ class Graphiti:
         source: EpisodeType = EpisodeType.message,
         group_id: str | None = None,
         uuid: str | None = None,
+        update_communities: bool = False,
     ):
         """
         Process an episode and update the graph.
@@ -247,6 +259,8 @@ class Graphiti:
             An id for the graph partition the episode is a part of.
         uuid : str | None
             Optional uuid of the episode.
+        update_communities : bool
+            Optional. Whether to update communities with new node information
 
         Returns
         -------
@@ -272,7 +286,6 @@ class Graphiti:
         try:
             start = time()
 
-            nodes: list[EntityNode] = []
             entity_edges: list[EntityEdge] = []
             embedder = self.llm_client.get_embedder()
             now = datetime.now()
@@ -291,6 +304,8 @@ class Graphiti:
                 valid_at=reference_time,
             )
             episode.uuid = uuid if uuid is not None else episode.uuid
+            if not self.store_raw_episode_content:
+                episode.content = ''
 
             # Extract entities as nodes
 
@@ -319,7 +334,7 @@ class Graphiti:
                 ),
             )
             logger.info(f'Adjusted mentioned nodes: {[(n.name, n.uuid) for n in mentioned_nodes]}')
-            nodes.extend(mentioned_nodes)
+            nodes = mentioned_nodes
 
             extracted_edges_with_resolved_pointers = resolve_edge_pointers(
                 extracted_edges, uuid_map
@@ -409,12 +424,22 @@ class Graphiti:
 
             logger.info(f'Built episodic edges: {episodic_edges}')
 
+            episode.entity_edges = [edge.uuid for edge in entity_edges]
+
             # Future optimization would be using batch operations to save nodes and edges
             await episode.save(self.driver)
             await asyncio.gather(*[node.save(self.driver) for node in nodes])
             await asyncio.gather(*[edge.save(self.driver) for edge in episodic_edges])
             await asyncio.gather(*[edge.save(self.driver) for edge in entity_edges])
 
+            # Update any communities
+            if update_communities:
+                await asyncio.gather(
+                    *[
+                        update_community(self.driver, self.llm_client, embedder, node)
+                        for node in nodes
+                    ]
+                )
             end = time()
             logger.info(f'Completed add_episode in {(end - start) * 1000} ms')
 
@@ -554,7 +579,7 @@ class Graphiti:
         center_node_uuid: str | None = None,
         group_ids: list[str | None] | None = None,
         num_results=DEFAULT_SEARCH_LIMIT,
-    ):
+    ) -> list[EntityEdge]:
         """
         Perform a hybrid search on the knowledge graph.
 
@@ -569,7 +594,7 @@ class Graphiti:
             Facts will be reranked based on proximity to this node
         group_ids : list[str | None] | None, optional
             The graph partitions to return data from.
-        limit : int, optional
+        num_results : int, optional
             The maximum number of results to return. Defaults to 10.
 
         Returns
@@ -668,3 +693,19 @@ class Graphiti:
             await search(self.driver, embedder, query, group_ids, search_config, center_node_uuid)
         ).nodes
         return nodes
+
+
+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[EntityEdge] = [edge for lst in edges_list for edge in lst]
+
+    nodes = await get_mentioned_nodes(self.driver, episodes)
+
+    communities = await get_communities_by_nodes(self.driver, nodes)
+
+    return SearchResults(edges=edges, nodes=nodes, communities=communities)