graphiti-core 0.12.0rc1__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/embedder/azure_openai.py

@@ -0,0 +1,71 @@
+"""
+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 Any
+
+from openai import AsyncAzureOpenAI, AsyncOpenAI
+
+from .client import EmbedderClient
+
+logger = logging.getLogger(__name__)
+
+
+class AzureOpenAIEmbedderClient(EmbedderClient):
+    """Wrapper class for Azure OpenAI that implements the EmbedderClient interface.
+
+    Supports both AsyncAzureOpenAI and AsyncOpenAI (with Azure v1 API endpoint).
+    """
+
+    def __init__(
+        self,
+        azure_client: AsyncAzureOpenAI | AsyncOpenAI,
+        model: str = 'text-embedding-3-small',
+    ):
+        self.azure_client = azure_client
+        self.model = model
+
+    async def create(self, input_data: str | list[str] | Any) -> list[float]:
+        """Create embeddings using Azure OpenAI client."""
+        try:
+            # Handle different input types
+            if isinstance(input_data, str):
+                text_input = [input_data]
+            elif isinstance(input_data, list) and all(isinstance(item, str) for item in input_data):
+                text_input = input_data
+            else:
+                # Convert to string list for other types
+                text_input = [str(input_data)]
+
+            response = await self.azure_client.embeddings.create(model=self.model, input=text_input)
+
+            # Return the first embedding as a list of floats
+            return response.data[0].embedding
+        except Exception as e:
+            logger.error(f'Error in Azure OpenAI embedding: {e}')
+            raise
+
+    async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
+        """Create batch embeddings using Azure OpenAI client."""
+        try:
+            response = await self.azure_client.embeddings.create(
+                model=self.model, input=input_data_list
+            )
+
+            return [embedding.embedding for embedding in response.data]
+        except Exception as e:
+            logger.error(f'Error in Azure OpenAI batch embedding: {e}')
+            raise

graphiti_core/embedder/client.py

@@ -14,12 +14,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 
 from pydantic import BaseModel, Field
 
-EMBEDDING_DIM = 1024
+EMBEDDING_DIM = int(os.getenv('EMBEDDING_DIM', 1024))
 
 
 class EmbedderConfig(BaseModel):

graphiti_core/embedder/gemini.py

@@ -14,15 +14,32 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import logging
 from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from google import genai
+    from google.genai import types
+else:
+    try:
+        from google import genai
+        from google.genai import types
+    except ImportError:
+        raise ImportError(
+            'google-genai is required for GeminiEmbedder. '
+            'Install it with: pip install graphiti-core[google-genai]'
+        ) from None
 
-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'
+logger = logging.getLogger(__name__)
+
+DEFAULT_EMBEDDING_MODEL = 'text-embedding-001'  # gemini-embedding-001 or text-embedding-005
+
+DEFAULT_BATCH_SIZE = 100
 
 
 class GeminiEmbedderConfig(EmbedderConfig):
@@ -35,15 +52,38 @@ class GeminiEmbedder(EmbedderClient):
     Google Gemini Embedder Client
     """
 
-    def __init__(self, config: GeminiEmbedderConfig | None = None):
+    def __init__(
+        self,
+        config: GeminiEmbedderConfig | None = None,
+        client: 'genai.Client | None' = None,
+        batch_size: int | None = None,
+    ):
+        """
+        Initialize the GeminiEmbedder with the provided configuration and client.
+
+        Args:
+            config (GeminiEmbedderConfig | None): The configuration for the GeminiEmbedder, including API key, model, base URL, temperature, and max tokens.
+            client (genai.Client | None): An optional async client instance to use. If not provided, a new genai.Client is created.
+            batch_size (int | None): An optional batch size to use. If not provided, the default batch size will be used.
+        """
         if config is None:
             config = GeminiEmbedderConfig()
+
         self.config = config
 
-        # Configure the Gemini API
-        self.client = genai.Client(
-            api_key=config.api_key,
-        )
+        if client is None:
+            self.client = genai.Client(api_key=config.api_key)
+        else:
+            self.client = client
+
+        if batch_size is None and self.config.embedding_model == 'gemini-embedding-001':
+            # Gemini API has a limit on the number of instances per request
+            # https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api
+            self.batch_size = 1
+        elif batch_size is None:
+            self.batch_size = DEFAULT_BATCH_SIZE
+        else:
+            self.batch_size = batch_size
 
     async def create(
         self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
@@ -71,19 +111,73 @@ class GeminiEmbedder(EmbedderClient):
         return result.embeddings[0].values
 
     async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
-        # Generate embeddings
-        result = await self.client.aio.models.embed_content(
-            model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
-            contents=input_data_list,  # type: ignore[arg-type]  # mypy fails on broad union type
-            config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
-        )
+        """
+        Create embeddings for a batch of input data using Google's Gemini embedding model.
 
-        if not result.embeddings or len(result.embeddings) == 0:
-            raise Exception('No embeddings returned')
+        This method handles batching to respect the Gemini API's limits on the number
+        of instances that can be processed in a single request.
 
-        embeddings = []
-        for embedding in result.embeddings:
-            if not embedding.values:
-                raise ValueError('Empty embedding values returned')
-            embeddings.append(embedding.values)
-        return embeddings
+        Args:
+            input_data_list: A list of strings to create embeddings for.
+
+        Returns:
+            A list of embedding vectors (each vector is a list of floats).
+        """
+        if not input_data_list:
+            return []
+
+        batch_size = self.batch_size
+        all_embeddings = []
+
+        # Process inputs in batches
+        for i in range(0, len(input_data_list), batch_size):
+            batch = input_data_list[i : i + batch_size]
+
+            try:
+                # Generate embeddings for this batch
+                result = await self.client.aio.models.embed_content(
+                    model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
+                    contents=batch,  # type: ignore[arg-type]  # mypy fails on broad union type
+                    config=types.EmbedContentConfig(
+                        output_dimensionality=self.config.embedding_dim
+                    ),
+                )
+
+                if not result.embeddings or len(result.embeddings) == 0:
+                    raise Exception('No embeddings returned')
+
+                # Process embeddings from this batch
+                for embedding in result.embeddings:
+                    if not embedding.values:
+                        raise ValueError('Empty embedding values returned')
+                    all_embeddings.append(embedding.values)
+
+            except Exception as e:
+                # If batch processing fails, fall back to individual processing
+                logger.warning(
+                    f'Batch embedding failed for batch {i // batch_size + 1}, falling back to individual processing: {e}'
+                )
+
+                for item in batch:
+                    try:
+                        # Process each item individually
+                        result = await self.client.aio.models.embed_content(
+                            model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
+                            contents=[item],  # type: ignore[arg-type]  # mypy fails on broad union type
+                            config=types.EmbedContentConfig(
+                                output_dimensionality=self.config.embedding_dim
+                            ),
+                        )
+
+                        if not result.embeddings or len(result.embeddings) == 0:
+                            raise ValueError('No embeddings returned from Gemini API')
+                        if not result.embeddings[0].values:
+                            raise ValueError('Empty embedding values returned')
+
+                        all_embeddings.append(result.embeddings[0].values)
+
+                    except Exception as individual_error:
+                        logger.error(f'Failed to embed individual item: {individual_error}')
+                        raise individual_error
+
+        return all_embeddings

graphiti_core/embedder/voyage.py

@@ -15,8 +15,19 @@ limitations under the License.
 """
 
 from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import voyageai
+else:
+    try:
+        import voyageai
+    except ImportError:
+        raise ImportError(
+            'voyageai is required for VoyageAIEmbedderClient. '
+            'Install it with: pip install graphiti-core[voyageai]'
+        ) from None
 
-import voyageai  # type: ignore
 from pydantic import Field
 
 from .client import EmbedderClient, EmbedderConfig
@@ -38,7 +49,7 @@ class VoyageAIEmbedder(EmbedderClient):
         if config is None:
             config = VoyageAIEmbedderConfig()
         self.config = config
-        self.client = voyageai.AsyncClient(api_key=config.api_key)
+        self.client = voyageai.AsyncClient(api_key=config.api_key)  # type: ignore[reportUnknownMemberType]
 
     async def create(
         self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]

graphiti_core/errors.py

@@ -73,3 +73,11 @@ class EntityTypeValidationError(GraphitiError):
     def __init__(self, entity_type: str, entity_type_attribute: str):
         self.message = f'{entity_type_attribute} cannot be used as an attribute for {entity_type} as it is a protected attribute name.'
         super().__init__(self.message)
+
+
+class GroupIdValidationError(GraphitiError):
+    """Raised when a group_id contains invalid characters."""
+
+    def __init__(self, group_id: str):
+        self.message = f'group_id "{group_id}" must contain only alphanumeric characters, dashes, or underscores'
+        super().__init__(self.message)

graphiti_core/graph_queries.py

@@ -0,0 +1,162 @@
+"""
+Database query utilities for different graph database backends.
+
+This module provides database-agnostic query generation for Neo4j and FalkorDB,
+supporting index creation, fulltext search, and bulk operations.
+"""
+
+from typing_extensions import LiteralString
+
+from graphiti_core.driver.driver import GraphProvider
+
+# Mapping from Neo4j fulltext index names to FalkorDB node labels
+NEO4J_TO_FALKORDB_MAPPING = {
+    'node_name_and_summary': 'Entity',
+    'community_name': 'Community',
+    'episode_content': 'Episodic',
+    'edge_name_and_fact': 'RELATES_TO',
+}
+# Mapping from fulltext index names to Kuzu node labels
+INDEX_TO_LABEL_KUZU_MAPPING = {
+    'node_name_and_summary': 'Entity',
+    'community_name': 'Community',
+    'episode_content': 'Episodic',
+    'edge_name_and_fact': 'RelatesToNode_',
+}
+
+
+def get_range_indices(provider: GraphProvider) -> list[LiteralString]:
+    if provider == GraphProvider.FALKORDB:
+        return [
+            # Entity node
+            'CREATE INDEX FOR (n:Entity) ON (n.uuid, n.group_id, n.name, n.created_at)',
+            # Episodic node
+            'CREATE INDEX FOR (n:Episodic) ON (n.uuid, n.group_id, n.created_at, n.valid_at)',
+            # Community node
+            'CREATE INDEX FOR (n:Community) ON (n.uuid)',
+            # RELATES_TO edge
+            'CREATE INDEX FOR ()-[e:RELATES_TO]-() ON (e.uuid, e.group_id, e.name, e.created_at, e.expired_at, e.valid_at, e.invalid_at)',
+            # MENTIONS edge
+            'CREATE INDEX FOR ()-[e:MENTIONS]-() ON (e.uuid, e.group_id)',
+            # HAS_MEMBER edge
+            'CREATE INDEX FOR ()-[e:HAS_MEMBER]-() ON (e.uuid)',
+        ]
+
+    if provider == GraphProvider.KUZU:
+        return []
+
+    return [
+        'CREATE INDEX entity_uuid IF NOT EXISTS FOR (n:Entity) ON (n.uuid)',
+        'CREATE INDEX episode_uuid IF NOT EXISTS FOR (n:Episodic) ON (n.uuid)',
+        'CREATE INDEX community_uuid IF NOT EXISTS FOR (n:Community) ON (n.uuid)',
+        'CREATE INDEX relation_uuid IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.uuid)',
+        'CREATE INDEX mention_uuid IF NOT EXISTS FOR ()-[e:MENTIONS]-() ON (e.uuid)',
+        'CREATE INDEX has_member_uuid IF NOT EXISTS FOR ()-[e:HAS_MEMBER]-() ON (e.uuid)',
+        'CREATE INDEX entity_group_id IF NOT EXISTS FOR (n:Entity) ON (n.group_id)',
+        'CREATE INDEX episode_group_id IF NOT EXISTS FOR (n:Episodic) ON (n.group_id)',
+        'CREATE INDEX community_group_id IF NOT EXISTS FOR (n:Community) ON (n.group_id)',
+        'CREATE INDEX relation_group_id IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.group_id)',
+        'CREATE INDEX mention_group_id IF NOT EXISTS FOR ()-[e:MENTIONS]-() ON (e.group_id)',
+        'CREATE INDEX name_entity_index IF NOT EXISTS FOR (n:Entity) ON (n.name)',
+        'CREATE INDEX created_at_entity_index IF NOT EXISTS FOR (n:Entity) ON (n.created_at)',
+        'CREATE INDEX created_at_episodic_index IF NOT EXISTS FOR (n:Episodic) ON (n.created_at)',
+        'CREATE INDEX valid_at_episodic_index IF NOT EXISTS FOR (n:Episodic) ON (n.valid_at)',
+        'CREATE INDEX name_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.name)',
+        'CREATE INDEX created_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.created_at)',
+        'CREATE INDEX expired_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.expired_at)',
+        'CREATE INDEX valid_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.valid_at)',
+        'CREATE INDEX invalid_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.invalid_at)',
+    ]
+
+
+def get_fulltext_indices(provider: GraphProvider) -> list[LiteralString]:
+    if provider == GraphProvider.FALKORDB:
+        from typing import cast
+
+        from graphiti_core.driver.falkordb_driver import STOPWORDS
+
+        # Convert to string representation for embedding in queries
+        stopwords_str = str(STOPWORDS)
+
+        # Use type: ignore to satisfy LiteralString requirement while maintaining single source of truth
+        return cast(
+            list[LiteralString],
+            [
+                f"""CALL db.idx.fulltext.createNodeIndex(
+                                                {{
+                                                    label: 'Episodic',
+                                                    stopwords: {stopwords_str}
+                                                }},
+                                                'content', 'source', 'source_description', 'group_id'
+                                                )""",
+                f"""CALL db.idx.fulltext.createNodeIndex(
+                                                {{
+                                                    label: 'Entity',
+                                                    stopwords: {stopwords_str}
+                                                }},
+                                                'name', 'summary', 'group_id'
+                                                )""",
+                f"""CALL db.idx.fulltext.createNodeIndex(
+                                                {{
+                                                    label: 'Community',
+                                                    stopwords: {stopwords_str}
+                                                }},
+                                                'name', 'group_id'
+                                                )""",
+                """CREATE FULLTEXT INDEX FOR ()-[e:RELATES_TO]-() ON (e.name, e.fact, e.group_id)""",
+            ],
+        )
+
+    if provider == GraphProvider.KUZU:
+        return [
+            "CALL CREATE_FTS_INDEX('Episodic', 'episode_content', ['content', 'source', 'source_description']);",
+            "CALL CREATE_FTS_INDEX('Entity', 'node_name_and_summary', ['name', 'summary']);",
+            "CALL CREATE_FTS_INDEX('Community', 'community_name', ['name']);",
+            "CALL CREATE_FTS_INDEX('RelatesToNode_', 'edge_name_and_fact', ['name', 'fact']);",
+        ]
+
+    return [
+        """CREATE FULLTEXT INDEX episode_content IF NOT EXISTS
+        FOR (e:Episodic) ON EACH [e.content, e.source, e.source_description, e.group_id]""",
+        """CREATE FULLTEXT INDEX node_name_and_summary IF NOT EXISTS
+        FOR (n:Entity) ON EACH [n.name, n.summary, n.group_id]""",
+        """CREATE FULLTEXT INDEX community_name IF NOT EXISTS
+        FOR (n:Community) ON EACH [n.name, n.group_id]""",
+        """CREATE FULLTEXT INDEX edge_name_and_fact IF NOT EXISTS
+        FOR ()-[e:RELATES_TO]-() ON EACH [e.name, e.fact, e.group_id]""",
+    ]
+
+
+def get_nodes_query(name: str, query: str, limit: int, provider: GraphProvider) -> str:
+    if provider == GraphProvider.FALKORDB:
+        label = NEO4J_TO_FALKORDB_MAPPING[name]
+        return f"CALL db.idx.fulltext.queryNodes('{label}', {query})"
+
+    if provider == GraphProvider.KUZU:
+        label = INDEX_TO_LABEL_KUZU_MAPPING[name]
+        return f"CALL QUERY_FTS_INDEX('{label}', '{name}', {query}, TOP := $limit)"
+
+    return f'CALL db.index.fulltext.queryNodes("{name}", {query}, {{limit: $limit}})'
+
+
+def get_vector_cosine_func_query(vec1, vec2, provider: GraphProvider) -> str:
+    if provider == GraphProvider.FALKORDB:
+        # FalkorDB uses a different syntax for regular cosine similarity and Neo4j uses normalized cosine similarity
+        return f'(2 - vec.cosineDistance({vec1}, vecf32({vec2})))/2'
+
+    if provider == GraphProvider.KUZU:
+        return f'array_cosine_similarity({vec1}, {vec2})'
+
+    return f'vector.similarity.cosine({vec1}, {vec2})'
+
+
+def get_relationships_query(name: str, limit: int, provider: GraphProvider) -> str:
+    if provider == GraphProvider.FALKORDB:
+        label = NEO4J_TO_FALKORDB_MAPPING[name]
+        return f"CALL db.idx.fulltext.queryRelationships('{label}', $query)"
+
+    if provider == GraphProvider.KUZU:
+        label = INDEX_TO_LABEL_KUZU_MAPPING[name]
+        return f"CALL QUERY_FTS_INDEX('{label}', '{name}', cast($query AS STRING), TOP := $limit)"
+
+    return f'CALL db.index.fulltext.queryRelationships("{name}", $query, {{limit: $limit}})'