graphiti-core 0.8.0__tar.gz → 0.8.2__tar.gz

{graphiti_core-0.8.0 → graphiti_core-0.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.8.0
+Version: 0.8.2
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -29,7 +29,7 @@ Description-Content-Type: text/markdown
 <h1 align="center">
 Graphiti
 </h1>
-<h2 align="center"> Temporal Knowledge Graphs for Agentic Applications</h2>
+<h2 align="center"> Build Real-Time Knowledge Graphs for AI Agents</h2>
 <br />
 
 [![Discord](https://dcbadge.vercel.app/api/server/W8Kw6bsgXQ?style=flat)](https://discord.com/invite/W8Kw6bsgXQ)
@@ -38,12 +38,16 @@ Graphiti
 [![MyPy Check](https://github.com/getzep/Graphiti/actions/workflows/typecheck.yml/badge.svg)](https://github.com/getzep/Graphiti/actions/workflows/typecheck.yml)
 [![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/getzep/Graphiti)
 
-:star: *Help us reach more developers and grow the Graphiti community. Star this repo!*
+:star: _Help us reach more developers and grow the Graphiti community. Star this repo!_
 <br />
 
-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 is a framework for building and querying temporally-aware knowledge graphs, specifically tailored for AI agents operating in dynamic environments. Unlike traditional retrieval-augmented generation (RAG) methods, Graphiti continuously integrates user interactions, structured and unstructured enterprise data, and external information into a coherent, queryable graph. The framework supports incremental data updates, efficient retrieval, and precise historical queries without requiring complete graph recomputation, making it suitable for developing interactive, context-aware AI applications.
+
+Use Graphiti to:
+
+- Integrate and maintain dynamic user interactions and business data.
+- Facilitate state-based reasoning and task automation for agents.
+- Query complex, evolving data with semantic, keyword, and graph-based search methods.
 
 <br />
 
@@ -53,31 +57,21 @@ a fusion of time, full-text, semantic, and graph algorithm approaches.
 
 <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
+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.
-- 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 and Zep Memory
 
-Graphiti powers the core of [Zep's memory layer](https://www.getzep.com) for LLM-powered Assistants and Agents.
+Graphiti powers the core of [Zep's memory layer](https://www.getzep.com) for AI Agents.
 
 Using Graphiti, we've demonstrated Zep is
 the [State of the Art in Agent Memory](https://blog.getzep.com/state-of-the-art-agent-memory/).
 
 Read our paper: [Zep: A Temporal Knowledge Graph Architecture for Agent Memory](https://arxiv.org/abs/2501.13956).
 
-We're excited to open-source Graphiti, believing its potential reaches far beyond memory applications.
+We're excited to open-source Graphiti, believing its potential reaches far beyond AI memory applications.
 
 <p align="center">
     <a href="https://arxiv.org/abs/2501.13956"><img src="images/arxiv-screenshot.png" alt="Zep: A Temporal Knowledge Graph Architecture for Agent Memory" width="700px"></a>
@@ -85,34 +79,41 @@ We're excited to open-source Graphiti, believing its potential reaches far beyon
 
 ## 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:
+Traditional RAG approaches often rely on batch processing and static data summarization, making them inefficient for frequently changing data. Graphiti addresses these challenges by providing:
 
-- **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.
+- **Real-Time Incremental Updates:** Immediate integration of new data episodes without batch recomputation.
+- **Bi-Temporal Data Model:** Explicit tracking of event occurrence and ingestion times, allowing accurate point-in-time queries.
+- **Efficient Hybrid Retrieval:** Combines semantic embeddings, keyword (BM25), and graph traversal to achieve low-latency queries without reliance on LLM summarization.
+- **Custom Entity Definitions:** Flexible ontology creation and support for developer-defined entities through straightforward Pydantic models.
+- **Scalability:** Efficiently manages large datasets with parallel processing, suitable for enterprise environments.
 
 <p align="center">
     <img src="/images/graphiti-intro-slides-stock-2.gif" alt="Graphiti structured + unstructured demo" width="700px">   
 </p>
 
+## Graphiti vs. GraphRAG
+
+| Aspect                     | GraphRAG                              | Graphiti                                         |
+| -------------------------- | ------------------------------------- | ------------------------------------------------ |
+| **Primary Use**            | Static document summarization         | Dynamic data management                          |
+| **Data Handling**          | Batch-oriented processing             | Continuous, incremental updates                  |
+| **Knowledge Structure**    | Entity clusters & community summaries | Episodic data, semantic entities, communities    |
+| **Retrieval Method**       | Sequential LLM summarization          | Hybrid semantic, keyword, and graph-based search |
+| **Adaptability**           | Low                                   | High                                             |
+| **Temporal Handling**      | Basic timestamp tracking              | Explicit bi-temporal tracking                    |
+| **Contradiction Handling** | LLM-driven summarization judgments    | Temporal edge invalidation                       |
+| **Query Latency**          | Seconds to tens of seconds            | Typically sub-second latency                     |
+| **Custom Entity Types**    | No                                    | Yes, customizable                                |
+| **Scalability**            | Moderate                              | High, optimized for large datasets               |
+
+Graphiti is specifically designed to address the challenges of dynamic and frequently updated datasets, making it particularly suitable for applications requiring real-time interaction and precise historical queries.
+
 ## Installation
 
 Requirements:
 
 - Python 3.10 or higher
-- Neo4j 5.26 or higher
+- Neo4j 5.26 or higher (serves as the embeddings storage backend)
 - OpenAI API key (for LLM inference and embedding)
 
 Optional:
@@ -145,7 +146,7 @@ from graphiti_core import Graphiti
 from graphiti_core.nodes import EpisodeType
 from datetime import datetime, timezone
 
-# Initialize Graphiti
+# Initialize Graphiti as Your Memory Layer
 graphiti = Graphiti("bolt://localhost:7687", "neo4j", "password")
 
 # Initialize the graph database with Graphiti's indices. This only needs to be done once.
@@ -166,7 +167,7 @@ for i, episode in enumerate(episodes):
         reference_time=datetime.now(timezone.utc)
     )
 
-# Search the graph
+# Search the graph for semantic memory retrieval
 # Execute a hybrid search combining semantic similarity and BM25 retrieval
 # Results are combined and reranked using Reciprocal Rank Fusion
 results = await graphiti.search('Who was the California Attorney General?')
@@ -198,11 +199,12 @@ results = await graphiti.search('Who was the California Attorney General?')
 # Results are weighted by their proximity, with distant edges receiving lower scores.
 await graphiti.search('Who was the California Attorney General?', center_node_uuid)
 
-# Close the connection
+# Close the connection when chat state management is complete
 graphiti.close()
 ```
 
 ## Graph Service
+
 The `server` directory contains an API service for interacting with the Graphiti API. It is built using FastAPI.
 
 Please see the [server README](./server/README.md) for more information.
@@ -218,6 +220,54 @@ to enable Neo4j's parallel runtime feature for several of our search queries.
 Note that this feature is not supported for Neo4j Community edition or for smaller AuraDB instances,
 as such this feature is off by default.
 
+## Using Graphiti with Azure OpenAI
+
+Graphiti supports Azure OpenAI for both LLM inference and embeddings. To use Azure OpenAI, you'll need to configure both the LLM client and embedder with your Azure OpenAI credentials.
+
+```python
+from openai import AsyncAzureOpenAI
+from graphiti_core import Graphiti
+from graphiti_core.llm_client import OpenAIClient
+from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
+from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
+
+# Azure OpenAI configuration
+api_key = "<your-api-key>"
+api_version = "<your-api-version>"
+azure_endpoint = "<your-azure-endpoint>"
+
+# Create Azure OpenAI client for LLM
+azure_openai_client = AsyncAzureOpenAI(
+    api_key=api_key,
+    api_version=api_version,
+    azure_endpoint=azure_endpoint
+)
+
+# Initialize Graphiti with Azure OpenAI clients
+graphiti = Graphiti(
+    "bolt://localhost:7687",
+    "neo4j",
+    "password",
+    llm_client=OpenAIClient(
+        client=azure_openai_client
+    ),
+    embedder=OpenAIEmbedder(
+        config=OpenAIEmbedderConfig(
+            embedding_model="text-embedding-3-small"  # Use your Azure deployed embedding model name
+        ),
+        client=azure_openai_client
+    ),
+    # Optional: Configure the OpenAI cross encoder with Azure OpenAI
+    cross_encoder=OpenAIRerankerClient(
+        client=azure_openai_client
+    )
+)
+
+# Now you can use Graphiti with Azure OpenAI
+```
+
+Make sure to replace the placeholder values with your actual Azure OpenAI credentials and specify the correct embedding model name that's deployed in your Azure OpenAI service.
+
 ## Documentation
 
 - [Guides and API documentation](https://help.getzep.com/graphiti).
@@ -228,10 +278,11 @@ as such this feature is off by default.
 
 Graphiti is under active development. We aim to maintain API stability while working on:
 
-- [ ] 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
+- [x] 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
 - [x] Enhancing retrieval capabilities with more robust and configurable options
+- [ ] Graphiti MCP Server
 - [ ] Expanding test coverage to ensure reliability and catch edge cases
 
 ## Contributing

{graphiti_core-0.8.0 → graphiti_core-0.8.2}/README.md

@@ -7,7 +7,7 @@
 <h1 align="center">
 Graphiti
 </h1>
-<h2 align="center"> Temporal Knowledge Graphs for Agentic Applications</h2>
+<h2 align="center"> Build Real-Time Knowledge Graphs for AI Agents</h2>
 <br />
 
 [![Discord](https://dcbadge.vercel.app/api/server/W8Kw6bsgXQ?style=flat)](https://discord.com/invite/W8Kw6bsgXQ)
@@ -16,12 +16,16 @@ Graphiti
 [![MyPy Check](https://github.com/getzep/Graphiti/actions/workflows/typecheck.yml/badge.svg)](https://github.com/getzep/Graphiti/actions/workflows/typecheck.yml)
 [![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/getzep/Graphiti)
 
-:star: *Help us reach more developers and grow the Graphiti community. Star this repo!*
+:star: _Help us reach more developers and grow the Graphiti community. Star this repo!_
 <br />
 
-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 is a framework for building and querying temporally-aware knowledge graphs, specifically tailored for AI agents operating in dynamic environments. Unlike traditional retrieval-augmented generation (RAG) methods, Graphiti continuously integrates user interactions, structured and unstructured enterprise data, and external information into a coherent, queryable graph. The framework supports incremental data updates, efficient retrieval, and precise historical queries without requiring complete graph recomputation, making it suitable for developing interactive, context-aware AI applications.
+
+Use Graphiti to:
+
+- Integrate and maintain dynamic user interactions and business data.
+- Facilitate state-based reasoning and task automation for agents.
+- Query complex, evolving data with semantic, keyword, and graph-based search methods.
 
 <br />
 
@@ -31,31 +35,21 @@ a fusion of time, full-text, semantic, and graph algorithm approaches.
 
 <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
+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.
-- 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 and Zep Memory
 
-Graphiti powers the core of [Zep's memory layer](https://www.getzep.com) for LLM-powered Assistants and Agents.
+Graphiti powers the core of [Zep's memory layer](https://www.getzep.com) for AI Agents.
 
 Using Graphiti, we've demonstrated Zep is
 the [State of the Art in Agent Memory](https://blog.getzep.com/state-of-the-art-agent-memory/).
 
 Read our paper: [Zep: A Temporal Knowledge Graph Architecture for Agent Memory](https://arxiv.org/abs/2501.13956).
 
-We're excited to open-source Graphiti, believing its potential reaches far beyond memory applications.
+We're excited to open-source Graphiti, believing its potential reaches far beyond AI memory applications.
 
 <p align="center">
     <a href="https://arxiv.org/abs/2501.13956"><img src="images/arxiv-screenshot.png" alt="Zep: A Temporal Knowledge Graph Architecture for Agent Memory" width="700px"></a>
@@ -63,34 +57,41 @@ We're excited to open-source Graphiti, believing its potential reaches far beyon
 
 ## 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:
+Traditional RAG approaches often rely on batch processing and static data summarization, making them inefficient for frequently changing data. Graphiti addresses these challenges by providing:
 
-- **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.
+- **Real-Time Incremental Updates:** Immediate integration of new data episodes without batch recomputation.
+- **Bi-Temporal Data Model:** Explicit tracking of event occurrence and ingestion times, allowing accurate point-in-time queries.
+- **Efficient Hybrid Retrieval:** Combines semantic embeddings, keyword (BM25), and graph traversal to achieve low-latency queries without reliance on LLM summarization.
+- **Custom Entity Definitions:** Flexible ontology creation and support for developer-defined entities through straightforward Pydantic models.
+- **Scalability:** Efficiently manages large datasets with parallel processing, suitable for enterprise environments.
 
 <p align="center">
     <img src="/images/graphiti-intro-slides-stock-2.gif" alt="Graphiti structured + unstructured demo" width="700px">   
 </p>
 
+## Graphiti vs. GraphRAG
+
+| Aspect                     | GraphRAG                              | Graphiti                                         |
+| -------------------------- | ------------------------------------- | ------------------------------------------------ |
+| **Primary Use**            | Static document summarization         | Dynamic data management                          |
+| **Data Handling**          | Batch-oriented processing             | Continuous, incremental updates                  |
+| **Knowledge Structure**    | Entity clusters & community summaries | Episodic data, semantic entities, communities    |
+| **Retrieval Method**       | Sequential LLM summarization          | Hybrid semantic, keyword, and graph-based search |
+| **Adaptability**           | Low                                   | High                                             |
+| **Temporal Handling**      | Basic timestamp tracking              | Explicit bi-temporal tracking                    |
+| **Contradiction Handling** | LLM-driven summarization judgments    | Temporal edge invalidation                       |
+| **Query Latency**          | Seconds to tens of seconds            | Typically sub-second latency                     |
+| **Custom Entity Types**    | No                                    | Yes, customizable                                |
+| **Scalability**            | Moderate                              | High, optimized for large datasets               |
+
+Graphiti is specifically designed to address the challenges of dynamic and frequently updated datasets, making it particularly suitable for applications requiring real-time interaction and precise historical queries.
+
 ## Installation
 
 Requirements:
 
 - Python 3.10 or higher
-- Neo4j 5.26 or higher
+- Neo4j 5.26 or higher (serves as the embeddings storage backend)
 - OpenAI API key (for LLM inference and embedding)
 
 Optional:
@@ -123,7 +124,7 @@ from graphiti_core import Graphiti
 from graphiti_core.nodes import EpisodeType
 from datetime import datetime, timezone
 
-# Initialize Graphiti
+# Initialize Graphiti as Your Memory Layer
 graphiti = Graphiti("bolt://localhost:7687", "neo4j", "password")
 
 # Initialize the graph database with Graphiti's indices. This only needs to be done once.
@@ -144,7 +145,7 @@ for i, episode in enumerate(episodes):
         reference_time=datetime.now(timezone.utc)
     )
 
-# Search the graph
+# Search the graph for semantic memory retrieval
 # Execute a hybrid search combining semantic similarity and BM25 retrieval
 # Results are combined and reranked using Reciprocal Rank Fusion
 results = await graphiti.search('Who was the California Attorney General?')
@@ -176,11 +177,12 @@ results = await graphiti.search('Who was the California Attorney General?')
 # Results are weighted by their proximity, with distant edges receiving lower scores.
 await graphiti.search('Who was the California Attorney General?', center_node_uuid)
 
-# Close the connection
+# Close the connection when chat state management is complete
 graphiti.close()
 ```
 
 ## Graph Service
+
 The `server` directory contains an API service for interacting with the Graphiti API. It is built using FastAPI.
 
 Please see the [server README](./server/README.md) for more information.
@@ -196,6 +198,54 @@ to enable Neo4j's parallel runtime feature for several of our search queries.
 Note that this feature is not supported for Neo4j Community edition or for smaller AuraDB instances,
 as such this feature is off by default.
 
+## Using Graphiti with Azure OpenAI
+
+Graphiti supports Azure OpenAI for both LLM inference and embeddings. To use Azure OpenAI, you'll need to configure both the LLM client and embedder with your Azure OpenAI credentials.
+
+```python
+from openai import AsyncAzureOpenAI
+from graphiti_core import Graphiti
+from graphiti_core.llm_client import OpenAIClient
+from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
+from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
+
+# Azure OpenAI configuration
+api_key = "<your-api-key>"
+api_version = "<your-api-version>"
+azure_endpoint = "<your-azure-endpoint>"
+
+# Create Azure OpenAI client for LLM
+azure_openai_client = AsyncAzureOpenAI(
+    api_key=api_key,
+    api_version=api_version,
+    azure_endpoint=azure_endpoint
+)
+
+# Initialize Graphiti with Azure OpenAI clients
+graphiti = Graphiti(
+    "bolt://localhost:7687",
+    "neo4j",
+    "password",
+    llm_client=OpenAIClient(
+        client=azure_openai_client
+    ),
+    embedder=OpenAIEmbedder(
+        config=OpenAIEmbedderConfig(
+            embedding_model="text-embedding-3-small"  # Use your Azure deployed embedding model name
+        ),
+        client=azure_openai_client
+    ),
+    # Optional: Configure the OpenAI cross encoder with Azure OpenAI
+    cross_encoder=OpenAIRerankerClient(
+        client=azure_openai_client
+    )
+)
+
+# Now you can use Graphiti with Azure OpenAI
+```
+
+Make sure to replace the placeholder values with your actual Azure OpenAI credentials and specify the correct embedding model name that's deployed in your Azure OpenAI service.
+
 ## Documentation
 
 - [Guides and API documentation](https://help.getzep.com/graphiti).
@@ -206,10 +256,11 @@ as such this feature is off by default.
 
 Graphiti is under active development. We aim to maintain API stability while working on:
 
-- [ ] 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
+- [x] 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
 - [x] Enhancing retrieval capabilities with more robust and configurable options
+- [ ] Graphiti MCP Server
 - [ ] Expanding test coverage to ensure reliability and catch edge cases
 
 ## Contributing

graphiti_core-0.8.2/graphiti_core/cross_encoder/__init__.py

@@ -0,0 +1,20 @@
+"""
+Copyright 2025, 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 .client import CrossEncoderClient
+from .openai_reranker_client import OpenAIRerankerClient
+
+__all__ = ['CrossEncoderClient', 'OpenAIRerankerClient']

{graphiti_core-0.8.0 → graphiti_core-0.8.2}/graphiti_core/cross_encoder/openai_reranker_client.py

@@ -18,7 +18,7 @@ import logging
 from typing import Any
 
 import openai
-from openai import AsyncOpenAI
+from openai import AsyncAzureOpenAI, AsyncOpenAI
 from pydantic import BaseModel
 
 from ..helpers import semaphore_gather
@@ -36,21 +36,29 @@ class BooleanClassifier(BaseModel):
 
 
 class OpenAIRerankerClient(CrossEncoderClient):
-    def __init__(self, config: LLMConfig | None = None):
+    def __init__(
+        self,
+        config: LLMConfig | None = None,
+        client: AsyncOpenAI | AsyncAzureOpenAI | None = None,
+    ):
         """
-        Initialize the OpenAIClient with the provided configuration, cache setting, and client.
+        Initialize the OpenAIRerankerClient with the provided configuration and client.
+
+        This reranker uses the OpenAI API to run a simple boolean classifier prompt concurrently
+        for each passage. Log-probabilities are used to rank the passages.
 
         Args:
             config (LLMConfig | None): The configuration for the LLM client, including API key, model, base URL, temperature, and max tokens.
-            cache (bool): Whether to use caching for responses. Defaults to False.
-            client (Any | None): An optional async client instance to use. If not provided, a new AsyncOpenAI client is created.
-
+            client (AsyncOpenAI | AsyncAzureOpenAI | None): An optional async client instance to use. If not provided, a new AsyncOpenAI client is created.
         """
         if config is None:
             config = LLMConfig()
 
         self.config = config
-        self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
+        if client is None:
+            self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
+        else:
+            self.client = client
 
     async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]:
         openai_messages_list: Any = [
@@ -62,7 +70,7 @@ class OpenAIRerankerClient(CrossEncoderClient):
                 Message(
                     role='user',
                     content=f"""
-                           Respond with "True" if PASSAGE is relevant to QUERY and "False" otherwise. 
+                           Respond with "True" if PASSAGE is relevant to QUERY and "False" otherwise.
                            <PASSAGE>
                            {passage}
                            </PASSAGE>

{graphiti_core-0.8.0 → graphiti_core-0.8.2}/graphiti_core/embedder/openai.py

@@ -15,8 +15,9 @@ limitations under the License.
 """
 
 from collections.abc import Iterable
+from typing import Union
 
-from openai import AsyncOpenAI
+from openai import AsyncAzureOpenAI, AsyncOpenAI
 from openai.types import EmbeddingModel
 
 from .client import EmbedderClient, EmbedderConfig
@@ -33,13 +34,23 @@ class OpenAIEmbedderConfig(EmbedderConfig):
 class OpenAIEmbedder(EmbedderClient):
     """
     OpenAI Embedder Client
+
+    This client supports both AsyncOpenAI and AsyncAzureOpenAI clients.
     """
 
-    def __init__(self, config: OpenAIEmbedderConfig | None = None):
+    def __init__(
+        self,
+        config: OpenAIEmbedderConfig | None = None,
+        client: Union[AsyncOpenAI, AsyncAzureOpenAI, None] = None,
+    ):
         if config is None:
             config = OpenAIEmbedderConfig()
         self.config = config
-        self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
+
+        if client is not None:
+            self.client = client
+        else:
+            self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
 
     async def create(
         self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]

{graphiti_core-0.8.0 → graphiti_core-0.8.2}/graphiti_core/models/nodes/node_db_queries.py

@@ -30,10 +30,10 @@ EPISODIC_NODE_SAVE_BULK = """
 """
 
 ENTITY_NODE_SAVE = """
-        MERGE (n:Entity {uuid: $uuid})
+        MERGE (n:Entity {uuid: $entity_data.uuid})
         SET n:$($labels)
         SET n = $entity_data
-        WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
+        WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $entity_data.name_embedding)
         RETURN n.uuid AS uuid"""
 
 ENTITY_NODE_SAVE_BULK = """

{graphiti_core-0.8.0 → graphiti_core-0.8.2}/graphiti_core/utils/maintenance/node_operations.py

@@ -189,6 +189,9 @@ async def extract_nodes(
     new_nodes = []
     for name in extracted_node_names:
         entity_type = node_classifications.get(name)
+        if entity_types is not None and entity_type not in entity_types:
+            entity_type = None
+
         labels = (
             ['Entity']
             if entity_type is None or entity_type == 'None' or entity_type == 'null'

{graphiti_core-0.8.0 → graphiti_core-0.8.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphiti-core"
-version = "0.8.0"
+version = "0.8.2"
 description = "A temporal graph building library"
 authors = [
     "Paul Paliychuk <paul@getzep.com>",