graphiti-core 0.8.8__tar.gz → 0.9.1__tar.gz

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.8.8
+Version: 0.9.1
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -11,7 +11,9 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: anthropic (>=0.49.0,<0.50.0)
+Provides-Extra: anthropic
+Provides-Extra: google-genai
+Provides-Extra: groq
 Requires-Dist: diskcache (>=5.6.3,<6.0.0)
 Requires-Dist: neo4j (>=5.23.0,<6.0.0)
 Requires-Dist: numpy (>=1.0.0)
@@ -58,8 +60,8 @@ Use Graphiti to:
 
 <br />
 
-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
+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.
 
@@ -119,7 +121,7 @@ Requirements:
 
 Optional:
 
-- Anthropic or Groq API key (for alternative LLM providers)
+- Google Gemini, 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
@@ -135,6 +137,22 @@ or
 poetry add graphiti-core
 ```
 
+You can also install optional LLM providers as extras:
+
+```bash
+# Install with Anthropic support
+pip install graphiti-core[anthropic]
+
+# Install with Groq support
+pip install graphiti-core[groq]
+
+# Install with Google Gemini support
+pip install graphiti-core[google-genai]
+
+# Install with multiple providers
+pip install graphiti-core[anthropic,groq,google-genai]
+```
+
 ## Quick Start
 
 > [!IMPORTANT]
@@ -142,67 +160,16 @@ poetry add graphiti-core
 > Support for Anthropic and Groq LLM inferences is available, too. Other LLM providers may be supported via OpenAI
 > compatible APIs.
 
-```python
-from graphiti_core import Graphiti
-from graphiti_core.nodes import EpisodeType
-from datetime import datetime, timezone
-
-# 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.
-graphiti.build_indices_and_constraints()
-
-# Add episodes
-episodes = [
-    "Kamala Harris is the Attorney General of California. She was previously "
-    "the district attorney for San Francisco.",
-    "As AG, Harris was in office from January 3, 2011 – January 3, 2017",
-]
-for i, episode in enumerate(episodes):
-    await graphiti.add_episode(
-        name=f"Freakonomics Radio {i}",
-        episode_body=episode,
-        source=EpisodeType.text,
-        source_description="podcast",
-        reference_time=datetime.now(timezone.utc)
-    )
+For a complete working example, see the [Quickstart Example](./examples/quickstart/README.md) in the examples directory. The quickstart demonstrates:
 
-# 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?')
-[
-    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 >)
-)
-]
+1. Connecting to a Neo4j database
+2. Initializing Graphiti indices and constraints
+3. Adding episodes to the graph (both text and structured JSON)
+4. Searching for relationships (edges) using hybrid search
+5. Reranking search results using graph distance
+6. Searching for nodes using predefined search recipes
 
-# Rerank search results based on graph distance
-# Provide a node UUID to prioritize results closer to that node in the graph.
-# 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 when chat state management is complete
-graphiti.close()
-```
+The example is fully documented with clear explanations of each functionality and includes a comprehensive README with setup instructions and next steps.
 
 ## Graph Service
 
@@ -285,6 +252,42 @@ graphiti = Graphiti(
 
 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.
 
+## Using Graphiti with Google Gemini
+
+Graphiti supports Google's Gemini models for both LLM inference and embeddings. To use Gemini, you'll need to configure both the LLM client and embedder with your Google API key.
+
+```python
+from graphiti_core import Graphiti
+from graphiti_core.llm_client.gemini_client import GeminiClient, LLMConfig
+from graphiti_core.embedder.gemini import GeminiEmbedder, GeminiEmbedderConfig
+
+# Google API key configuration
+api_key = "<your-google-api-key>"
+
+# Initialize Graphiti with Gemini clients
+graphiti = Graphiti(
+    "bolt://localhost:7687",
+    "neo4j",
+    "password",
+    llm_client=GeminiClient(
+        config=LLMConfig(
+            api_key=api_key,
+            model="gemini-2.0-flash"
+        )
+    ),
+    embedder=GeminiEmbedder(
+        config=GeminiEmbedderConfig(
+            api_key=api_key,
+            embedding_model="embedding-001"
+        )
+    )
+)
+
+# Now you can use Graphiti with Google Gemini
+```
+
+Make sure to replace the placeholder value with your actual Google API key. You can find more details in the example file at `examples/gemini_example.py`.
+
 ## Documentation
 
 - [Guides and API documentation](https://help.getzep.com/graphiti).

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/README.md

@@ -35,8 +35,8 @@ Use Graphiti to:
 
 <br />
 
-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
+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.
 
@@ -96,7 +96,7 @@ Requirements:
 
 Optional:
 
-- Anthropic or Groq API key (for alternative LLM providers)
+- Google Gemini, 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
@@ -112,6 +112,22 @@ or
 poetry add graphiti-core
 ```
 
+You can also install optional LLM providers as extras:
+
+```bash
+# Install with Anthropic support
+pip install graphiti-core[anthropic]
+
+# Install with Groq support
+pip install graphiti-core[groq]
+
+# Install with Google Gemini support
+pip install graphiti-core[google-genai]
+
+# Install with multiple providers
+pip install graphiti-core[anthropic,groq,google-genai]
+```
+
 ## Quick Start
 
 > [!IMPORTANT]
@@ -119,67 +135,16 @@ poetry add graphiti-core
 > Support for Anthropic and Groq LLM inferences is available, too. Other LLM providers may be supported via OpenAI
 > compatible APIs.
 
-```python
-from graphiti_core import Graphiti
-from graphiti_core.nodes import EpisodeType
-from datetime import datetime, timezone
-
-# 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.
-graphiti.build_indices_and_constraints()
-
-# Add episodes
-episodes = [
-    "Kamala Harris is the Attorney General of California. She was previously "
-    "the district attorney for San Francisco.",
-    "As AG, Harris was in office from January 3, 2011 – January 3, 2017",
-]
-for i, episode in enumerate(episodes):
-    await graphiti.add_episode(
-        name=f"Freakonomics Radio {i}",
-        episode_body=episode,
-        source=EpisodeType.text,
-        source_description="podcast",
-        reference_time=datetime.now(timezone.utc)
-    )
+For a complete working example, see the [Quickstart Example](./examples/quickstart/README.md) in the examples directory. The quickstart demonstrates:
 
-# 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?')
-[
-    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 >)
-)
-]
+1. Connecting to a Neo4j database
+2. Initializing Graphiti indices and constraints
+3. Adding episodes to the graph (both text and structured JSON)
+4. Searching for relationships (edges) using hybrid search
+5. Reranking search results using graph distance
+6. Searching for nodes using predefined search recipes
 
-# Rerank search results based on graph distance
-# Provide a node UUID to prioritize results closer to that node in the graph.
-# 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 when chat state management is complete
-graphiti.close()
-```
+The example is fully documented with clear explanations of each functionality and includes a comprehensive README with setup instructions and next steps.
 
 ## Graph Service
 
@@ -262,6 +227,42 @@ graphiti = Graphiti(
 
 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.
 
+## Using Graphiti with Google Gemini
+
+Graphiti supports Google's Gemini models for both LLM inference and embeddings. To use Gemini, you'll need to configure both the LLM client and embedder with your Google API key.
+
+```python
+from graphiti_core import Graphiti
+from graphiti_core.llm_client.gemini_client import GeminiClient, LLMConfig
+from graphiti_core.embedder.gemini import GeminiEmbedder, GeminiEmbedderConfig
+
+# Google API key configuration
+api_key = "<your-google-api-key>"
+
+# Initialize Graphiti with Gemini clients
+graphiti = Graphiti(
+    "bolt://localhost:7687",
+    "neo4j",
+    "password",
+    llm_client=GeminiClient(
+        config=LLMConfig(
+            api_key=api_key,
+            model="gemini-2.0-flash"
+        )
+    ),
+    embedder=GeminiEmbedder(
+        config=GeminiEmbedderConfig(
+            api_key=api_key,
+            embedding_model="embedding-001"
+        )
+    )
+)
+
+# Now you can use Graphiti with Google Gemini
+```
+
+Make sure to replace the placeholder value with your actual Google API key. You can find more details in the example file at `examples/gemini_example.py`.
+
 ## Documentation
 
 - [Guides and API documentation](https://help.getzep.com/graphiti).

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/graphiti_core/edges.py

@@ -143,9 +143,9 @@ class EpisodicEdge(Edge):
         driver: AsyncDriver,
         group_ids: list[str],
         limit: int | None = None,
-        created_at: datetime | None = None,
+        uuid_cursor: str | None = None,
     ):
-        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+        cursor_query: LiteralString = 'AND e.uuid < $uuid' if uuid_cursor else ''
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
 
         records, _, _ = await driver.execute_query(
@@ -165,7 +165,7 @@ class EpisodicEdge(Edge):
         """
             + limit_query,
             group_ids=group_ids,
-            created_at=created_at,
+            uuid=uuid_cursor,
             limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
@@ -297,9 +297,9 @@ class EntityEdge(Edge):
         driver: AsyncDriver,
         group_ids: list[str],
         limit: int | None = None,
-        created_at: datetime | None = None,
+        uuid_cursor: str | None = None,
     ):
-        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+        cursor_query: LiteralString = 'AND e.uuid < $uuid' if uuid_cursor else ''
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
 
         records, _, _ = await driver.execute_query(
@@ -326,7 +326,7 @@ class EntityEdge(Edge):
         """
             + limit_query,
             group_ids=group_ids,
-            created_at=created_at,
+            uuid=uuid_cursor,
             limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
@@ -430,9 +430,9 @@ class CommunityEdge(Edge):
         driver: AsyncDriver,
         group_ids: list[str],
         limit: int | None = None,
-        created_at: datetime | None = None,
+        uuid_cursor: str | None = None,
     ):
-        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+        cursor_query: LiteralString = 'AND e.uuid < $uuid' if uuid_cursor else ''
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
 
         records, _, _ = await driver.execute_query(
@@ -452,7 +452,7 @@ class CommunityEdge(Edge):
         """
             + limit_query,
             group_ids=group_ids,
-            created_at=created_at,
+            uuid=uuid_cursor,
             limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/graphiti_core/embedder/__init__.py

@@ -1,4 +1,8 @@
 from .client import EmbedderClient
 from .openai import OpenAIEmbedder, OpenAIEmbedderConfig
 
-__all__ = ['EmbedderClient', 'OpenAIEmbedder', 'OpenAIEmbedderConfig']
+__all__ = [
+    'EmbedderClient',
+    'OpenAIEmbedder',
+    'OpenAIEmbedderConfig',
+]

graphiti_core-0.9.1/graphiti_core/embedder/gemini.py

@@ -0,0 +1,68 @@
+"""
+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 collections.abc import Iterable
+
+from google import genai  # type: ignore
+from google.genai import types  # type: ignore
+from pydantic import Field
+
+from .client import EmbedderClient, EmbedderConfig
+
+DEFAULT_EMBEDDING_MODEL = 'embedding-001'
+
+
+class GeminiEmbedderConfig(EmbedderConfig):
+    embedding_model: str = Field(default=DEFAULT_EMBEDDING_MODEL)
+    api_key: str | None = None
+
+
+class GeminiEmbedder(EmbedderClient):
+    """
+    Google Gemini Embedder Client
+    """
+
+    def __init__(self, config: GeminiEmbedderConfig | None = None):
+        if config is None:
+            config = GeminiEmbedderConfig()
+        self.config = config
+
+        # Configure the Gemini API
+        self.client = genai.Client(
+            api_key=config.api_key,
+        )
+
+    async def create(
+        self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
+    ) -> list[float]:
+        """
+        Create embeddings for the given input data using Google's Gemini embedding model.
+
+        Args:
+            input_data: The input data to create embeddings for. Can be a string, list of strings,
+                       or an iterable of integers or iterables of integers.
+
+        Returns:
+            A list of floats representing the embedding vector.
+        """
+        # Generate embeddings
+        result = await self.client.aio.models.embed_content(
+            model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
+            contents=[input_data],
+            config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
+        )
+
+        return result.embeddings[0].values

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/graphiti_core/graphiti.py

@@ -71,10 +71,7 @@ from graphiti_core.utils.maintenance.graph_data_operations import (
     build_indices_and_constraints,
     retrieve_episodes,
 )
-from graphiti_core.utils.maintenance.node_operations import (
-    extract_nodes,
-    resolve_extracted_nodes,
-)
+from graphiti_core.utils.maintenance.node_operations import extract_nodes, resolve_extracted_nodes
 from graphiti_core.utils.maintenance.temporal_operations import get_edge_contradictions
 from graphiti_core.utils.ontology_utils.entity_types_utils import validate_entity_types
 

graphiti_core-0.9.1/graphiti_core/llm_client/gemini_client.py

@@ -0,0 +1,186 @@
+"""
+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 json
+import logging
+import typing
+
+from google import genai  # type: ignore
+from google.genai import types  # type: ignore
+from pydantic import BaseModel
+
+from ..prompts.models import Message
+from .client import LLMClient
+from .config import DEFAULT_MAX_TOKENS, LLMConfig
+from .errors import RateLimitError
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = 'gemini-2.0-flash'
+
+
+class GeminiClient(LLMClient):
+    """
+    GeminiClient is a client class for interacting with Google's Gemini language models.
+
+    This class extends the LLMClient and provides methods to initialize the client
+    and generate responses from the Gemini language model.
+
+    Attributes:
+        model (str): The model name to use for generating responses.
+        temperature (float): The temperature to use for generating responses.
+        max_tokens (int): The maximum number of tokens to generate in a response.
+
+    Methods:
+        __init__(config: LLMConfig | None = None, cache: bool = False):
+            Initializes the GeminiClient with the provided configuration and cache setting.
+
+        _generate_response(messages: list[Message]) -> dict[str, typing.Any]:
+            Generates a response from the language model based on the provided messages.
+    """
+
+    def __init__(
+        self,
+        config: LLMConfig | None = None,
+        cache: bool = False,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ):
+        """
+        Initialize the GeminiClient with the provided configuration and cache setting.
+
+        Args:
+            config (LLMConfig | None): The configuration for the LLM client, including API key, model, temperature, and max tokens.
+            cache (bool): Whether to use caching for responses. Defaults to False.
+        """
+        if config is None:
+            config = LLMConfig()
+
+        super().__init__(config, cache)
+
+        self.model = config.model
+        # Configure the Gemini API
+        self.client = genai.Client(
+            api_key=config.api_key,
+        )
+        self.max_tokens = max_tokens
+
+    async def _generate_response(
+        self,
+        messages: list[Message],
+        response_model: type[BaseModel] | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> dict[str, typing.Any]:
+        """
+        Generate a response from the Gemini language model.
+
+        Args:
+            messages (list[Message]): A list of messages to send to the language model.
+            response_model (type[BaseModel] | None): An optional Pydantic model to parse the response into.
+            max_tokens (int): The maximum number of tokens to generate in the response.
+
+        Returns:
+            dict[str, typing.Any]: The response from the language model.
+
+        Raises:
+            RateLimitError: If the API rate limit is exceeded.
+            RefusalError: If the content is blocked by the model.
+            Exception: If there is an error generating the response.
+        """
+        try:
+            gemini_messages: list[types.Content] = []
+            # If a response model is provided, add schema for structured output
+            system_prompt = ''
+            if response_model is not None:
+                # Get the schema from the Pydantic model
+                pydantic_schema = response_model.model_json_schema()
+
+                # Create instruction to output in the desired JSON format
+                system_prompt += (
+                    f'Output ONLY valid JSON matching this schema: {json.dumps(pydantic_schema)}.\n'
+                    'Do not include any explanatory text before or after the JSON.\n\n'
+                )
+
+            # Add messages content
+            # First check for a system message
+            if messages and messages[0].role == 'system':
+                system_prompt = f'{messages[0].content}\n\n {system_prompt}'
+                messages = messages[1:]
+
+            # Add the rest of the messages
+            for m in messages:
+                m.content = self._clean_input(m.content)
+                gemini_messages.append(
+                    types.Content(role=m.role, parts=[types.Part.from_text(text=m.content)])
+                )
+
+            # Create generation config
+            generation_config = types.GenerateContentConfig(
+                temperature=self.temperature,
+                max_output_tokens=max_tokens or self.max_tokens,
+                response_mime_type='application/json' if response_model else None,
+                response_schema=response_model if response_model else None,
+                system_instruction=system_prompt,
+            )
+
+            # Generate content using the simple string approach
+            response = await self.client.aio.models.generate_content(
+                model=self.model or DEFAULT_MODEL,
+                contents=gemini_messages,
+                config=generation_config,
+            )
+
+            # If this was a structured output request, parse the response into the Pydantic model
+            if response_model is not None:
+                try:
+                    validated_model = response_model.model_validate(json.loads(response.text))
+
+                    # Return as a dictionary for API consistency
+                    return validated_model.model_dump()
+                except Exception as e:
+                    raise Exception(f'Failed to parse structured response: {e}') from e
+
+            # Otherwise, return the response text as a dictionary
+            return {'content': response.text}
+
+        except Exception as e:
+            # Check if it's a rate limit error
+            if 'rate limit' in str(e).lower() or 'quota' in str(e).lower():
+                raise RateLimitError from e
+            logger.error(f'Error in generating LLM response: {e}')
+            raise
+
+    async def generate_response(
+        self,
+        messages: list[Message],
+        response_model: type[BaseModel] | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> dict[str, typing.Any]:
+        """
+        Generate a response from the Gemini language model.
+        This method overrides the parent class method to provide a direct implementation.
+
+        Args:
+            messages (list[Message]): A list of messages to send to the language model.
+            response_model (type[BaseModel] | None): An optional Pydantic model to parse the response into.
+            max_tokens (int): The maximum number of tokens to generate in the response.
+
+        Returns:
+            dict[str, typing.Any]: The response from the language model.
+        """
+        # Call the internal _generate_response method
+        return await self._generate_response(
+            messages=messages, response_model=response_model, max_tokens=max_tokens
+        )

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/graphiti_core/llm_client/openai_client.py

@@ -87,6 +87,8 @@ class OpenAIClient(LLMClient):
         else:
             self.client = client
 
+        self.max_tokens = max_tokens
+
     async def _generate_response(
         self,
         messages: list[Message],

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/graphiti_core/nodes.py

@@ -216,9 +216,9 @@ class EpisodicNode(Node):
         driver: AsyncDriver,
         group_ids: list[str],
         limit: int | None = None,
-        created_at: datetime | None = None,
+        uuid_cursor: str | None = None,
     ):
-        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+        cursor_query: LiteralString = 'AND e.uuid < $uuid' if uuid_cursor else ''
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
 
         records, _, _ = await driver.execute_query(
@@ -241,7 +241,7 @@ class EpisodicNode(Node):
         """
             + limit_query,
             group_ids=group_ids,
-            created_at=created_at,
+            uuid=uuid_cursor,
             limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
@@ -348,9 +348,9 @@ class EntityNode(Node):
         driver: AsyncDriver,
         group_ids: list[str],
         limit: int | None = None,
-        created_at: datetime | None = None,
+        uuid_cursor: str | None = None,
     ):
-        cursor_query: LiteralString = 'AND n.created_at < $created_at' if created_at else ''
+        cursor_query: LiteralString = 'AND n.uuid < $uuid' if uuid_cursor else ''
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
 
         records, _, _ = await driver.execute_query(
@@ -372,7 +372,7 @@ class EntityNode(Node):
         """
             + limit_query,
             group_ids=group_ids,
-            created_at=created_at,
+            uuid=uuid_cursor,
             limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
@@ -465,9 +465,9 @@ class CommunityNode(Node):
         driver: AsyncDriver,
         group_ids: list[str],
         limit: int | None = None,
-        created_at: datetime | None = None,
+        uuid_cursor: str | None = None,
     ):
-        cursor_query: LiteralString = 'AND n.created_at < $created_at' if created_at else ''
+        cursor_query: LiteralString = 'AND n.uuid < $uuid' if uuid_cursor else ''
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
 
         records, _, _ = await driver.execute_query(
@@ -487,7 +487,7 @@ class CommunityNode(Node):
         """
             + limit_query,
             group_ids=group_ids,
-            created_at=created_at,
+            uuid=uuid_cursor,
             limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/graphiti_core/prompts/lib.py

@@ -16,63 +16,27 @@ 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 .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 .eval import Prompt as EvalPrompt
 from .eval import Versions as EvalVersions
 from .eval import versions as eval_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 .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
 from .prompt_helpers import DO_NOT_ESCAPE_UNICODE
 from .summarize_nodes import Prompt as SummarizeNodesPrompt

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/graphiti_core/utils/maintenance/__init__.py

@@ -1,8 +1,5 @@
 from .edge_operations import build_episodic_edges, extract_edges
-from .graph_data_operations import (
-    clear_data,
-    retrieve_episodes,
-)
+from .graph_data_operations import clear_data, retrieve_episodes
 from .node_operations import extract_nodes
 
 __all__ = [

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/graphiti_core/utils/maintenance/community_operations.py

@@ -9,11 +9,7 @@ from graphiti_core.edges import CommunityEdge
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
 from graphiti_core.llm_client import LLMClient
-from graphiti_core.nodes import (
-    CommunityNode,
-    EntityNode,
-    get_community_node_from_record,
-)
+from graphiti_core.nodes import CommunityNode, EntityNode, get_community_node_from_record
 from graphiti_core.prompts import prompt_library
 from graphiti_core.prompts.summarize_nodes import Summary, SummaryDescription
 from graphiti_core.utils.datetime_utils import utc_now

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/graphiti_core/utils/maintenance/node_operations.py

@@ -26,11 +26,7 @@ from graphiti_core.llm_client import LLMClient
 from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.prompts import prompt_library
 from graphiti_core.prompts.dedupe_nodes import NodeDuplicate
-from graphiti_core.prompts.extract_nodes import (
-    EntityClassification,
-    ExtractedNodes,
-    MissedEntities,
-)
+from graphiti_core.prompts.extract_nodes import EntityClassification, ExtractedNodes, MissedEntities
 from graphiti_core.prompts.summarize_nodes import Summary
 from graphiti_core.utils.datetime_utils import utc_now
 

{graphiti_core-0.8.8 → graphiti_core-0.9.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphiti-core"
-version = "0.8.8"
+version = "0.9.1"
 description = "A temporal graph building library"
 authors = [
     "Paul Paliychuk <paul@getzep.com>",
@@ -21,7 +21,11 @@ openai = "^1.53.0"
 tenacity = "9.0.0"
 numpy = ">=1.0.0"
 python-dotenv = "^1.0.1"
-anthropic = "~0.49.0"
+
+[tool.poetry.extras]
+anthropic = ["anthropic"]
+groq = ["groq"]
+google-genai = ["google-genai"]
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.11.1"