graphiti-core 0.3.6__py3-none-any.whl → 0.3.8__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of graphiti-core might be problematic. Click here for more details.

graphiti_core/edges.py CHANGED
@@ -24,9 +24,9 @@ from uuid import uuid4
24
24
  from neo4j import AsyncDriver
25
25
  from pydantic import BaseModel, Field
26
26
 
27
+ from graphiti_core.embedder import EmbedderClient
27
28
  from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
28
29
  from graphiti_core.helpers import parse_db_date
29
- from graphiti_core.llm_client.config import EMBEDDING_DIM
30
30
  from graphiti_core.nodes import Node
31
31
 
32
32
  logger = logging.getLogger(__name__)
@@ -171,17 +171,16 @@ class EntityEdge(Edge):
171
171
  default=None, description='datetime of when the fact stopped being true'
172
172
  )
173
173
 
174
- async def generate_embedding(self, embedder, model='text-embedding-3-small'):
174
+ async def generate_embedding(self, embedder: EmbedderClient):
175
175
  start = time()
176
176
 
177
177
  text = self.fact.replace('\n', ' ')
178
- embedding = (await embedder.create(input=[text], model=model)).data[0].embedding
179
- self.fact_embedding = embedding[:EMBEDDING_DIM]
178
+ self.fact_embedding = await embedder.create(input=[text])
180
179
 
181
180
  end = time()
182
181
  logger.info(f'embedded {text} in {end - start} ms')
183
182
 
184
- return embedding
183
+ return self.fact_embedding
185
184
 
186
185
  async def save(self, driver: AsyncDriver):
187
186
  result = await driver.execute_query(
@@ -0,0 +1,4 @@
1
+ from .client import EmbedderClient
2
+ from .openai import OpenAIEmbedder, OpenAIEmbedderConfig
3
+
4
+ __all__ = ['EmbedderClient', 'OpenAIEmbedder', 'OpenAIEmbedderConfig']
@@ -0,0 +1,34 @@
1
+ """
2
+ Copyright 2024, Zep Software, Inc.
3
+
4
+ Licensed under the Apache License, Version 2.0 (the "License");
5
+ you may not use this file except in compliance with the License.
6
+ You may obtain a copy of the License at
7
+
8
+ http://www.apache.org/licenses/LICENSE-2.0
9
+
10
+ Unless required by applicable law or agreed to in writing, software
11
+ distributed under the License is distributed on an "AS IS" BASIS,
12
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ See the License for the specific language governing permissions and
14
+ limitations under the License.
15
+ """
16
+
17
+ from abc import ABC, abstractmethod
18
+ from typing import Iterable, List, Literal
19
+
20
+ from pydantic import BaseModel, Field
21
+
22
+ EMBEDDING_DIM = 1024
23
+
24
+
25
+ class EmbedderConfig(BaseModel):
26
+ embedding_dim: Literal[1024] = Field(default=EMBEDDING_DIM, frozen=True)
27
+
28
+
29
+ class EmbedderClient(ABC):
30
+ @abstractmethod
31
+ async def create(
32
+ self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
33
+ ) -> list[float]:
34
+ pass
@@ -0,0 +1,48 @@
1
+ """
2
+ Copyright 2024, Zep Software, Inc.
3
+
4
+ Licensed under the Apache License, Version 2.0 (the "License");
5
+ you may not use this file except in compliance with the License.
6
+ You may obtain a copy of the License at
7
+
8
+ http://www.apache.org/licenses/LICENSE-2.0
9
+
10
+ Unless required by applicable law or agreed to in writing, software
11
+ distributed under the License is distributed on an "AS IS" BASIS,
12
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ See the License for the specific language governing permissions and
14
+ limitations under the License.
15
+ """
16
+
17
+ from typing import Iterable, List
18
+
19
+ from openai import AsyncOpenAI
20
+ from openai.types import EmbeddingModel
21
+
22
+ from .client import EmbedderClient, EmbedderConfig
23
+
24
+ DEFAULT_EMBEDDING_MODEL = 'text-embedding-3-small'
25
+
26
+
27
+ class OpenAIEmbedderConfig(EmbedderConfig):
28
+ embedding_model: EmbeddingModel | str = DEFAULT_EMBEDDING_MODEL
29
+ api_key: str | None = None
30
+ base_url: str | None = None
31
+
32
+
33
+ class OpenAIEmbedder(EmbedderClient):
34
+ """
35
+ OpenAI Embedder Client
36
+ """
37
+
38
+ def __init__(self, config: OpenAIEmbedderConfig | None = None):
39
+ if config is None:
40
+ config = OpenAIEmbedderConfig()
41
+ self.config = config
42
+ self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
43
+
44
+ async def create(
45
+ self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
46
+ ) -> list[float]:
47
+ result = await self.client.embeddings.create(input=input, model=self.config.embedding_model)
48
+ return result.data[0].embedding[: self.config.embedding_dim]
@@ -0,0 +1,47 @@
1
+ """
2
+ Copyright 2024, Zep Software, Inc.
3
+
4
+ Licensed under the Apache License, Version 2.0 (the "License");
5
+ you may not use this file except in compliance with the License.
6
+ You may obtain a copy of the License at
7
+
8
+ http://www.apache.org/licenses/LICENSE-2.0
9
+
10
+ Unless required by applicable law or agreed to in writing, software
11
+ distributed under the License is distributed on an "AS IS" BASIS,
12
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ See the License for the specific language governing permissions and
14
+ limitations under the License.
15
+ """
16
+
17
+ from typing import Iterable, List
18
+
19
+ import voyageai # type: ignore
20
+ from pydantic import Field
21
+
22
+ from .client import EmbedderClient, EmbedderConfig
23
+
24
+ DEFAULT_EMBEDDING_MODEL = 'voyage-3'
25
+
26
+
27
+ class VoyageAIEmbedderConfig(EmbedderConfig):
28
+ embedding_model: str = Field(default=DEFAULT_EMBEDDING_MODEL)
29
+ api_key: str | None = None
30
+
31
+
32
+ class VoyageAIEmbedder(EmbedderClient):
33
+ """
34
+ VoyageAI Embedder Client
35
+ """
36
+
37
+ def __init__(self, config: VoyageAIEmbedderConfig | None = None):
38
+ if config is None:
39
+ config = VoyageAIEmbedderConfig()
40
+ self.config = config
41
+ self.client = voyageai.AsyncClient(api_key=config.api_key)
42
+
43
+ async def create(
44
+ self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
45
+ ) -> list[float]:
46
+ result = await self.client.embed(input, model=self.config.embedding_model)
47
+ return result.embeddings[0][: self.config.embedding_dim]
graphiti_core/graphiti.py CHANGED
@@ -23,6 +23,7 @@ from dotenv import load_dotenv
23
23
  from neo4j import AsyncGraphDatabase
24
24
 
25
25
  from graphiti_core.edges import EntityEdge, EpisodicEdge
26
+ from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
26
27
  from graphiti_core.llm_client import LLMClient, OpenAIClient
27
28
  from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
28
29
  from graphiti_core.search.search import SearchConfig, search
@@ -83,6 +84,7 @@ class Graphiti:
83
84
  user: str,
84
85
  password: str,
85
86
  llm_client: LLMClient | None = None,
87
+ embedder: EmbedderClient | None = None,
86
88
  store_raw_episode_content: bool = True,
87
89
  ):
88
90
  """
@@ -128,6 +130,10 @@ class Graphiti:
128
130
  self.llm_client = llm_client
129
131
  else:
130
132
  self.llm_client = OpenAIClient()
133
+ if embedder:
134
+ self.embedder = embedder
135
+ else:
136
+ self.embedder = OpenAIEmbedder()
131
137
 
132
138
  async def close(self):
133
139
  """
@@ -161,7 +167,7 @@ class Graphiti:
161
167
  """
162
168
  await self.driver.close()
163
169
 
164
- async def build_indices_and_constraints(self):
170
+ async def build_indices_and_constraints(self, delete_existing: bool = False):
165
171
  """
166
172
  Build indices and constraints in the Neo4j database.
167
173
 
@@ -171,6 +177,9 @@ class Graphiti:
171
177
  Parameters
172
178
  ----------
173
179
  self
180
+ delete_existing : bool, optional
181
+ Whether to clear existing indices before creating new ones.
182
+
174
183
 
175
184
  Returns
176
185
  -------
@@ -191,7 +200,7 @@ class Graphiti:
191
200
  Caution: Running this method on a large existing database may take some time
192
201
  and could impact database performance during execution.
193
202
  """
194
- await build_indices_and_constraints(self.driver)
203
+ await build_indices_and_constraints(self.driver, delete_existing)
195
204
 
196
205
  async def retrieve_episodes(
197
206
  self,
@@ -287,7 +296,6 @@ class Graphiti:
287
296
  start = time()
288
297
 
289
298
  entity_edges: list[EntityEdge] = []
290
- embedder = self.llm_client.get_embedder()
291
299
  now = datetime.now()
292
300
 
293
301
  previous_episodes = await self.retrieve_episodes(
@@ -315,7 +323,7 @@ class Graphiti:
315
323
  # Calculate Embeddings
316
324
 
317
325
  await asyncio.gather(
318
- *[node.generate_name_embedding(embedder) for node in extracted_nodes]
326
+ *[node.generate_name_embedding(self.embedder) for node in extracted_nodes]
319
327
  )
320
328
 
321
329
  # Resolve extracted nodes with nodes already in the graph and extract facts
@@ -343,7 +351,7 @@ class Graphiti:
343
351
  # calculate embeddings
344
352
  await asyncio.gather(
345
353
  *[
346
- edge.generate_embedding(embedder)
354
+ edge.generate_embedding(self.embedder)
347
355
  for edge in extracted_edges_with_resolved_pointers
348
356
  ]
349
357
  )
@@ -436,7 +444,7 @@ class Graphiti:
436
444
  if update_communities:
437
445
  await asyncio.gather(
438
446
  *[
439
- update_community(self.driver, self.llm_client, embedder, node)
447
+ update_community(self.driver, self.llm_client, self.embedder, node)
440
448
  for node in nodes
441
449
  ]
442
450
  )
@@ -485,7 +493,6 @@ class Graphiti:
485
493
  """
486
494
  try:
487
495
  start = time()
488
- embedder = self.llm_client.get_embedder()
489
496
  now = datetime.now()
490
497
 
491
498
  episodes = [
@@ -517,8 +524,8 @@ class Graphiti:
517
524
 
518
525
  # Generate embeddings
519
526
  await asyncio.gather(
520
- *[node.generate_name_embedding(embedder) for node in extracted_nodes],
521
- *[edge.generate_embedding(embedder) for edge in extracted_edges],
527
+ *[node.generate_name_embedding(self.embedder) for node in extracted_nodes],
528
+ *[edge.generate_embedding(self.embedder) for edge in extracted_edges],
522
529
  )
523
530
 
524
531
  # Dedupe extracted nodes, compress extracted edges
@@ -561,14 +568,14 @@ class Graphiti:
561
568
  raise e
562
569
 
563
570
  async def build_communities(self):
564
- embedder = self.llm_client.get_embedder()
565
-
566
571
  # Clear existing communities
567
572
  await remove_communities(self.driver)
568
573
 
569
574
  community_nodes, community_edges = await build_communities(self.driver, self.llm_client)
570
575
 
571
- await asyncio.gather(*[node.generate_name_embedding(embedder) for node in community_nodes])
576
+ await asyncio.gather(
577
+ *[node.generate_name_embedding(self.embedder) for node in community_nodes]
578
+ )
572
579
 
573
580
  await asyncio.gather(*[node.save(self.driver) for node in community_nodes])
574
581
  await asyncio.gather(*[edge.save(self.driver) for edge in community_edges])
@@ -619,7 +626,7 @@ class Graphiti:
619
626
  edges = (
620
627
  await search(
621
628
  self.driver,
622
- self.llm_client.get_embedder(),
629
+ self.embedder,
623
630
  query,
624
631
  group_ids,
625
632
  search_config,
@@ -636,9 +643,7 @@ class Graphiti:
636
643
  group_ids: list[str] | None = None,
637
644
  center_node_uuid: str | None = None,
638
645
  ) -> SearchResults:
639
- return await search(
640
- self.driver, self.llm_client.get_embedder(), query, group_ids, config, center_node_uuid
641
- )
646
+ return await search(self.driver, self.embedder, query, group_ids, config, center_node_uuid)
642
647
 
643
648
  async def get_nodes_by_query(
644
649
  self,
@@ -683,14 +688,15 @@ class Graphiti:
683
688
  to each individual search method before results are combined and deduplicated.
684
689
  If not specified, a default limit (defined in the search functions) will be used.
685
690
  """
686
- embedder = self.llm_client.get_embedder()
687
691
  search_config = (
688
692
  NODE_HYBRID_SEARCH_RRF if center_node_uuid is None else NODE_HYBRID_SEARCH_NODE_DISTANCE
689
693
  )
690
694
  search_config.limit = limit
691
695
 
692
696
  nodes = (
693
- await search(self.driver, embedder, query, group_ids, search_config, center_node_uuid)
697
+ await search(
698
+ self.driver, self.embedder, query, group_ids, search_config, center_node_uuid
699
+ )
694
700
  ).nodes
695
701
  return nodes
696
702
 
graphiti_core/helpers.py CHANGED
@@ -21,3 +21,33 @@ from neo4j import time as neo4j_time
21
21
 
22
22
  def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:
23
23
  return neo_date.to_native() if neo_date else None
24
+
25
+
26
+ def lucene_sanitize(query: str) -> str:
27
+ # Escape special characters from a query before passing into Lucene
28
+ # + - && || ! ( ) { } [ ] ^ " ~ * ? : \
29
+ escape_map = str.maketrans(
30
+ {
31
+ '+': r'\+',
32
+ '-': r'\-',
33
+ '&': r'\&',
34
+ '|': r'\|',
35
+ '!': r'\!',
36
+ '(': r'\(',
37
+ ')': r'\)',
38
+ '{': r'\{',
39
+ '}': r'\}',
40
+ '[': r'\[',
41
+ ']': r'\]',
42
+ '^': r'\^',
43
+ '"': r'\"',
44
+ '~': r'\~',
45
+ '*': r'\*',
46
+ '?': r'\?',
47
+ ':': r'\:',
48
+ '\\': r'\\',
49
+ }
50
+ )
51
+
52
+ sanitized = query.translate(escape_map)
53
+ return sanitized
@@ -20,7 +20,6 @@ import typing
20
20
 
21
21
  import anthropic
22
22
  from anthropic import AsyncAnthropic
23
- from openai import AsyncOpenAI
24
23
 
25
24
  from ..prompts.models import Message
26
25
  from .client import LLMClient
@@ -47,10 +46,6 @@ class AnthropicClient(LLMClient):
47
46
  max_retries=1,
48
47
  )
49
48
 
50
- def get_embedder(self) -> typing.Any:
51
- openai_client = AsyncOpenAI()
52
- return openai_client.embeddings
53
-
54
49
  async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
55
50
  system_message = messages[0]
56
51
  user_messages = [{'role': m.role, 'content': m.content} for m in messages[1:]] + [
@@ -55,10 +55,6 @@ class LLMClient(ABC):
55
55
  self.cache_enabled = cache
56
56
  self.cache_dir = Cache(DEFAULT_CACHE_DIR) # Create a cache directory
57
57
 
58
- @abstractmethod
59
- def get_embedder(self) -> typing.Any:
60
- pass
61
-
62
58
  @retry(
63
59
  stop=stop_after_attempt(4),
64
60
  wait=wait_random_exponential(multiplier=10, min=5, max=120),
@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
14
14
  limitations under the License.
15
15
  """
16
16
 
17
- EMBEDDING_DIM = 1024
18
17
  DEFAULT_MAX_TOKENS = 16384
19
18
  DEFAULT_TEMPERATURE = 0
20
19
 
@@ -21,7 +21,6 @@ import typing
21
21
  import groq
22
22
  from groq import AsyncGroq
23
23
  from groq.types.chat import ChatCompletionMessageParam
24
- from openai import AsyncOpenAI
25
24
 
26
25
  from ..prompts.models import Message
27
26
  from .client import LLMClient
@@ -44,10 +43,6 @@ class GroqClient(LLMClient):
44
43
 
45
44
  self.client = AsyncGroq(api_key=config.api_key)
46
45
 
47
- def get_embedder(self) -> typing.Any:
48
- openai_client = AsyncOpenAI()
49
- return openai_client.embeddings
50
-
51
46
  async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
52
47
  msgs: list[ChatCompletionMessageParam] = []
53
48
  for m in messages:
@@ -49,9 +49,6 @@ class OpenAIClient(LLMClient):
49
49
  __init__(config: LLMConfig | None = None, cache: bool = False, client: typing.Any = None):
50
50
  Initializes the OpenAIClient with the provided configuration, cache setting, and client.
51
51
 
52
- get_embedder() -> typing.Any:
53
- Returns the embedder from the OpenAI client.
54
-
55
52
  _generate_response(messages: list[Message]) -> dict[str, typing.Any]:
56
53
  Generates a response from the language model based on the provided messages.
57
54
  """
@@ -78,9 +75,6 @@ class OpenAIClient(LLMClient):
78
75
  else:
79
76
  self.client = client
80
77
 
81
- def get_embedder(self) -> typing.Any:
82
- return self.client.embeddings
83
-
84
78
  async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
85
79
  openai_messages: list[ChatCompletionMessageParam] = []
86
80
  for m in messages:
@@ -15,22 +15,18 @@ limitations under the License.
15
15
  """
16
16
 
17
17
  import logging
18
- import typing
19
18
  from time import time
20
19
 
21
- from graphiti_core.llm_client.config import EMBEDDING_DIM
20
+ from graphiti_core.embedder.client import EmbedderClient
22
21
 
23
22
  logger = logging.getLogger(__name__)
24
23
 
25
24
 
26
- async def generate_embedding(
27
- embedder: typing.Any, text: str, model: str = 'text-embedding-3-small'
28
- ):
25
+ async def generate_embedding(embedder: EmbedderClient, text: str):
29
26
  start = time()
30
27
 
31
28
  text = text.replace('\n', ' ')
32
- embedding = (await embedder.create(input=[text], model=model)).data[0].embedding
33
- embedding = embedding[:EMBEDDING_DIM]
29
+ embedding = await embedder.create(input=[text])
34
30
 
35
31
  end = time()
36
32
  logger.debug(f'embedded text of length {len(text)} in {end - start} ms')
graphiti_core/nodes.py CHANGED
@@ -25,8 +25,8 @@ from uuid import uuid4
25
25
  from neo4j import AsyncDriver
26
26
  from pydantic import BaseModel, Field
27
27
 
28
+ from graphiti_core.embedder import EmbedderClient
28
29
  from graphiti_core.errors import NodeNotFoundError
29
- from graphiti_core.llm_client.config import EMBEDDING_DIM
30
30
 
31
31
  logger = logging.getLogger(__name__)
32
32
 
@@ -212,15 +212,14 @@ class EntityNode(Node):
212
212
  name_embedding: list[float] | None = Field(default=None, description='embedding of the name')
213
213
  summary: str = Field(description='regional summary of surrounding edges', default_factory=str)
214
214
 
215
- async def generate_name_embedding(self, embedder, model='text-embedding-3-small'):
215
+ async def generate_name_embedding(self, embedder: EmbedderClient):
216
216
  start = time()
217
217
  text = self.name.replace('\n', ' ')
218
- embedding = (await embedder.create(input=[text], model=model)).data[0].embedding
219
- self.name_embedding = embedding[:EMBEDDING_DIM]
218
+ self.name_embedding = await embedder.create(input=[text])
220
219
  end = time()
221
220
  logger.info(f'embedded {text} in {end - start} ms')
222
221
 
223
- return embedding
222
+ return self.name_embedding
224
223
 
225
224
  async def save(self, driver: AsyncDriver):
226
225
  result = await driver.execute_query(
@@ -323,15 +322,14 @@ class CommunityNode(Node):
323
322
 
324
323
  return result
325
324
 
326
- async def generate_name_embedding(self, embedder, model='text-embedding-3-small'):
325
+ async def generate_name_embedding(self, embedder: EmbedderClient):
327
326
  start = time()
328
327
  text = self.name.replace('\n', ' ')
329
- embedding = (await embedder.create(input=[text], model=model)).data[0].embedding
330
- self.name_embedding = embedding[:EMBEDDING_DIM]
328
+ self.name_embedding = await embedder.create(input=[text])
331
329
  end = time()
332
330
  logger.info(f'embedded {text} in {end - start} ms')
333
331
 
334
- return embedding
332
+ return self.name_embedding
335
333
 
336
334
  @classmethod
337
335
  async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
@@ -0,0 +1,90 @@
1
+ """
2
+ Copyright 2024, Zep Software, Inc.
3
+
4
+ Licensed under the Apache License, Version 2.0 (the "License");
5
+ you may not use this file except in compliance with the License.
6
+ You may obtain a copy of the License at
7
+
8
+ http://www.apache.org/licenses/LICENSE-2.0
9
+
10
+ Unless required by applicable law or agreed to in writing, software
11
+ distributed under the License is distributed on an "AS IS" BASIS,
12
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ See the License for the specific language governing permissions and
14
+ limitations under the License.
15
+ """
16
+
17
+ import json
18
+ from typing import Any, Protocol, TypedDict
19
+
20
+ from .models import Message, PromptFunction, PromptVersion
21
+
22
+
23
+ class Prompt(Protocol):
24
+ qa_prompt: PromptVersion
25
+ eval_prompt: PromptVersion
26
+
27
+
28
+ class Versions(TypedDict):
29
+ qa_prompt: PromptFunction
30
+ eval_prompt: PromptFunction
31
+
32
+
33
+ def qa_prompt(context: dict[str, Any]) -> list[Message]:
34
+ sys_prompt = """You are Alice and should respond to all questions from the first person perspective of Alice"""
35
+
36
+ user_prompt = f"""
37
+ Your task is to briefly answer the question in the way that you think Alice would answer the question.
38
+ You are given the following entity summaries and facts to help you determine the answer to your question.
39
+ <ENTITY_SUMMARIES>
40
+ {json.dumps(context['entity_summaries'])}
41
+ </ENTITY_SUMMARIES
42
+ <FACTS>
43
+ {json.dumps(context['facts'])}
44
+ </FACTS>
45
+ <QUESTION>
46
+ {context['query']}
47
+ </QUESTION>
48
+ respond with a JSON object in the following format:
49
+ {{
50
+ "ANSWER": "how Alice would answer the question"
51
+ }}
52
+ """
53
+ return [
54
+ Message(role='system', content=sys_prompt),
55
+ Message(role='user', content=user_prompt),
56
+ ]
57
+
58
+
59
+ def eval_prompt(context: dict[str, Any]) -> list[Message]:
60
+ sys_prompt = (
61
+ """You are a judge that determines if answers to questions match a gold standard answer"""
62
+ )
63
+
64
+ user_prompt = f"""
65
+ Given the QUESTION and the gold standard ANSWER determine if the RESPONSE to the question is correct or incorrect.
66
+ Although the RESPONSE may be more verbose, mark it as correct as long as it references the same topic
67
+ as the gold standard ANSWER. Also include your reasoning for the grade.
68
+ <QUESTION>
69
+ {context['query']}
70
+ </QUESTION>
71
+ <ANSWER>
72
+ {context['answer']}
73
+ </ANSWER>
74
+ <RESPONSE>
75
+ {context['response']}
76
+ </RESPONSE>
77
+
78
+ respond with a JSON object in the following format:
79
+ {{
80
+ "is_correct": "boolean if the answer is correct or incorrect"
81
+ "reasoning": "why you determined the response was correct or incorrect"
82
+ }}
83
+ """
84
+ return [
85
+ Message(role='system', content=sys_prompt),
86
+ Message(role='user', content=user_prompt),
87
+ ]
88
+
89
+
90
+ versions: Versions = {'qa_prompt': qa_prompt, 'eval_prompt': eval_prompt}
@@ -34,6 +34,9 @@ from .dedupe_nodes import (
34
34
  from .dedupe_nodes import (
35
35
  versions as dedupe_nodes_versions,
36
36
  )
37
+ from .eval import Prompt as EvalPrompt
38
+ from .eval import Versions as EvalVersions
39
+ from .eval import versions as eval_versions
37
40
  from .extract_edge_dates import (
38
41
  Prompt as ExtractEdgeDatesPrompt,
39
42
  )
@@ -84,6 +87,7 @@ class PromptLibrary(Protocol):
84
87
  invalidate_edges: InvalidateEdgesPrompt
85
88
  extract_edge_dates: ExtractEdgeDatesPrompt
86
89
  summarize_nodes: SummarizeNodesPrompt
90
+ eval: EvalPrompt
87
91
 
88
92
 
89
93
  class PromptLibraryImpl(TypedDict):
@@ -94,6 +98,7 @@ class PromptLibraryImpl(TypedDict):
94
98
  invalidate_edges: InvalidateEdgesVersions
95
99
  extract_edge_dates: ExtractEdgeDatesVersions
96
100
  summarize_nodes: SummarizeNodesVersions
101
+ eval: EvalVersions
97
102
 
98
103
 
99
104
  class VersionWrapper:
@@ -124,5 +129,6 @@ PROMPT_LIBRARY_IMPL: PromptLibraryImpl = {
124
129
  'invalidate_edges': invalidate_edges_versions,
125
130
  'extract_edge_dates': extract_edge_dates_versions,
126
131
  'summarize_nodes': summarize_nodes_versions,
132
+ 'eval': eval_versions,
127
133
  }
128
134
  prompt_library: PromptLibrary = PromptLibraryWrapper(PROMPT_LIBRARY_IMPL) # type: ignore[assignment]
@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
14
14
  limitations under the License.
15
15
  """
16
16
 
17
+ import asyncio
17
18
  import logging
18
19
  from collections import defaultdict
19
20
  from time import time
@@ -21,8 +22,8 @@ from time import time
21
22
  from neo4j import AsyncDriver
22
23
 
23
24
  from graphiti_core.edges import EntityEdge
25
+ from graphiti_core.embedder import EmbedderClient
24
26
  from graphiti_core.errors import SearchRerankerError
25
- from graphiti_core.llm_client.config import EMBEDDING_DIM
26
27
  from graphiti_core.nodes import CommunityNode, EntityNode
27
28
  from graphiti_core.search.search_config import (
28
29
  DEFAULT_SEARCH_LIMIT,
@@ -55,7 +56,7 @@ logger = logging.getLogger(__name__)
55
56
 
56
57
  async def search(
57
58
  driver: AsyncDriver,
58
- embedder,
59
+ embedder: EmbedderClient,
59
60
  query: str,
60
61
  group_ids: list[str] | None,
61
62
  config: SearchConfig,
@@ -65,32 +66,39 @@ async def search(
65
66
  query = query.replace('\n', ' ')
66
67
  # if group_ids is empty, set it to None
67
68
  group_ids = group_ids if group_ids else None
68
- edges = (
69
- await edge_search(
70
- driver, embedder, query, group_ids, config.edge_config, center_node_uuid, config.limit
71
- )
72
- if config.edge_config is not None
73
- else []
74
- )
75
- nodes = (
76
- await node_search(
77
- driver, embedder, query, group_ids, config.node_config, center_node_uuid, config.limit
78
- )
79
- if config.node_config is not None
80
- else []
81
- )
82
- communities = (
83
- await community_search(
84
- driver, embedder, query, group_ids, config.community_config, config.limit
85
- )
86
- if config.community_config is not None
87
- else []
69
+ edges, nodes, communities = await asyncio.gather(
70
+ edge_search(
71
+ driver,
72
+ embedder,
73
+ query,
74
+ group_ids,
75
+ config.edge_config,
76
+ center_node_uuid,
77
+ config.limit,
78
+ ),
79
+ node_search(
80
+ driver,
81
+ embedder,
82
+ query,
83
+ group_ids,
84
+ config.node_config,
85
+ center_node_uuid,
86
+ config.limit,
87
+ ),
88
+ community_search(
89
+ driver,
90
+ embedder,
91
+ query,
92
+ group_ids,
93
+ config.community_config,
94
+ config.limit,
95
+ ),
88
96
  )
89
97
 
90
98
  results = SearchResults(
91
- edges=edges[: config.limit],
92
- nodes=nodes[: config.limit],
93
- communities=communities[: config.limit],
99
+ edges=edges,
100
+ nodes=nodes,
101
+ communities=communities,
94
102
  )
95
103
 
96
104
  end = time()
@@ -102,13 +110,16 @@ async def search(
102
110
 
103
111
  async def edge_search(
104
112
  driver: AsyncDriver,
105
- embedder,
113
+ embedder: EmbedderClient,
106
114
  query: str,
107
115
  group_ids: list[str] | None,
108
- config: EdgeSearchConfig,
116
+ config: EdgeSearchConfig | None,
109
117
  center_node_uuid: str | None = None,
110
118
  limit=DEFAULT_SEARCH_LIMIT,
111
119
  ) -> list[EntityEdge]:
120
+ if config is None:
121
+ return []
122
+
112
123
  search_results: list[list[EntityEdge]] = []
113
124
 
114
125
  if EdgeSearchMethod.bm25 in config.search_methods:
@@ -116,11 +127,7 @@ async def edge_search(
116
127
  search_results.append(text_search)
117
128
 
118
129
  if EdgeSearchMethod.cosine_similarity in config.search_methods:
119
- search_vector = (
120
- (await embedder.create(input=[query], model='text-embedding-3-small'))
121
- .data[0]
122
- .embedding[:EMBEDDING_DIM]
123
- )
130
+ search_vector = await embedder.create(input=[query])
124
131
 
125
132
  similarity_search = await edge_similarity_search(
126
133
  driver, search_vector, None, None, group_ids, 2 * limit
@@ -162,18 +169,21 @@ async def edge_search(
162
169
  if config.reranker == EdgeReranker.episode_mentions:
163
170
  reranked_edges.sort(reverse=True, key=lambda edge: len(edge.episodes))
164
171
 
165
- return reranked_edges
172
+ return reranked_edges[:limit]
166
173
 
167
174
 
168
175
  async def node_search(
169
176
  driver: AsyncDriver,
170
- embedder,
177
+ embedder: EmbedderClient,
171
178
  query: str,
172
179
  group_ids: list[str] | None,
173
- config: NodeSearchConfig,
180
+ config: NodeSearchConfig | None,
174
181
  center_node_uuid: str | None = None,
175
182
  limit=DEFAULT_SEARCH_LIMIT,
176
183
  ) -> list[EntityNode]:
184
+ if config is None:
185
+ return []
186
+
177
187
  search_results: list[list[EntityNode]] = []
178
188
 
179
189
  if NodeSearchMethod.bm25 in config.search_methods:
@@ -181,11 +191,7 @@ async def node_search(
181
191
  search_results.append(text_search)
182
192
 
183
193
  if NodeSearchMethod.cosine_similarity in config.search_methods:
184
- search_vector = (
185
- (await embedder.create(input=[query], model='text-embedding-3-small'))
186
- .data[0]
187
- .embedding[:EMBEDDING_DIM]
188
- )
194
+ search_vector = await embedder.create(input=[query])
189
195
 
190
196
  similarity_search = await node_similarity_search(
191
197
  driver, search_vector, group_ids, 2 * limit
@@ -212,17 +218,20 @@ async def node_search(
212
218
 
213
219
  reranked_nodes = [node_uuid_map[uuid] for uuid in reranked_uuids]
214
220
 
215
- return reranked_nodes
221
+ return reranked_nodes[:limit]
216
222
 
217
223
 
218
224
  async def community_search(
219
225
  driver: AsyncDriver,
220
- embedder,
226
+ embedder: EmbedderClient,
221
227
  query: str,
222
228
  group_ids: list[str] | None,
223
- config: CommunitySearchConfig,
229
+ config: CommunitySearchConfig | None,
224
230
  limit=DEFAULT_SEARCH_LIMIT,
225
231
  ) -> list[CommunityNode]:
232
+ if config is None:
233
+ return []
234
+
226
235
  search_results: list[list[CommunityNode]] = []
227
236
 
228
237
  if CommunitySearchMethod.bm25 in config.search_methods:
@@ -230,11 +239,7 @@ async def community_search(
230
239
  search_results.append(text_search)
231
240
 
232
241
  if CommunitySearchMethod.cosine_similarity in config.search_methods:
233
- search_vector = (
234
- (await embedder.create(input=[query], model='text-embedding-3-small'))
235
- .data[0]
236
- .embedding[:EMBEDDING_DIM]
237
- )
242
+ search_vector = await embedder.create(input=[query])
238
243
 
239
244
  similarity_search = await community_similarity_search(
240
245
  driver, search_vector, group_ids, 2 * limit
@@ -255,4 +260,4 @@ async def community_search(
255
260
 
256
261
  reranked_communities = [community_uuid_map[uuid] for uuid in reranked_uuids]
257
262
 
258
- return reranked_communities
263
+ return reranked_communities[:limit]
@@ -16,13 +16,13 @@ limitations under the License.
16
16
 
17
17
  import asyncio
18
18
  import logging
19
- import re
20
19
  from collections import defaultdict
21
20
  from time import time
22
21
 
23
22
  from neo4j import AsyncDriver, Query
24
23
 
25
24
  from graphiti_core.edges import EntityEdge, get_entity_edge_from_record
25
+ from graphiti_core.helpers import lucene_sanitize
26
26
  from graphiti_core.nodes import (
27
27
  CommunityNode,
28
28
  EntityNode,
@@ -36,6 +36,22 @@ logger = logging.getLogger(__name__)
36
36
  RELEVANT_SCHEMA_LIMIT = 3
37
37
 
38
38
 
39
+ def fulltext_query(query: str, group_ids: list[str] | None = None):
40
+ group_ids_filter_list = (
41
+ [f'group_id:"{lucene_sanitize(g)}"' for g in group_ids] if group_ids is not None else []
42
+ )
43
+ group_ids_filter = ''
44
+ for f in group_ids_filter_list:
45
+ group_ids_filter += f if not group_ids_filter else f'OR {f}'
46
+
47
+ group_ids_filter += ' AND ' if group_ids_filter else ''
48
+
49
+ fuzzy_query = lucene_sanitize(query) + '~'
50
+ full_query = group_ids_filter + fuzzy_query
51
+
52
+ return full_query
53
+
54
+
39
55
  async def get_mentioned_nodes(
40
56
  driver: AsyncDriver, episodes: list[EpisodicNode]
41
57
  ) -> list[EntityNode]:
@@ -91,11 +107,15 @@ async def edge_fulltext_search(
91
107
  limit=RELEVANT_SCHEMA_LIMIT,
92
108
  ) -> list[EntityEdge]:
93
109
  # fulltext search over facts
110
+ fuzzy_query = fulltext_query(query, group_ids)
111
+
94
112
  cypher_query = Query("""
95
- CALL db.index.fulltext.queryRelationships("name_and_fact", $query)
113
+ CALL db.index.fulltext.queryRelationships("edge_name_and_fact", $query)
96
114
  YIELD relationship AS rel, score
97
- MATCH (n:Entity {uuid: $source_uuid})-[r {uuid: rel.uuid}]-(m:Entity {uuid: $target_uuid})
98
- WHERE $group_ids IS NULL OR n.group_id IN $group_ids
115
+ MATCH (n:Entity)-[r {uuid: rel.uuid}]-(m:Entity)
116
+ WHERE ($source_uuid IS NULL OR n.uuid = $source_uuid)
117
+ AND ($target_uuid IS NULL OR m.uuid = $target_uuid)
118
+ AND ($group_ids IS NULL OR n.group_id IN $group_ids)
99
119
  RETURN
100
120
  r.uuid AS uuid,
101
121
  r.group_id AS group_id,
@@ -112,72 +132,6 @@ async def edge_fulltext_search(
112
132
  ORDER BY score DESC LIMIT $limit
113
133
  """)
114
134
 
115
- if source_node_uuid is None and target_node_uuid is None:
116
- cypher_query = Query("""
117
- CALL db.index.fulltext.queryRelationships("name_and_fact", $query)
118
- YIELD relationship AS rel, score
119
- MATCH (n:Entity)-[r {uuid: rel.uuid}]-(m:Entity)
120
- WHERE $group_ids IS NULL OR r.group_id IN $group_ids
121
- RETURN
122
- r.uuid AS uuid,
123
- r.group_id AS group_id,
124
- n.uuid AS source_node_uuid,
125
- m.uuid AS target_node_uuid,
126
- r.created_at AS created_at,
127
- r.name AS name,
128
- r.fact AS fact,
129
- r.fact_embedding AS fact_embedding,
130
- r.episodes AS episodes,
131
- r.expired_at AS expired_at,
132
- r.valid_at AS valid_at,
133
- r.invalid_at AS invalid_at
134
- ORDER BY score DESC LIMIT $limit
135
- """)
136
- elif source_node_uuid is None:
137
- cypher_query = Query("""
138
- CALL db.index.fulltext.queryRelationships("name_and_fact", $query)
139
- YIELD relationship AS rel, score
140
- MATCH (n:Entity)-[r {uuid: rel.uuid}]-(m:Entity {uuid: $target_uuid})
141
- WHERE $group_ids IS NULL OR r.group_id IN $group_ids
142
- RETURN
143
- r.uuid AS uuid,
144
- r.group_id AS group_id,
145
- n.uuid AS source_node_uuid,
146
- m.uuid AS target_node_uuid,
147
- r.created_at AS created_at,
148
- r.name AS name,
149
- r.fact AS fact,
150
- r.fact_embedding AS fact_embedding,
151
- r.episodes AS episodes,
152
- r.expired_at AS expired_at,
153
- r.valid_at AS valid_at,
154
- r.invalid_at AS invalid_at
155
- ORDER BY score DESC LIMIT $limit
156
- """)
157
- elif target_node_uuid is None:
158
- cypher_query = Query("""
159
- CALL db.index.fulltext.queryRelationships("name_and_fact", $query)
160
- YIELD relationship AS rel, score
161
- MATCH (n:Entity {uuid: $source_uuid})-[r {uuid: rel.uuid}]-(m:Entity)
162
- WHERE $group_ids IS NULL OR r.group_id IN $group_ids
163
- RETURN
164
- r.uuid AS uuid,
165
- r.group_id AS group_id,
166
- n.uuid AS source_node_uuid,
167
- m.uuid AS target_node_uuid,
168
- r.created_at AS created_at,
169
- r.name AS name,
170
- r.fact AS fact,
171
- r.fact_embedding AS fact_embedding,
172
- r.episodes AS episodes,
173
- r.expired_at AS expired_at,
174
- r.valid_at AS valid_at,
175
- r.invalid_at AS invalid_at
176
- ORDER BY score DESC LIMIT $limit
177
- """)
178
-
179
- fuzzy_query = re.sub(r'[^\w\s]', '', query) + '~'
180
-
181
135
  records, _, _ = await driver.execute_query(
182
136
  cypher_query,
183
137
  query=fuzzy_query,
@@ -202,11 +156,12 @@ async def edge_similarity_search(
202
156
  ) -> list[EntityEdge]:
203
157
  # vector similarity search over embedded facts
204
158
  query = Query("""
205
- CALL db.index.vector.queryRelationships("fact_embedding", $limit, $search_vector)
206
- YIELD relationship AS rel, score
207
- MATCH (n:Entity {uuid: $source_uuid})-[r {uuid: rel.uuid}]-(m:Entity {uuid: $target_uuid})
208
- WHERE $group_ids IS NULL OR r.group_id IN $group_ids
159
+ MATCH (n:Entity)-[r:RELATES_TO]-(m:Entity)
160
+ WHERE ($group_ids IS NULL OR r.group_id IN $group_ids)
161
+ AND ($source_uuid IS NULL OR n.uuid = $source_uuid)
162
+ AND ($target_uuid IS NULL OR m.uuid = $target_uuid)
209
163
  RETURN
164
+ vector.similarity.cosine(r.fact_embedding, $search_vector) AS score,
210
165
  r.uuid AS uuid,
211
166
  r.group_id AS group_id,
212
167
  n.uuid AS source_node_uuid,
@@ -220,72 +175,9 @@ async def edge_similarity_search(
220
175
  r.valid_at AS valid_at,
221
176
  r.invalid_at AS invalid_at
222
177
  ORDER BY score DESC
178
+ LIMIT $limit
223
179
  """)
224
180
 
225
- if source_node_uuid is None and target_node_uuid is None:
226
- query = Query("""
227
- CALL db.index.vector.queryRelationships("fact_embedding", $limit, $search_vector)
228
- YIELD relationship AS rel, score
229
- MATCH (n:Entity)-[r {uuid: rel.uuid}]-(m:Entity)
230
- WHERE $group_ids IS NULL OR r.group_id IN $group_ids
231
- RETURN
232
- r.uuid AS uuid,
233
- r.group_id AS group_id,
234
- n.uuid AS source_node_uuid,
235
- m.uuid AS target_node_uuid,
236
- r.created_at AS created_at,
237
- r.name AS name,
238
- r.fact AS fact,
239
- r.fact_embedding AS fact_embedding,
240
- r.episodes AS episodes,
241
- r.expired_at AS expired_at,
242
- r.valid_at AS valid_at,
243
- r.invalid_at AS invalid_at
244
- ORDER BY score DESC
245
- """)
246
- elif source_node_uuid is None:
247
- query = Query("""
248
- CALL db.index.vector.queryRelationships("fact_embedding", $limit, $search_vector)
249
- YIELD relationship AS rel, score
250
- MATCH (n:Entity)-[r {uuid: rel.uuid}]-(m:Entity {uuid: $target_uuid})
251
- WHERE $group_ids IS NULL OR r.group_id IN $group_ids
252
- RETURN
253
- r.uuid AS uuid,
254
- r.group_id AS group_id,
255
- n.uuid AS source_node_uuid,
256
- m.uuid AS target_node_uuid,
257
- r.created_at AS created_at,
258
- r.name AS name,
259
- r.fact AS fact,
260
- r.fact_embedding AS fact_embedding,
261
- r.episodes AS episodes,
262
- r.expired_at AS expired_at,
263
- r.valid_at AS valid_at,
264
- r.invalid_at AS invalid_at
265
- ORDER BY score DESC
266
- """)
267
- elif target_node_uuid is None:
268
- query = Query("""
269
- CALL db.index.vector.queryRelationships("fact_embedding", $limit, $search_vector)
270
- YIELD relationship AS rel, score
271
- MATCH (n:Entity {uuid: $source_uuid})-[r {uuid: rel.uuid}]-(m:Entity)
272
- WHERE $group_ids IS NULL OR r.group_id IN $group_ids
273
- RETURN
274
- r.uuid AS uuid,
275
- r.group_id AS group_id,
276
- n.uuid AS source_node_uuid,
277
- m.uuid AS target_node_uuid,
278
- r.created_at AS created_at,
279
- r.name AS name,
280
- r.fact AS fact,
281
- r.fact_embedding AS fact_embedding,
282
- r.episodes AS episodes,
283
- r.expired_at AS expired_at,
284
- r.valid_at AS valid_at,
285
- r.invalid_at AS invalid_at
286
- ORDER BY score DESC
287
- """)
288
-
289
181
  records, _, _ = await driver.execute_query(
290
182
  query,
291
183
  search_vector=search_vector,
@@ -307,10 +199,11 @@ async def node_fulltext_search(
307
199
  limit=RELEVANT_SCHEMA_LIMIT,
308
200
  ) -> list[EntityNode]:
309
201
  # BM25 search to get top nodes
310
- fuzzy_query = re.sub(r'[^\w\s]', '', query) + '~'
202
+ fuzzy_query = fulltext_query(query, group_ids)
203
+
311
204
  records, _, _ = await driver.execute_query(
312
205
  """
313
- CALL db.index.fulltext.queryNodes("name_and_summary", $query)
206
+ CALL db.index.fulltext.queryNodes("node_name_and_summary", $query)
314
207
  YIELD node AS n, score
315
208
  WHERE $group_ids IS NULL OR n.group_id IN $group_ids
316
209
  RETURN
@@ -341,11 +234,10 @@ async def node_similarity_search(
341
234
  # vector similarity search over entity names
342
235
  records, _, _ = await driver.execute_query(
343
236
  """
344
- CALL db.index.vector.queryNodes("name_embedding", $limit, $search_vector)
345
- YIELD node AS n, score
346
237
  MATCH (n:Entity)
347
238
  WHERE $group_ids IS NULL OR n.group_id IN $group_ids
348
239
  RETURN
240
+ vector.similarity.cosine(n.name_embedding, $search_vector) AS score,
349
241
  n.uuid As uuid,
350
242
  n.group_id AS group_id,
351
243
  n.name AS name,
@@ -353,6 +245,7 @@ async def node_similarity_search(
353
245
  n.created_at AS created_at,
354
246
  n.summary AS summary
355
247
  ORDER BY score DESC
248
+ LIMIT $limit
356
249
  """,
357
250
  search_vector=search_vector,
358
251
  group_ids=group_ids,
@@ -370,7 +263,8 @@ async def community_fulltext_search(
370
263
  limit=RELEVANT_SCHEMA_LIMIT,
371
264
  ) -> list[CommunityNode]:
372
265
  # BM25 search to get top communities
373
- fuzzy_query = re.sub(r'[^\w\s]', '', query) + '~'
266
+ fuzzy_query = fulltext_query(query, group_ids)
267
+
374
268
  records, _, _ = await driver.execute_query(
375
269
  """
376
270
  CALL db.index.fulltext.queryNodes("community_name", $query)
@@ -405,11 +299,10 @@ async def community_similarity_search(
405
299
  # vector similarity search over entity names
406
300
  records, _, _ = await driver.execute_query(
407
301
  """
408
- CALL db.index.vector.queryNodes("community_name_embedding", $limit, $search_vector)
409
- YIELD node AS comm, score
410
302
  MATCH (comm:Community)
411
- WHERE $group_ids IS NULL OR comm.group_id IN $group_ids
303
+ WHERE ($group_ids IS NULL OR comm.group_id IN $group_ids)
412
304
  RETURN
305
+ vector.similarity.cosine(comm.name_embedding, $search_vector) AS score,
413
306
  comm.uuid As uuid,
414
307
  comm.group_id AS group_id,
415
308
  comm.name AS name,
@@ -417,6 +310,7 @@ async def community_similarity_search(
417
310
  comm.created_at AS created_at,
418
311
  comm.summary AS summary
419
312
  ORDER BY score DESC
313
+ LIMIT $limit
420
314
  """,
421
315
  search_vector=search_vector,
422
316
  group_ids=group_ids,
@@ -7,6 +7,7 @@ from neo4j import AsyncDriver
7
7
  from pydantic import BaseModel
8
8
 
9
9
  from graphiti_core.edges import CommunityEdge
10
+ from graphiti_core.embedder import EmbedderClient
10
11
  from graphiti_core.llm_client import LLMClient
11
12
  from graphiti_core.nodes import CommunityNode, EntityNode, get_community_node_from_record
12
13
  from graphiti_core.prompts import prompt_library
@@ -288,7 +289,7 @@ async def determine_entity_community(
288
289
 
289
290
 
290
291
  async def update_community(
291
- driver: AsyncDriver, llm_client: LLMClient, embedder, entity: EntityNode
292
+ driver: AsyncDriver, llm_client: LLMClient, embedder: EmbedderClient, entity: EntityNode
292
293
  ):
293
294
  community, is_new = await determine_entity_community(driver, entity)
294
295
 
@@ -28,7 +28,16 @@ EPISODE_WINDOW_LEN = 3
28
28
  logger = logging.getLogger(__name__)
29
29
 
30
30
 
31
- async def build_indices_and_constraints(driver: AsyncDriver):
31
+ async def build_indices_and_constraints(driver: AsyncDriver, delete_existing: bool = False):
32
+ if delete_existing:
33
+ records, _, _ = await driver.execute_query("""
34
+ SHOW INDEXES YIELD name
35
+ """)
36
+ index_names = [record['name'] for record in records]
37
+ await asyncio.gather(
38
+ *[driver.execute_query("""DROP INDEX $name""", name=name) for name in index_names]
39
+ )
40
+
32
41
  range_indices: list[LiteralString] = [
33
42
  'CREATE INDEX entity_uuid IF NOT EXISTS FOR (n:Entity) ON (n.uuid)',
34
43
  'CREATE INDEX episode_uuid IF NOT EXISTS FOR (n:Episodic) ON (n.uuid)',
@@ -52,38 +61,15 @@ async def build_indices_and_constraints(driver: AsyncDriver):
52
61
  ]
53
62
 
54
63
  fulltext_indices: list[LiteralString] = [
55
- 'CREATE FULLTEXT INDEX name_and_summary IF NOT EXISTS FOR (n:Entity) ON EACH [n.name, n.summary]',
56
- 'CREATE FULLTEXT INDEX community_name IF NOT EXISTS FOR (n:Community) ON EACH [n.name]',
57
- 'CREATE FULLTEXT INDEX name_and_fact IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON EACH [e.name, e.fact]',
64
+ """CREATE FULLTEXT INDEX node_name_and_summary IF NOT EXISTS
65
+ FOR (n:Entity) ON EACH [n.name, n.summary, n.group_id]""",
66
+ """CREATE FULLTEXT INDEX community_name IF NOT EXISTS
67
+ FOR (n:Community) ON EACH [n.name, n.group_id]""",
68
+ """CREATE FULLTEXT INDEX edge_name_and_fact IF NOT EXISTS
69
+ FOR ()-[e:RELATES_TO]-() ON EACH [e.name, e.fact, e.group_id]""",
58
70
  ]
59
71
 
60
- vector_indices: list[LiteralString] = [
61
- """
62
- CREATE VECTOR INDEX fact_embedding IF NOT EXISTS
63
- FOR ()-[r:RELATES_TO]-() ON (r.fact_embedding)
64
- OPTIONS {indexConfig: {
65
- `vector.dimensions`: 1024,
66
- `vector.similarity_function`: 'cosine'
67
- }}
68
- """,
69
- """
70
- CREATE VECTOR INDEX name_embedding IF NOT EXISTS
71
- FOR (n:Entity) ON (n.name_embedding)
72
- OPTIONS {indexConfig: {
73
- `vector.dimensions`: 1024,
74
- `vector.similarity_function`: 'cosine'
75
- }}
76
- """,
77
- """
78
- CREATE VECTOR INDEX community_name_embedding IF NOT EXISTS
79
- FOR (n:Community) ON (n.name_embedding)
80
- OPTIONS {indexConfig: {
81
- `vector.dimensions`: 1024,
82
- `vector.similarity_function`: 'cosine'
83
- }}
84
- """,
85
- ]
86
- index_queries: list[LiteralString] = range_indices + fulltext_indices + vector_indices
72
+ index_queries: list[LiteralString] = range_indices + fulltext_indices
87
73
 
88
74
  await asyncio.gather(*[driver.execute_query(query) for query in index_queries])
89
75
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: graphiti-core
3
- Version: 0.3.6
3
+ Version: 0.3.8
4
4
  Summary: A temporal graph building library
5
5
  License: Apache-2.0
6
6
  Author: Paul Paliychuk
@@ -14,7 +14,7 @@ Classifier: Programming Language :: Python :: 3.12
14
14
  Requires-Dist: diskcache (>=5.6.3,<6.0.0)
15
15
  Requires-Dist: neo4j (>=5.23.0,<6.0.0)
16
16
  Requires-Dist: numpy (>=1.0.0)
17
- Requires-Dist: openai (>=1.38.0,<2.0.0)
17
+ Requires-Dist: openai (>=1.50.2,<2.0.0)
18
18
  Requires-Dist: pydantic (>=2.8.2,<3.0.0)
19
19
  Requires-Dist: tenacity (<9.0.0)
20
20
  Description-Content-Type: text/markdown
@@ -1,43 +1,48 @@
1
1
  graphiti_core/__init__.py,sha256=e5SWFkRiaUwfprYIeIgVIh7JDedNiloZvd3roU-0aDY,55
2
- graphiti_core/edges.py,sha256=ePQfFHYXwYXqcTH_im9OjU74OH78f0fhCBo08aEXLoo,13471
2
+ graphiti_core/edges.py,sha256=lLuRKjSHTk1GvTS06OUw2lSMiDAB4TQSXgnLq1fU3n8,13378
3
+ graphiti_core/embedder/__init__.py,sha256=eWd-0sPxflnYXLoWNT9sxwCIFun5JNO9Fk4E-ZXXf8Y,164
4
+ graphiti_core/embedder/client.py,sha256=Sd9CyYXaqRazdOH8opKackrTx-y9y-T54M78XTVMzxs,1006
5
+ graphiti_core/embedder/openai.py,sha256=28cl4qQCQeu6EGxVVPw3lPesA-Z_Cpvuhozyc1jdqVg,1586
6
+ graphiti_core/embedder/voyage.py,sha256=pGrSquGnSiYl4nXGnutbdWchtYgZb0Fi_yW3c90dPlI,1497
3
7
  graphiti_core/errors.py,sha256=iJrkk5sTgc2z16ABS6TziPylEabdBJcpk0x9KyBUmxs,1527
4
- graphiti_core/graphiti.py,sha256=z6a4tCyDID_o6gloXDuUFmbL22bRiUE7A22JPYGVIyI,25947
5
- graphiti_core/helpers.py,sha256=qQqZJBkc_z5f3x5axPfCKK_QHLRybvWNFb57WXNENfQ,769
8
+ graphiti_core/graphiti.py,sha256=5E2UbYlbl65D3MZyagEUPgoPrb_kVYDIqIw7KVlU_NM,26162
9
+ graphiti_core/helpers.py,sha256=_wTSDcYmeXT3u0AwX15iSLuTRa_SR4jJdT10rxfl1_E,1484
6
10
  graphiti_core/llm_client/__init__.py,sha256=PA80TSMeX-sUXITXEAxMDEt3gtfZgcJrGJUcyds1mSo,207
7
- graphiti_core/llm_client/anthropic_client.py,sha256=16cWm_fQSUvJTSmgISBcNF8vytIhk-c5mmMK0Xd7SPE,2557
8
- graphiti_core/llm_client/client.py,sha256=g3vEBNV0E24HdKR3DmqjY8cqqr1CDlvrdh7SaiCUkDc,3470
9
- graphiti_core/llm_client/config.py,sha256=YIuR5XTINvxsEGDcpPXCqDWfWXGHTB4GB0k5DSRD7Rg,2360
11
+ graphiti_core/llm_client/anthropic_client.py,sha256=4l2PbCjIoeRr7UJ2DUh2grYLTtE2vNaWlo72IIRQDeI,2405
12
+ graphiti_core/llm_client/client.py,sha256=WAnX0e4EuCFHXdFHeq_O1HZsW1STSByvDCFUHMAHEFU,3394
13
+ graphiti_core/llm_client/config.py,sha256=VwtvD0B7TNqE6Cl-rvH5v-bAfmjMLhEUuFmHSPt10EI,2339
10
14
  graphiti_core/llm_client/errors.py,sha256=-qlWwv1X-UjfsFIiNl-7yJIYvPwi7z8srVRfX4-s6uk,814
11
- graphiti_core/llm_client/groq_client.py,sha256=j467rL2tNaKplpTOP9pZNUCxG3rrHAEE26CBDk24jzw,2481
12
- graphiti_core/llm_client/openai_client.py,sha256=LlvhAI5nfrLUWehQ0TSPkeqzgV4wJMDiK56XUQxR21A,4002
13
- graphiti_core/llm_client/utils.py,sha256=0KT4XxTVw3c0__HLDj3F8kNR4K_qY0hT0TH-pQZ_IZw,1126
14
- graphiti_core/nodes.py,sha256=4RdkzvaiqEIqppkaYq53JLZ3tr6AskeKVOHOCJLR2BA,13847
15
+ graphiti_core/llm_client/groq_client.py,sha256=5uGWeQ903EuNxuRiaeH-_J1U2Le_b7Q1UGV_K8bQAiw,2329
16
+ graphiti_core/llm_client/openai_client.py,sha256=xLkbpusRVFRK0zPr3kOqY31HK_XCXrpO5rqUSpcEqEU,3825
17
+ graphiti_core/llm_client/utils.py,sha256=Ms-QhA5X9rps7NBdJeQZUgQLD3vaZRWPiTlhJa6BjXM,995
18
+ graphiti_core/nodes.py,sha256=wIYeRspoRErcX0vvesk_fxhdXKCYn4rpgjgm3PdwSkI,13669
15
19
  graphiti_core/prompts/__init__.py,sha256=EA-x9xUki9l8wnu2l8ek_oNf75-do5tq5hVq7Zbv8Kw,101
16
20
  graphiti_core/prompts/dedupe_edges.py,sha256=DUNHdIudj50FAjkla4nc68tSFSD2yjmYHBw-Bb7ph20,6529
17
21
  graphiti_core/prompts/dedupe_nodes.py,sha256=BZ9S-PB9SSGjc5Oo8ivdgA6rZx3OGOFhKtwrBlQ0bm0,7269
22
+ graphiti_core/prompts/eval.py,sha256=fYLY2nKwgE9dB7mtYMNKyn1tQXM8B-tOeYmSzB5Bxk8,2844
18
23
  graphiti_core/prompts/extract_edge_dates.py,sha256=oOCR8mC_3gI1bumrmIjUbkNO-WTuLTXXAalPDYnDXeM,3655
19
24
  graphiti_core/prompts/extract_edges.py,sha256=AQ8xYbAv_RKXAT6WMwXs1_GvUdLtM_lhLNbt3SkOAmk,5348
20
25
  graphiti_core/prompts/extract_nodes.py,sha256=VIr0Nh0mSiodI3iGOQFszh7DOni4mufOKJDuGkMysl8,6889
21
26
  graphiti_core/prompts/invalidate_edges.py,sha256=8SHt3iPTdmqk8A52LxgdMtI39w4USKqVDMOS2i6lRQ4,4342
22
- graphiti_core/prompts/lib.py,sha256=lIgVAxu4U4R9gnJVPkqxT4hcAXfErHECtM_Uceh55VA,3857
27
+ graphiti_core/prompts/lib.py,sha256=ZOE6nNoI_wQ12Sufx7rQkQtkIm_eTAL7pCiYGU2hcMI,4054
23
28
  graphiti_core/prompts/models.py,sha256=cvx_Bv5RMFUD_5IUawYrbpOKLPHogai7_bm7YXrSz84,867
24
29
  graphiti_core/prompts/summarize_nodes.py,sha256=FLuZpGTABgcxuIDkx_IKH115nHEw0rIaFhcGlWveAMc,2357
25
30
  graphiti_core/py.typed,sha256=vlmmzQOt7bmeQl9L3XJP4W6Ry0iiELepnOrinKz5KQg,79
26
31
  graphiti_core/search/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
27
- graphiti_core/search/search.py,sha256=QDgxWihEQuQ9H9BNGfWOEEDH0aGXUtwfpb7inhN88fk,8773
32
+ graphiti_core/search/search.py,sha256=odxpm6MJw5ihEDjbBQ2Icvtr5Mf2oG8Yj6LpNqO3gFw,8620
28
33
  graphiti_core/search/search_config.py,sha256=d8w9RDO55G2bwbjYQBaD6gXqEWK1-NsDANrNibYB6t8,2165
29
34
  graphiti_core/search/search_config_recipes.py,sha256=_VJqvYB70e8Jke3hsbeQF3Bdogn2MubpYeAQe15M2Jo,3450
30
- graphiti_core/search/search_utils.py,sha256=LKFskMPqgRoxFpn5cdNYjAGTMy6z-FybNRBhNVtpJZM,23497
35
+ graphiti_core/search/search_utils.py,sha256=WE-iVPI92AWR13aM3JQxtHaYoiPzDMtOOo8rEob8QEI,17844
31
36
  graphiti_core/utils/__init__.py,sha256=cJAcMnBZdHBQmWrZdU1PQ1YmaL75bhVUkyVpIPuOyns,260
32
37
  graphiti_core/utils/bulk_utils.py,sha256=JtoYTZPCigPa3n2E43Oe7QhFZRTA_QKNGy1jVgklHag,12614
33
38
  graphiti_core/utils/maintenance/__init__.py,sha256=4b9sfxqyFZMLwxxS2lnQ6_wBr3xrJRIqfAWOidK8EK0,388
34
- graphiti_core/utils/maintenance/community_operations.py,sha256=2jtA0ZwjwZyDiC1Es8d4p0KafT98w1fSkQOvi1IdT80,9809
39
+ graphiti_core/utils/maintenance/community_operations.py,sha256=Z2lVrTmUh42sEPqSDZq4fXbcj507BuZrHZKV1vJk6tU,9875
35
40
  graphiti_core/utils/maintenance/edge_operations.py,sha256=lSeesSnWQ3vpeD2dIY0tSiHEHRMK6fiirEhNNT-s5os,11438
36
- graphiti_core/utils/maintenance/graph_data_operations.py,sha256=zk-Ir7msJIbdQj-8KTl0As9a8zYUG-e-dVdbVacxlf8,6515
41
+ graphiti_core/utils/maintenance/graph_data_operations.py,sha256=RgdqYSau9Mr-f7IUSD1sSPztxlyO0C80C3MPPmPBRi0,6100
37
42
  graphiti_core/utils/maintenance/node_operations.py,sha256=QAg4KQkSAOXx9QRaUp7t6DCaztZlzeOBC3__57FCs_o,9025
38
43
  graphiti_core/utils/maintenance/temporal_operations.py,sha256=BzfGDm96w4HcUEsaWTHUBt5S8dNmDQL1eX6AuBL-XFM,8135
39
44
  graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
40
- graphiti_core-0.3.6.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
41
- graphiti_core-0.3.6.dist-info/METADATA,sha256=y5lM6qDP9MbjeNExduVvVaxyi7yP7vw7mxqsbOE7k5s,9395
42
- graphiti_core-0.3.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
43
- graphiti_core-0.3.6.dist-info/RECORD,,
45
+ graphiti_core-0.3.8.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
46
+ graphiti_core-0.3.8.dist-info/METADATA,sha256=Gn4kxyPz5d-tByddy6fAUz0dUA3jaDg4gf2f04saK8Y,9395
47
+ graphiti_core-0.3.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
48
+ graphiti_core-0.3.8.dist-info/RECORD,,