graphiti-core 0.13.0__tar.gz → 0.13.1__tar.gz

{graphiti_core-0.13.0 → graphiti_core-0.13.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: graphiti-core
-Version: 0.13.0
+Version: 0.13.1
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -23,7 +23,7 @@ Requires-Dist: google-genai (>=1.8.0) ; extra == "google-genai"
 Requires-Dist: groq (>=0.2.0) ; extra == "groq"
 Requires-Dist: neo4j (>=5.26.0)
 Requires-Dist: numpy (>=1.0.0)
-Requires-Dist: openai (>=1.53.0)
+Requires-Dist: openai (>=1.91.0)
 Requires-Dist: pydantic (>=2.11.5)
 Requires-Dist: python-dotenv (>=1.0.1)
 Requires-Dist: tenacity (>=9.0.0)
@@ -331,6 +331,55 @@ graphiti = Graphiti(
 # Now you can use Graphiti with Google Gemini
 ```
 
+## Using Graphiti with Ollama (Local LLM)
+
+Graphiti supports Ollama for running local LLMs and embedding models via Ollama's OpenAI-compatible API. This is ideal for privacy-focused applications or when you want to avoid API costs.
+
+
+Install the models:
+ollama pull deepseek-r1:7b         # LLM
+ollama pull nomic-embed-text       # embeddings
+
+```python
+from graphiti_core import Graphiti
+from graphiti_core.llm_client.config import LLMConfig
+from graphiti_core.llm_client.openai_client import OpenAIClient
+from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
+from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
+
+# Configure Ollama LLM client
+llm_config = LLMConfig(
+    api_key="abc",  # Ollama doesn't require a real API key
+    model="deepseek-r1:7b",
+    small_model="deepseek-r1:7b",
+    base_url="http://localhost:11434/v1", # Ollama provides this port
+)
+
+llm_client = OpenAIClient(config=llm_config)
+
+# Initialize Graphiti with Ollama clients
+graphiti = Graphiti(
+    "bolt://localhost:7687",
+    "neo4j",
+    "password",
+    llm_client=llm_client,
+    embedder=OpenAIEmbedder(
+        config=OpenAIEmbedderConfig(
+            api_key="abc",
+            embedding_model="nomic-embed-text",
+            embedding_dim=768,
+            base_url="http://localhost:11434/v1",
+        )
+    ),
+    cross_encoder=OpenAIRerankerClient(client=llm_client, config=llm_config),
+)
+
+# Now you can use Graphiti with local Ollama models
+```
+
+Ensure Ollama is running (`ollama serve`) and that you have pulled the models you want to use.
+
+
 ## Documentation
 
 - [Guides and API documentation](https://help.getzep.com/graphiti).

{graphiti_core-0.13.0 → graphiti_core-0.13.1}/README.md

@@ -298,6 +298,55 @@ graphiti = Graphiti(
 # Now you can use Graphiti with Google Gemini
 ```
 
+## Using Graphiti with Ollama (Local LLM)
+
+Graphiti supports Ollama for running local LLMs and embedding models via Ollama's OpenAI-compatible API. This is ideal for privacy-focused applications or when you want to avoid API costs.
+
+
+Install the models:
+ollama pull deepseek-r1:7b         # LLM
+ollama pull nomic-embed-text       # embeddings
+
+```python
+from graphiti_core import Graphiti
+from graphiti_core.llm_client.config import LLMConfig
+from graphiti_core.llm_client.openai_client import OpenAIClient
+from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
+from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
+
+# Configure Ollama LLM client
+llm_config = LLMConfig(
+    api_key="abc",  # Ollama doesn't require a real API key
+    model="deepseek-r1:7b",
+    small_model="deepseek-r1:7b",
+    base_url="http://localhost:11434/v1", # Ollama provides this port
+)
+
+llm_client = OpenAIClient(config=llm_config)
+
+# Initialize Graphiti with Ollama clients
+graphiti = Graphiti(
+    "bolt://localhost:7687",
+    "neo4j",
+    "password",
+    llm_client=llm_client,
+    embedder=OpenAIEmbedder(
+        config=OpenAIEmbedderConfig(
+            api_key="abc",
+            embedding_model="nomic-embed-text",
+            embedding_dim=768,
+            base_url="http://localhost:11434/v1",
+        )
+    ),
+    cross_encoder=OpenAIRerankerClient(client=llm_client, config=llm_config),
+)
+
+# Now you can use Graphiti with local Ollama models
+```
+
+Ensure Ollama is running (`ollama serve`) and that you have pulled the models you want to use.
+
+
 ## Documentation
 
 - [Guides and API documentation](https://help.getzep.com/graphiti).

graphiti_core-0.13.1/graphiti_core/llm_client/azure_openai_client.py

@@ -0,0 +1,77 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import logging
+from typing import ClassVar
+
+from openai import AsyncAzureOpenAI
+from openai.types.chat import ChatCompletionMessageParam
+from pydantic import BaseModel
+
+from .config import DEFAULT_MAX_TOKENS, LLMConfig
+from .openai_base_client import BaseOpenAIClient
+
+logger = logging.getLogger(__name__)
+
+
+class AzureOpenAILLMClient(BaseOpenAIClient):
+    """Wrapper class for AsyncAzureOpenAI that implements the LLMClient interface."""
+
+    # Class-level constants
+    MAX_RETRIES: ClassVar[int] = 2
+
+    def __init__(
+        self,
+        azure_client: AsyncAzureOpenAI,
+        config: LLMConfig | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ):
+        super().__init__(config, cache=False, max_tokens=max_tokens)
+        self.client = azure_client
+
+    async def _create_structured_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel],
+    ):
+        """Create a structured completion using Azure OpenAI's beta parse API."""
+        return await self.client.beta.chat.completions.parse(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            response_format=response_model,  # type: ignore
+        )
+
+    async def _create_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel] | None = None,
+    ):
+        """Create a regular completion with JSON format using Azure OpenAI."""
+        return await self.client.chat.completions.create(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            response_format={'type': 'json_object'},
+        )

graphiti_core-0.13.0/graphiti_core/llm_client/openai_client.py → graphiti_core-0.13.1/graphiti_core/llm_client/openai_base_client.py

@@ -14,12 +14,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import json
 import logging
 import typing
-from typing import ClassVar
+from abc import abstractmethod
+from typing import Any, ClassVar
 
 import openai
-from openai import AsyncOpenAI
 from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
@@ -34,25 +35,12 @@ DEFAULT_MODEL = 'gpt-4.1-mini'
 DEFAULT_SMALL_MODEL = 'gpt-4.1-nano'
 
 
-class OpenAIClient(LLMClient):
+class BaseOpenAIClient(LLMClient):
     """
-    OpenAIClient is a client class for interacting with OpenAI's language models.
+    Base client class for OpenAI-compatible APIs (OpenAI and Azure OpenAI).
 
-    This class extends the LLMClient and provides methods to initialize the client,
-    get an embedder, and generate responses from the language model.
-
-    Attributes:
-        client (AsyncOpenAI): The OpenAI client used to interact with the API.
-        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, client: typing.Any = None):
-            Initializes the OpenAIClient with the provided configuration, cache setting, and client.
-
-        _generate_response(messages: list[Message]) -> dict[str, typing.Any]:
-            Generates a response from the language model based on the provided messages.
+    This class contains shared logic for both OpenAI and Azure OpenAI clients,
+    reducing code duplication while allowing for implementation-specific differences.
     """
 
     # Class-level constants
@@ -62,41 +50,45 @@ class OpenAIClient(LLMClient):
         self,
         config: LLMConfig | None = None,
         cache: bool = False,
-        client: typing.Any = None,
         max_tokens: int = DEFAULT_MAX_TOKENS,
     ):
-        """
-        Initialize the OpenAIClient with the provided configuration, cache setting, and client.
-
-        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.
-
-        """
-        # removed caching to simplify the `generate_response` override
         if cache:
-            raise NotImplementedError('Caching is not implemented for OpenAI')
+            raise NotImplementedError('Caching is not implemented for OpenAI-based clients')
 
         if config is None:
             config = LLMConfig()
 
         super().__init__(config, cache)
-
-        if client is None:
-            self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
-        else:
-            self.client = client
-
         self.max_tokens = max_tokens
 
-    async def _generate_response(
+    @abstractmethod
+    async def _create_completion(
         self,
-        messages: list[Message],
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
         response_model: type[BaseModel] | None = None,
-        max_tokens: int = DEFAULT_MAX_TOKENS,
-        model_size: ModelSize = ModelSize.medium,
-    ) -> dict[str, typing.Any]:
+    ) -> Any:
+        """Create a completion using the specific client implementation."""
+        pass
+
+    @abstractmethod
+    async def _create_structured_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel],
+    ) -> Any:
+        """Create a structured completion using the specific client implementation."""
+        pass
+
+    def _convert_messages_to_openai_format(
+        self, messages: list[Message]
+    ) -> list[ChatCompletionMessageParam]:
+        """Convert internal Message format to OpenAI ChatCompletionMessageParam format."""
         openai_messages: list[ChatCompletionMessageParam] = []
         for m in messages:
             m.content = self._clean_input(m.content)
@@ -104,28 +96,61 @@ class OpenAIClient(LLMClient):
                 openai_messages.append({'role': 'user', 'content': m.content})
             elif m.role == 'system':
                 openai_messages.append({'role': 'system', 'content': m.content})
+        return openai_messages
+
+    def _get_model_for_size(self, model_size: ModelSize) -> str:
+        """Get the appropriate model name based on the requested size."""
+        if model_size == ModelSize.small:
+            return self.small_model or DEFAULT_SMALL_MODEL
+        else:
+            return self.model or DEFAULT_MODEL
+
+    def _handle_structured_response(self, response: Any) -> dict[str, Any]:
+        """Handle structured response parsing and validation."""
+        response_object = response.choices[0].message
+
+        if response_object.parsed:
+            return response_object.parsed.model_dump()
+        elif response_object.refusal:
+            raise RefusalError(response_object.refusal)
+        else:
+            raise Exception(f'Invalid response from LLM: {response_object.model_dump()}')
+
+    def _handle_json_response(self, response: Any) -> dict[str, Any]:
+        """Handle JSON response parsing."""
+        result = response.choices[0].message.content or '{}'
+        return json.loads(result)
+
+    async def _generate_response(
+        self,
+        messages: list[Message],
+        response_model: type[BaseModel] | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+        model_size: ModelSize = ModelSize.medium,
+    ) -> dict[str, Any]:
+        """Generate a response using the appropriate client implementation."""
+        openai_messages = self._convert_messages_to_openai_format(messages)
+        model = self._get_model_for_size(model_size)
+
         try:
-            if model_size == ModelSize.small:
-                model = self.small_model or DEFAULT_SMALL_MODEL
-            else:
-                model = self.model or DEFAULT_MODEL
-
-            response = await self.client.beta.chat.completions.parse(
-                model=model,
-                messages=openai_messages,
-                temperature=self.temperature,
-                max_tokens=max_tokens or self.max_tokens,
-                response_format=response_model,  # type: ignore
-            )
-
-            response_object = response.choices[0].message
-
-            if response_object.parsed:
-                return response_object.parsed.model_dump()
-            elif response_object.refusal:
-                raise RefusalError(response_object.refusal)
+            if response_model:
+                response = await self._create_structured_completion(
+                    model=model,
+                    messages=openai_messages,
+                    temperature=self.temperature,
+                    max_tokens=max_tokens or self.max_tokens,
+                    response_model=response_model,
+                )
+                return self._handle_structured_response(response)
             else:
-                raise Exception(f'Invalid response from LLM: {response_object.model_dump()}')
+                response = await self._create_completion(
+                    model=model,
+                    messages=openai_messages,
+                    temperature=self.temperature,
+                    max_tokens=max_tokens or self.max_tokens,
+                )
+                return self._handle_json_response(response)
+
         except openai.LengthFinishReasonError as e:
             raise Exception(f'Output length exceeded max tokens {self.max_tokens}: {e}') from e
         except openai.RateLimitError as e:
@@ -141,6 +166,7 @@ class OpenAIClient(LLMClient):
         max_tokens: int | None = None,
         model_size: ModelSize = ModelSize.medium,
     ) -> dict[str, typing.Any]:
+        """Generate a response with retry logic and error handling."""
         if max_tokens is None:
             max_tokens = self.max_tokens
 

graphiti_core-0.13.1/graphiti_core/llm_client/openai_client.py

@@ -0,0 +1,95 @@
+"""
+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 typing
+
+from openai import AsyncOpenAI
+from openai.types.chat import ChatCompletionMessageParam
+from pydantic import BaseModel
+
+from .config import DEFAULT_MAX_TOKENS, LLMConfig
+from .openai_base_client import BaseOpenAIClient
+
+
+class OpenAIClient(BaseOpenAIClient):
+    """
+    OpenAIClient is a client class for interacting with OpenAI's language models.
+
+    This class extends the BaseOpenAIClient and provides OpenAI-specific implementation
+    for creating completions.
+
+    Attributes:
+        client (AsyncOpenAI): The OpenAI client used to interact with the API.
+    """
+
+    def __init__(
+        self,
+        config: LLMConfig | None = None,
+        cache: bool = False,
+        client: typing.Any = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ):
+        """
+        Initialize the OpenAIClient with the provided configuration, cache setting, and client.
+
+        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.
+        """
+        super().__init__(config, cache, max_tokens)
+
+        if config is None:
+            config = LLMConfig()
+
+        if client is None:
+            self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
+        else:
+            self.client = client
+
+    async def _create_structured_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel],
+    ):
+        """Create a structured completion using OpenAI's beta parse API."""
+        return await self.client.beta.chat.completions.parse(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            response_format=response_model,  # type: ignore
+        )
+
+    async def _create_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel] | None = None,
+    ):
+        """Create a regular completion with JSON format."""
+        return await self.client.chat.completions.create(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            response_format={'type': 'json_object'},
+        )

{graphiti_core-0.13.0 → graphiti_core-0.13.1}/graphiti_core/search/search_utils.py

@@ -63,7 +63,7 @@ MAX_QUERY_LENGTH = 32
 
 def fulltext_query(query: str, group_ids: list[str] | None = None):
     group_ids_filter_list = (
-        [f"group_id-'{lucene_sanitize(g)}'" for g in group_ids] if group_ids is not None else []
+        [f'group_id:"{lucene_sanitize(g)}"' for g in group_ids] if group_ids is not None else []
     )
     group_ids_filter = ''
     for f in group_ids_filter_list:
@@ -301,12 +301,12 @@ async def edge_bfs_search(
 
     query = (
         """
-                            UNWIND $bfs_origin_node_uuids AS origin_uuid
-                            MATCH path = (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity)
-                            UNWIND relationships(path) AS rel
-                            MATCH (n:Entity)-[r:RELATES_TO]-(m:Entity)
-                            WHERE r.uuid = rel.uuid
-                            """
+                                    UNWIND $bfs_origin_node_uuids AS origin_uuid
+                                    MATCH path = (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity)
+                                    UNWIND relationships(path) AS rel
+                                    MATCH (n:Entity)-[r:RELATES_TO]-(m:Entity)
+                                    WHERE r.uuid = rel.uuid
+                                    """
         + filter_query
         + """  
                 RETURN DISTINCT
@@ -455,10 +455,10 @@ async def node_bfs_search(
 
     query = (
         """
-                    UNWIND $bfs_origin_node_uuids AS origin_uuid
-                    MATCH (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity)
-                    WHERE n.group_id = origin.group_id
-                    """
+                            UNWIND $bfs_origin_node_uuids AS origin_uuid
+                            MATCH (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity)
+                            WHERE n.group_id = origin.group_id
+                            """
         + filter_query
         + ENTITY_NODE_RETURN
         + """

{graphiti_core-0.13.0 → graphiti_core-0.13.1}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "graphiti-core"
 description = "A temporal graph building library"
-version = "0.13.0"
+version = "0.13.1"
 authors = [
     { "name" = "Paul Paliychuk", "email" = "paul@getzep.com" },
     { "name" = "Preston Rasmussen", "email" = "preston@getzep.com" },
@@ -15,7 +15,7 @@ dependencies = [
     "pydantic>=2.11.5",
     "neo4j>=5.26.0",
     "diskcache>=5.6.3",
-    "openai>=1.53.0",
+    "openai>=1.91.0",
     "tenacity>=9.0.0",
     "numpy>=1.0.0",
     "python-dotenv>=1.0.1",

graphiti_core-0.13.0/graphiti_core/llm_client/azure_openai_client.py

@@ -1,73 +0,0 @@
-"""
-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
-from typing import Any
-
-from openai import AsyncAzureOpenAI
-from openai.types.chat import ChatCompletionMessageParam
-from pydantic import BaseModel
-
-from ..prompts.models import Message
-from .client import LLMClient
-from .config import LLMConfig, ModelSize
-
-logger = logging.getLogger(__name__)
-
-
-class AzureOpenAILLMClient(LLMClient):
-    """Wrapper class for AsyncAzureOpenAI that implements the LLMClient interface."""
-
-    def __init__(self, azure_client: AsyncAzureOpenAI, config: LLMConfig | None = None):
-        super().__init__(config, cache=False)
-        self.azure_client = azure_client
-
-    async def _generate_response(
-        self,
-        messages: list[Message],
-        response_model: type[BaseModel] | None = None,
-        max_tokens: int = 1024,
-        model_size: ModelSize = ModelSize.medium,
-    ) -> dict[str, Any]:
-        """Generate response using Azure OpenAI client."""
-        # Convert messages to OpenAI format
-        openai_messages: list[ChatCompletionMessageParam] = []
-        for message in messages:
-            message.content = self._clean_input(message.content)
-            if message.role == 'user':
-                openai_messages.append({'role': 'user', 'content': message.content})
-            elif message.role == 'system':
-                openai_messages.append({'role': 'system', 'content': message.content})
-
-        # Ensure model is a string
-        model_name = self.model if self.model else 'gpt-4o-mini'
-
-        try:
-            response = await self.azure_client.chat.completions.create(
-                model=model_name,
-                messages=openai_messages,
-                temperature=float(self.temperature) if self.temperature is not None else 0.7,
-                max_tokens=max_tokens,
-                response_format={'type': 'json_object'},
-            )
-            result = response.choices[0].message.content or '{}'
-
-            # Parse JSON response
-            return json.loads(result)
-        except Exception as e:
-            logger.error(f'Error in Azure OpenAI LLM response: {e}')
-            raise