graphiti-core 0.1.0__py3-none-any.whl

graphiti_core/llm_client/client.py

@@ -0,0 +1,96 @@
+"""
+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 hashlib
+import json
+import logging
+import typing
+from abc import ABC, abstractmethod
+
+import httpx
+from diskcache import Cache
+from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
+
+from ..prompts.models import Message
+from .config import LLMConfig
+
+DEFAULT_TEMPERATURE = 0
+DEFAULT_CACHE_DIR = './llm_cache'
+
+logger = logging.getLogger(__name__)
+
+
+def is_server_error(exception):
+    return (
+        isinstance(exception, httpx.HTTPStatusError) and 500 <= exception.response.status_code < 600
+    )
+
+
+class LLMClient(ABC):
+    def __init__(self, config: LLMConfig | None, cache: bool = False):
+        if config is None:
+            config = LLMConfig()
+
+        self.config = config
+        self.model = config.model
+        self.temperature = config.temperature
+        self.max_tokens = config.max_tokens
+        self.cache_enabled = cache
+        self.cache_dir = Cache(DEFAULT_CACHE_DIR)  # Create a cache directory
+
+    @abstractmethod
+    def get_embedder(self) -> typing.Any:
+        pass
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception(is_server_error),
+    )
+    async def _generate_response_with_retry(self, messages: list[Message]) -> dict[str, typing.Any]:
+        try:
+            return await self._generate_response(messages)
+        except httpx.HTTPStatusError as e:
+            if not is_server_error(e):
+                raise Exception(f'LLM request error: {e}') from e
+            else:
+                raise
+
+    @abstractmethod
+    async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+        pass
+
+    def _get_cache_key(self, messages: list[Message]) -> str:
+        # Create a unique cache key based on the messages and model
+        message_str = json.dumps([m.model_dump() for m in messages], sort_keys=True)
+        key_str = f'{self.model}:{message_str}'
+        return hashlib.md5(key_str.encode()).hexdigest()
+
+    async def generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+        if self.cache_enabled:
+            cache_key = self._get_cache_key(messages)
+
+            cached_response = self.cache_dir.get(cache_key)
+            if cached_response is not None:
+                logger.debug(f'Cache hit for {cache_key}')
+                return cached_response
+
+        response = await self._generate_response_with_retry(messages)
+
+        if self.cache_enabled:
+            self.cache_dir.set(cache_key, response)
+
+        return response

graphiti_core/llm_client/config.py

@@ -0,0 +1,58 @@
+"""
+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.
+"""
+
+EMBEDDING_DIM = 1024
+DEFAULT_MAX_TOKENS = 4096
+DEFAULT_TEMPERATURE = 0
+
+
+class LLMConfig:
+    """
+    Configuration class for the Language Learning Model (LLM).
+
+    This class encapsulates the necessary parameters to interact with an LLM API,
+    such as OpenAI's GPT models. It stores the API key, model name, and base URL
+    for making requests to the LLM service.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model: str | None = None,
+        base_url: str | None = None,
+        temperature: float = DEFAULT_TEMPERATURE,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ):
+        """
+        Initialize the LLMConfig with the provided parameters.
+
+        Args:
+                api_key (str): The authentication key for accessing the LLM API.
+                                                This is required for making authorized requests.
+
+                model (str, optional): The specific LLM model to use for generating responses.
+                                                                Defaults to "gpt-4o-mini", which appears to be a custom model name.
+                                                                Common values might include "gpt-3.5-turbo" or "gpt-4".
+
+                base_url (str, optional): The base URL of the LLM API service.
+                                                                        Defaults to "https://api.openai.com", which is OpenAI's standard API endpoint.
+                                                                        This can be changed if using a different provider or a custom endpoint.
+        """
+        self.base_url = base_url
+        self.api_key = api_key
+        self.model = model
+        self.temperature = temperature
+        self.max_tokens = max_tokens

graphiti_core/llm_client/groq_client.py

@@ -0,0 +1,64 @@
+"""
+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 groq import AsyncGroq
+from groq.types.chat import ChatCompletionMessageParam
+from openai import AsyncOpenAI
+
+from ..prompts.models import Message
+from .client import LLMClient
+from .config import LLMConfig
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = 'llama-3.1-70b-versatile'
+
+
+class GroqClient(LLMClient):
+    def __init__(self, config: LLMConfig | None = None, cache: bool = False):
+        if config is None:
+            config = LLMConfig()
+        super().__init__(config, cache)
+        self.client = AsyncGroq(api_key=config.api_key)
+
+    def get_embedder(self) -> typing.Any:
+        openai_client = AsyncOpenAI()
+        return openai_client.embeddings
+
+    async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+        msgs: list[ChatCompletionMessageParam] = []
+        for m in messages:
+            if m.role == 'user':
+                msgs.append({'role': 'user', 'content': m.content})
+            elif m.role == 'system':
+                msgs.append({'role': 'system', 'content': m.content})
+        try:
+            response = await self.client.chat.completions.create(
+                model=self.model or DEFAULT_MODEL,
+                messages=msgs,
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+                response_format={'type': 'json_object'},
+            )
+            result = response.choices[0].message.content or ''
+            return json.loads(result)
+        except Exception as e:
+            logger.error(f'Error in generating LLM response: {e}')
+            raise

graphiti_core/llm_client/openai_client.py

@@ -0,0 +1,65 @@
+"""
+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 openai import AsyncOpenAI
+from openai.types.chat import ChatCompletionMessageParam
+
+from ..prompts.models import Message
+from .client import LLMClient
+from .config import LLMConfig
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = 'gpt-4o-2024-08-06'
+
+
+class OpenAIClient(LLMClient):
+    def __init__(self, config: LLMConfig | None = None, cache: bool = False):
+        if config is None:
+            config = LLMConfig()
+
+        super().__init__(config, cache)
+
+        self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
+
+    def get_embedder(self) -> typing.Any:
+        return self.client.embeddings
+
+    async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+        openai_messages: list[ChatCompletionMessageParam] = []
+        for m in messages:
+            if m.role == 'user':
+                openai_messages.append({'role': 'user', 'content': m.content})
+            elif m.role == 'system':
+                openai_messages.append({'role': 'system', 'content': m.content})
+        try:
+            response = await self.client.chat.completions.create(
+                model=self.model or DEFAULT_MODEL,
+                messages=openai_messages,
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+                response_format={'type': 'json_object'},
+            )
+            result = response.choices[0].message.content or ''
+            return json.loads(result)
+        except Exception as e:
+            print(openai_messages)
+            logger.error(f'Error in generating LLM response: {e}')
+            raise

graphiti_core/llm_client/utils.py

@@ -0,0 +1,22 @@
+import logging
+import typing
+from time import time
+
+from graphiti_core.llm_client.config import EMBEDDING_DIM
+
+logger = logging.getLogger(__name__)
+
+
+async def generate_embedding(
+    embedder: typing.Any, text: str, model: str = 'text-embedding-3-small'
+):
+    start = time()
+
+    text = text.replace('\n', ' ')
+    embedding = (await embedder.create(input=[text], model=model)).data[0].embedding
+    embedding = embedding[:EMBEDDING_DIM]
+
+    end = time()
+    logger.debug(f'embedded text of length {len(text)} in {end-start} ms')
+
+    return embedding

graphiti_core/nodes.py

@@ -0,0 +1,250 @@
+"""
+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 abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from time import time
+from uuid import uuid4
+
+from neo4j import AsyncDriver
+from openai import OpenAI
+from pydantic import BaseModel, Field
+
+from graphiti_core.llm_client.config import EMBEDDING_DIM
+
+logger = logging.getLogger(__name__)
+
+
+class EpisodeType(Enum):
+    """
+    Enumeration of different types of episodes that can be processed.
+
+    This enum defines the various sources or formats of episodes that the system
+    can handle. It's used to categorize and potentially handle different types
+    of input data differently.
+
+    Attributes:
+    -----------
+    message : str
+        Represents a standard message-type episode. The content for this type
+        should be formatted as "actor: content". For example, "user: Hello, how are you?"
+        or "assistant: I'm doing well, thank you for asking."
+    json : str
+        Represents an episode containing a JSON string object with structured data.
+    text : str
+        Represents a plain text episode.
+    """
+
+    message = 'message'
+    json = 'json'
+    text = 'text'
+
+    @staticmethod
+    def from_str(episode_type: str):
+        if episode_type == 'message':
+            return EpisodeType.message
+        if episode_type == 'json':
+            return EpisodeType.json
+        if episode_type == 'text':
+            return EpisodeType.text
+        logger.error(f'Episode type: {episode_type} not implemented')
+        raise NotImplementedError
+
+
+class Node(BaseModel, ABC):
+    uuid: str = Field(default_factory=lambda: uuid4().hex)
+    name: str = Field(description='name of the node')
+    labels: list[str] = Field(default_factory=list)
+    created_at: datetime = Field(default_factory=lambda: datetime.now())
+
+    @abstractmethod
+    async def save(self, driver: AsyncDriver): ...
+
+    @abstractmethod
+    async def delete(self, driver: AsyncDriver): ...
+
+    def __hash__(self):
+        return hash(self.uuid)
+
+    def __eq__(self, other):
+        if isinstance(other, Node):
+            return self.uuid == other.uuid
+        return False
+
+    @classmethod
+    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str): ...
+
+
+class EpisodicNode(Node):
+    source: EpisodeType = Field(description='source type')
+    source_description: str = Field(description='description of the data source')
+    content: str = Field(description='raw episode data')
+    valid_at: datetime = Field(
+        description='datetime of when the original document was created',
+    )
+    entity_edges: list[str] = Field(
+        description='list of entity edges referenced in this episode',
+        default_factory=list,
+    )
+
+    async def save(self, driver: AsyncDriver):
+        result = await driver.execute_query(
+            """
+        MERGE (n:Episodic {uuid: $uuid})
+        SET n = {uuid: $uuid, name: $name, source_description: $source_description, source: $source, content: $content, 
+        entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
+        RETURN n.uuid AS uuid""",
+            uuid=self.uuid,
+            name=self.name,
+            source_description=self.source_description,
+            content=self.content,
+            entity_edges=self.entity_edges,
+            created_at=self.created_at,
+            valid_at=self.valid_at,
+            source=self.source.value,
+        )
+
+        logger.info(f'Saved Node to neo4j: {self.uuid}')
+
+        return result
+
+    async def delete(self, driver: AsyncDriver):
+        result = await driver.execute_query(
+            """
+        MATCH (n:Episodic {uuid: $uuid})
+        DETACH DELETE n
+        """,
+            uuid=self.uuid,
+        )
+
+        logger.info(f'Deleted Node: {self.uuid}')
+
+        return result
+
+    @classmethod
+    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (e:Episodic {uuid: $uuid})
+            RETURN e.content as content,
+            e.created_at as created_at,
+            e.valid_at as valid_at,
+            e.uuid as uuid,
+            e.name as name,
+            e.source_description as source_description,
+            e.source as source
+        """,
+            uuid=uuid,
+        )
+
+        episodes = [
+            EpisodicNode(
+                content=record['content'],
+                created_at=record['created_at'].to_native().timestamp(),
+                valid_at=(record['valid_at'].to_native()),
+                uuid=record['uuid'],
+                source=EpisodeType.from_str(record['source']),
+                name=record['name'],
+                source_description=record['source_description'],
+            )
+            for record in records
+        ]
+
+        logger.info(f'Found Node: {uuid}')
+
+        return episodes[0]
+
+
+class EntityNode(Node):
+    name_embedding: list[float] | None = Field(default=None, description='embedding of the name')
+    summary: str = Field(description='regional summary of surrounding edges', default_factory=str)
+
+    async def update_summary(self, driver: AsyncDriver): ...
+
+    async def refresh_summary(self, driver: AsyncDriver, llm_client: OpenAI): ...
+
+    async def generate_name_embedding(self, embedder, model='text-embedding-3-small'):
+        start = time()
+        text = self.name.replace('\n', ' ')
+        embedding = (await embedder.create(input=[text], model=model)).data[0].embedding
+        self.name_embedding = embedding[:EMBEDDING_DIM]
+        end = time()
+        logger.info(f'embedded {text} in {end - start} ms')
+
+        return embedding
+
+    async def save(self, driver: AsyncDriver):
+        result = await driver.execute_query(
+            """
+        MERGE (n:Entity {uuid: $uuid})
+        SET n = {uuid: $uuid, name: $name, name_embedding: $name_embedding, summary: $summary, created_at: $created_at}
+        RETURN n.uuid AS uuid""",
+            uuid=self.uuid,
+            name=self.name,
+            summary=self.summary,
+            name_embedding=self.name_embedding,
+            created_at=self.created_at,
+        )
+
+        logger.info(f'Saved Node to neo4j: {self.uuid}')
+
+        return result
+
+    async def delete(self, driver: AsyncDriver):
+        result = await driver.execute_query(
+            """
+        MATCH (n:Entity {uuid: $uuid})
+        DETACH DELETE n
+        """,
+            uuid=self.uuid,
+        )
+
+        logger.info(f'Deleted Node: {self.uuid}')
+
+        return result
+
+    @classmethod
+    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Entity {uuid: $uuid})
+        RETURN
+            n.uuid As uuid, 
+            n.name AS name, 
+            n.created_at AS created_at, 
+            n.summary AS summary
+        """,
+            uuid=uuid,
+        )
+
+        nodes: list[EntityNode] = []
+
+        for record in records:
+            nodes.append(
+                EntityNode(
+                    uuid=record['uuid'],
+                    name=record['name'],
+                    labels=['Entity'],
+                    created_at=record['created_at'].to_native(),
+                    summary=record['summary'],
+                )
+            )
+
+        logger.info(f'Found Node: {uuid}')
+
+        return nodes[0]

graphiti_core/prompts/__init__.py

@@ -0,0 +1,4 @@
+from .lib import prompt_library
+from .models import Message
+
+__all__ = ['prompt_library', 'Message']