graphiti-core 0.11.3__py3-none-any.whl → 0.11.5__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/graphiti.py +1 -1
- graphiti_core/llm_client/anthropic_client.py +6 -2
- graphiti_core/llm_client/client.py +9 -3
- graphiti_core/llm_client/config.py +13 -2
- graphiti_core/llm_client/gemini_client.py +7 -2
- graphiti_core/llm_client/groq_client.py +2 -1
- graphiti_core/llm_client/openai_client.py +13 -3
- graphiti_core/llm_client/openai_generic_client.py +4 -2
- graphiti_core/prompts/dedupe_nodes.py +19 -15
- graphiti_core/prompts/extract_nodes.py +1 -1
- graphiti_core/utils/maintenance/edge_operations.py +4 -1
- graphiti_core/utils/maintenance/node_operations.py +28 -15
- graphiti_core/utils/maintenance/temporal_operations.py +4 -1
- {graphiti_core-0.11.3.dist-info → graphiti_core-0.11.5.dist-info}/METADATA +1 -1
- {graphiti_core-0.11.3.dist-info → graphiti_core-0.11.5.dist-info}/RECORD +17 -17
- {graphiti_core-0.11.3.dist-info → graphiti_core-0.11.5.dist-info}/LICENSE +0 -0
- {graphiti_core-0.11.3.dist-info → graphiti_core-0.11.5.dist-info}/WHEEL +0 -0
graphiti_core/graphiti.py
CHANGED
|
@@ -396,7 +396,7 @@ class Graphiti:
|
|
|
396
396
|
episode.content = ''
|
|
397
397
|
|
|
398
398
|
await add_nodes_and_edges_bulk(
|
|
399
|
-
self.driver, [episode], episodic_edges,
|
|
399
|
+
self.driver, [episode], episodic_edges, hydrated_nodes, entity_edges
|
|
400
400
|
)
|
|
401
401
|
|
|
402
402
|
# Update any communities
|
|
@@ -28,7 +28,7 @@ from pydantic import BaseModel, ValidationError
|
|
|
28
28
|
|
|
29
29
|
from ..prompts.models import Message
|
|
30
30
|
from .client import LLMClient
|
|
31
|
-
from .config import DEFAULT_MAX_TOKENS, LLMConfig
|
|
31
|
+
from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
|
|
32
32
|
from .errors import RateLimitError, RefusalError
|
|
33
33
|
|
|
34
34
|
logger = logging.getLogger(__name__)
|
|
@@ -173,6 +173,7 @@ class AnthropicClient(LLMClient):
|
|
|
173
173
|
messages: list[Message],
|
|
174
174
|
response_model: type[BaseModel] | None = None,
|
|
175
175
|
max_tokens: int | None = None,
|
|
176
|
+
model_size: ModelSize = ModelSize.medium,
|
|
176
177
|
) -> dict[str, typing.Any]:
|
|
177
178
|
"""
|
|
178
179
|
Generate a response from the Anthropic LLM using tool-based approach for all requests.
|
|
@@ -263,6 +264,7 @@ class AnthropicClient(LLMClient):
|
|
|
263
264
|
messages: list[Message],
|
|
264
265
|
response_model: type[BaseModel] | None = None,
|
|
265
266
|
max_tokens: int | None = None,
|
|
267
|
+
model_size: ModelSize = ModelSize.medium,
|
|
266
268
|
) -> dict[str, typing.Any]:
|
|
267
269
|
"""
|
|
268
270
|
Generate a response from the LLM.
|
|
@@ -289,7 +291,9 @@ class AnthropicClient(LLMClient):
|
|
|
289
291
|
|
|
290
292
|
while retry_count <= max_retries:
|
|
291
293
|
try:
|
|
292
|
-
response = await self._generate_response(
|
|
294
|
+
response = await self._generate_response(
|
|
295
|
+
messages, response_model, max_tokens, model_size
|
|
296
|
+
)
|
|
293
297
|
|
|
294
298
|
# If we have a response_model, attempt to validate the response
|
|
295
299
|
if response_model is not None:
|
|
@@ -26,7 +26,7 @@ from pydantic import BaseModel
|
|
|
26
26
|
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_random_exponential
|
|
27
27
|
|
|
28
28
|
from ..prompts.models import Message
|
|
29
|
-
from .config import DEFAULT_MAX_TOKENS, LLMConfig
|
|
29
|
+
from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
|
|
30
30
|
from .errors import RateLimitError
|
|
31
31
|
|
|
32
32
|
DEFAULT_TEMPERATURE = 0
|
|
@@ -55,6 +55,7 @@ class LLMClient(ABC):
|
|
|
55
55
|
|
|
56
56
|
self.config = config
|
|
57
57
|
self.model = config.model
|
|
58
|
+
self.small_model = config.small_model
|
|
58
59
|
self.temperature = config.temperature
|
|
59
60
|
self.max_tokens = config.max_tokens
|
|
60
61
|
self.cache_enabled = cache
|
|
@@ -102,9 +103,10 @@ class LLMClient(ABC):
|
|
|
102
103
|
messages: list[Message],
|
|
103
104
|
response_model: type[BaseModel] | None = None,
|
|
104
105
|
max_tokens: int = DEFAULT_MAX_TOKENS,
|
|
106
|
+
model_size: ModelSize = ModelSize.medium,
|
|
105
107
|
) -> dict[str, typing.Any]:
|
|
106
108
|
try:
|
|
107
|
-
return await self._generate_response(messages, response_model, max_tokens)
|
|
109
|
+
return await self._generate_response(messages, response_model, max_tokens, model_size)
|
|
108
110
|
except (httpx.HTTPStatusError, RateLimitError) as e:
|
|
109
111
|
raise e
|
|
110
112
|
|
|
@@ -114,6 +116,7 @@ class LLMClient(ABC):
|
|
|
114
116
|
messages: list[Message],
|
|
115
117
|
response_model: type[BaseModel] | None = None,
|
|
116
118
|
max_tokens: int = DEFAULT_MAX_TOKENS,
|
|
119
|
+
model_size: ModelSize = ModelSize.medium,
|
|
117
120
|
) -> dict[str, typing.Any]:
|
|
118
121
|
pass
|
|
119
122
|
|
|
@@ -128,6 +131,7 @@ class LLMClient(ABC):
|
|
|
128
131
|
messages: list[Message],
|
|
129
132
|
response_model: type[BaseModel] | None = None,
|
|
130
133
|
max_tokens: int | None = None,
|
|
134
|
+
model_size: ModelSize = ModelSize.medium,
|
|
131
135
|
) -> dict[str, typing.Any]:
|
|
132
136
|
if max_tokens is None:
|
|
133
137
|
max_tokens = self.max_tokens
|
|
@@ -154,7 +158,9 @@ class LLMClient(ABC):
|
|
|
154
158
|
for message in messages:
|
|
155
159
|
message.content = self._clean_input(message.content)
|
|
156
160
|
|
|
157
|
-
response = await self._generate_response_with_retry(
|
|
161
|
+
response = await self._generate_response_with_retry(
|
|
162
|
+
messages, response_model, max_tokens, model_size
|
|
163
|
+
)
|
|
158
164
|
|
|
159
165
|
if self.cache_enabled and self.cache_dir is not None:
|
|
160
166
|
cache_key = self._get_cache_key(messages)
|
|
@@ -14,10 +14,17 @@ See the License for the specific language governing permissions and
|
|
|
14
14
|
limitations under the License.
|
|
15
15
|
"""
|
|
16
16
|
|
|
17
|
+
from enum import Enum
|
|
18
|
+
|
|
17
19
|
DEFAULT_MAX_TOKENS = 8192
|
|
18
20
|
DEFAULT_TEMPERATURE = 0
|
|
19
21
|
|
|
20
22
|
|
|
23
|
+
class ModelSize(Enum):
|
|
24
|
+
small = 'small'
|
|
25
|
+
medium = 'medium'
|
|
26
|
+
|
|
27
|
+
|
|
21
28
|
class LLMConfig:
|
|
22
29
|
"""
|
|
23
30
|
Configuration class for the Language Learning Model (LLM).
|
|
@@ -34,6 +41,7 @@ class LLMConfig:
|
|
|
34
41
|
base_url: str | None = None,
|
|
35
42
|
temperature: float = DEFAULT_TEMPERATURE,
|
|
36
43
|
max_tokens: int = DEFAULT_MAX_TOKENS,
|
|
44
|
+
small_model: str | None = None,
|
|
37
45
|
):
|
|
38
46
|
"""
|
|
39
47
|
Initialize the LLMConfig with the provided parameters.
|
|
@@ -43,15 +51,18 @@ class LLMConfig:
|
|
|
43
51
|
This is required for making authorized requests.
|
|
44
52
|
|
|
45
53
|
model (str, optional): The specific LLM model to use for generating responses.
|
|
46
|
-
Defaults to "gpt-4.1-mini"
|
|
47
|
-
Common values might include "gpt-3.5-turbo" or "gpt-4".
|
|
54
|
+
Defaults to "gpt-4.1-mini".
|
|
48
55
|
|
|
49
56
|
base_url (str, optional): The base URL of the LLM API service.
|
|
50
57
|
Defaults to "https://api.openai.com", which is OpenAI's standard API endpoint.
|
|
51
58
|
This can be changed if using a different provider or a custom endpoint.
|
|
59
|
+
|
|
60
|
+
small_model (str, optional): The specific LLM model to use for generating responses of simpler prompts.
|
|
61
|
+
Defaults to "gpt-4.1-nano".
|
|
52
62
|
"""
|
|
53
63
|
self.base_url = base_url
|
|
54
64
|
self.api_key = api_key
|
|
55
65
|
self.model = model
|
|
66
|
+
self.small_model = small_model
|
|
56
67
|
self.temperature = temperature
|
|
57
68
|
self.max_tokens = max_tokens
|
|
@@ -24,7 +24,7 @@ from pydantic import BaseModel
|
|
|
24
24
|
|
|
25
25
|
from ..prompts.models import Message
|
|
26
26
|
from .client import LLMClient
|
|
27
|
-
from .config import DEFAULT_MAX_TOKENS, LLMConfig
|
|
27
|
+
from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
|
|
28
28
|
from .errors import RateLimitError
|
|
29
29
|
|
|
30
30
|
logger = logging.getLogger(__name__)
|
|
@@ -82,6 +82,7 @@ class GeminiClient(LLMClient):
|
|
|
82
82
|
messages: list[Message],
|
|
83
83
|
response_model: type[BaseModel] | None = None,
|
|
84
84
|
max_tokens: int = DEFAULT_MAX_TOKENS,
|
|
85
|
+
model_size: ModelSize = ModelSize.medium,
|
|
85
86
|
) -> dict[str, typing.Any]:
|
|
86
87
|
"""
|
|
87
88
|
Generate a response from the Gemini language model.
|
|
@@ -167,6 +168,7 @@ class GeminiClient(LLMClient):
|
|
|
167
168
|
messages: list[Message],
|
|
168
169
|
response_model: type[BaseModel] | None = None,
|
|
169
170
|
max_tokens: int | None = None,
|
|
171
|
+
model_size: ModelSize = ModelSize.medium,
|
|
170
172
|
) -> dict[str, typing.Any]:
|
|
171
173
|
"""
|
|
172
174
|
Generate a response from the Gemini language model.
|
|
@@ -185,5 +187,8 @@ class GeminiClient(LLMClient):
|
|
|
185
187
|
|
|
186
188
|
# Call the internal _generate_response method
|
|
187
189
|
return await self._generate_response(
|
|
188
|
-
messages=messages,
|
|
190
|
+
messages=messages,
|
|
191
|
+
response_model=response_model,
|
|
192
|
+
max_tokens=max_tokens,
|
|
193
|
+
model_size=model_size,
|
|
189
194
|
)
|
|
@@ -25,7 +25,7 @@ from pydantic import BaseModel
|
|
|
25
25
|
|
|
26
26
|
from ..prompts.models import Message
|
|
27
27
|
from .client import LLMClient
|
|
28
|
-
from .config import LLMConfig
|
|
28
|
+
from .config import LLMConfig, ModelSize
|
|
29
29
|
from .errors import RateLimitError
|
|
30
30
|
|
|
31
31
|
logger = logging.getLogger(__name__)
|
|
@@ -49,6 +49,7 @@ class GroqClient(LLMClient):
|
|
|
49
49
|
messages: list[Message],
|
|
50
50
|
response_model: type[BaseModel] | None = None,
|
|
51
51
|
max_tokens: int = DEFAULT_MAX_TOKENS,
|
|
52
|
+
model_size: ModelSize = ModelSize.medium,
|
|
52
53
|
) -> dict[str, typing.Any]:
|
|
53
54
|
msgs: list[ChatCompletionMessageParam] = []
|
|
54
55
|
for m in messages:
|
|
@@ -25,12 +25,13 @@ from pydantic import BaseModel
|
|
|
25
25
|
|
|
26
26
|
from ..prompts.models import Message
|
|
27
27
|
from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
|
|
28
|
-
from .config import DEFAULT_MAX_TOKENS, LLMConfig
|
|
28
|
+
from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
|
|
29
29
|
from .errors import RateLimitError, RefusalError
|
|
30
30
|
|
|
31
31
|
logger = logging.getLogger(__name__)
|
|
32
32
|
|
|
33
33
|
DEFAULT_MODEL = 'gpt-4.1-mini'
|
|
34
|
+
DEFAULT_SMALL_MODEL = 'gpt-4.1-nano'
|
|
34
35
|
|
|
35
36
|
|
|
36
37
|
class OpenAIClient(LLMClient):
|
|
@@ -94,6 +95,7 @@ class OpenAIClient(LLMClient):
|
|
|
94
95
|
messages: list[Message],
|
|
95
96
|
response_model: type[BaseModel] | None = None,
|
|
96
97
|
max_tokens: int = DEFAULT_MAX_TOKENS,
|
|
98
|
+
model_size: ModelSize = ModelSize.medium,
|
|
97
99
|
) -> dict[str, typing.Any]:
|
|
98
100
|
openai_messages: list[ChatCompletionMessageParam] = []
|
|
99
101
|
for m in messages:
|
|
@@ -103,8 +105,13 @@ class OpenAIClient(LLMClient):
|
|
|
103
105
|
elif m.role == 'system':
|
|
104
106
|
openai_messages.append({'role': 'system', 'content': m.content})
|
|
105
107
|
try:
|
|
108
|
+
if model_size == ModelSize.small:
|
|
109
|
+
model = self.small_model or DEFAULT_SMALL_MODEL
|
|
110
|
+
else:
|
|
111
|
+
model = self.model or DEFAULT_MODEL
|
|
112
|
+
|
|
106
113
|
response = await self.client.beta.chat.completions.parse(
|
|
107
|
-
model=
|
|
114
|
+
model=model,
|
|
108
115
|
messages=openai_messages,
|
|
109
116
|
temperature=self.temperature,
|
|
110
117
|
max_tokens=max_tokens or self.max_tokens,
|
|
@@ -132,6 +139,7 @@ class OpenAIClient(LLMClient):
|
|
|
132
139
|
messages: list[Message],
|
|
133
140
|
response_model: type[BaseModel] | None = None,
|
|
134
141
|
max_tokens: int | None = None,
|
|
142
|
+
model_size: ModelSize = ModelSize.medium,
|
|
135
143
|
) -> dict[str, typing.Any]:
|
|
136
144
|
if max_tokens is None:
|
|
137
145
|
max_tokens = self.max_tokens
|
|
@@ -144,7 +152,9 @@ class OpenAIClient(LLMClient):
|
|
|
144
152
|
|
|
145
153
|
while retry_count <= self.MAX_RETRIES:
|
|
146
154
|
try:
|
|
147
|
-
response = await self._generate_response(
|
|
155
|
+
response = await self._generate_response(
|
|
156
|
+
messages, response_model, max_tokens, model_size
|
|
157
|
+
)
|
|
148
158
|
return response
|
|
149
159
|
except (RateLimitError, RefusalError):
|
|
150
160
|
# These errors should not trigger retries
|
|
@@ -26,7 +26,7 @@ from pydantic import BaseModel
|
|
|
26
26
|
|
|
27
27
|
from ..prompts.models import Message
|
|
28
28
|
from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
|
|
29
|
-
from .config import DEFAULT_MAX_TOKENS, LLMConfig
|
|
29
|
+
from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
|
|
30
30
|
from .errors import RateLimitError, RefusalError
|
|
31
31
|
|
|
32
32
|
logger = logging.getLogger(__name__)
|
|
@@ -89,6 +89,7 @@ class OpenAIGenericClient(LLMClient):
|
|
|
89
89
|
messages: list[Message],
|
|
90
90
|
response_model: type[BaseModel] | None = None,
|
|
91
91
|
max_tokens: int = DEFAULT_MAX_TOKENS,
|
|
92
|
+
model_size: ModelSize = ModelSize.medium,
|
|
92
93
|
) -> dict[str, typing.Any]:
|
|
93
94
|
openai_messages: list[ChatCompletionMessageParam] = []
|
|
94
95
|
for m in messages:
|
|
@@ -118,6 +119,7 @@ class OpenAIGenericClient(LLMClient):
|
|
|
118
119
|
messages: list[Message],
|
|
119
120
|
response_model: type[BaseModel] | None = None,
|
|
120
121
|
max_tokens: int | None = None,
|
|
122
|
+
model_size: ModelSize = ModelSize.medium,
|
|
121
123
|
) -> dict[str, typing.Any]:
|
|
122
124
|
if max_tokens is None:
|
|
123
125
|
max_tokens = self.max_tokens
|
|
@@ -139,7 +141,7 @@ class OpenAIGenericClient(LLMClient):
|
|
|
139
141
|
while retry_count <= self.MAX_RETRIES:
|
|
140
142
|
try:
|
|
141
143
|
response = await self._generate_response(
|
|
142
|
-
messages, response_model, max_tokens=max_tokens
|
|
144
|
+
messages, response_model, max_tokens=max_tokens, model_size=model_size
|
|
143
145
|
)
|
|
144
146
|
return response
|
|
145
147
|
except (RateLimitError, RefusalError):
|
|
@@ -27,6 +27,7 @@ class NodeDuplicate(BaseModel):
|
|
|
27
27
|
...,
|
|
28
28
|
description='id of the duplicate node. If no duplicate nodes are found, default to -1.',
|
|
29
29
|
)
|
|
30
|
+
name: str = Field(..., description='Name of the entity.')
|
|
30
31
|
|
|
31
32
|
|
|
32
33
|
class Prompt(Protocol):
|
|
@@ -43,7 +44,7 @@ def node(context: dict[str, Any]) -> list[Message]:
|
|
|
43
44
|
return [
|
|
44
45
|
Message(
|
|
45
46
|
role='system',
|
|
46
|
-
content='You are a helpful assistant that de-duplicates
|
|
47
|
+
content='You are a helpful assistant that de-duplicates entities from entity lists.',
|
|
47
48
|
),
|
|
48
49
|
Message(
|
|
49
50
|
role='user',
|
|
@@ -54,25 +55,28 @@ def node(context: dict[str, Any]) -> list[Message]:
|
|
|
54
55
|
<CURRENT MESSAGE>
|
|
55
56
|
{context['episode_content']}
|
|
56
57
|
</CURRENT MESSAGE>
|
|
58
|
+
<NEW ENTITY>
|
|
59
|
+
{json.dumps(context['extracted_node'], indent=2)}
|
|
60
|
+
</NEW ENTITY>
|
|
61
|
+
<ENTITY TYPE DESCRIPTION>
|
|
62
|
+
{json.dumps(context['entity_type_description'], indent=2)}
|
|
63
|
+
</ENTITY TYPE DESCRIPTION>
|
|
57
64
|
|
|
58
|
-
<EXISTING
|
|
65
|
+
<EXISTING ENTITIES>
|
|
59
66
|
{json.dumps(context['existing_nodes'], indent=2)}
|
|
60
|
-
</EXISTING
|
|
67
|
+
</EXISTING ENTITIES>
|
|
61
68
|
|
|
62
|
-
Given the above EXISTING
|
|
63
|
-
is a duplicate entity of one of the EXISTING
|
|
69
|
+
Given the above EXISTING ENTITIES and their attributes, MESSAGE, and PREVIOUS MESSAGES; Determine if the NEW ENTITY extracted from the conversation
|
|
70
|
+
is a duplicate entity of one of the EXISTING ENTITIES.
|
|
71
|
+
|
|
72
|
+
The ENTITY TYPE DESCRIPTION gives more insight into what the entity type means for the NEW ENTITY.
|
|
64
73
|
|
|
65
|
-
<NEW NODE>
|
|
66
|
-
{json.dumps(context['extracted_node'], indent=2)}
|
|
67
|
-
</NEW NODE>
|
|
68
74
|
Task:
|
|
69
|
-
If the NEW
|
|
70
|
-
id of the EXISTING
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
1. Use the name, summary, and attributes of nodes to determine if the entities are duplicates,
|
|
75
|
-
duplicate nodes may have different names
|
|
75
|
+
If the NEW ENTITY represents a duplicate entity of any entity in EXISTING ENTITIES, set duplicate_entity_id to the
|
|
76
|
+
id of the EXISTING ENTITY that is the duplicate. If the NEW ENTITY is not a duplicate of any of the EXISTING ENTITIES,
|
|
77
|
+
duplicate_entity_id should be set to -1.
|
|
78
|
+
|
|
79
|
+
Also return the most complete name for the entity.
|
|
76
80
|
""",
|
|
77
81
|
),
|
|
78
82
|
]
|
|
@@ -256,7 +256,7 @@ def extract_attributes(context: dict[str, Any]) -> list[Message]:
|
|
|
256
256
|
1. Do not hallucinate entity property values if they cannot be found in the current context.
|
|
257
257
|
2. Only use the provided MESSAGES and ENTITY to set attribute values.
|
|
258
258
|
3. The summary attribute represents a summary of the ENTITY, and should be updated with new information about the Entity from the MESSAGES.
|
|
259
|
-
Summaries must be no longer than
|
|
259
|
+
Summaries must be no longer than 500 words.
|
|
260
260
|
|
|
261
261
|
<ENTITY>
|
|
262
262
|
{context['node']}
|
|
@@ -27,6 +27,7 @@ from graphiti_core.edges import (
|
|
|
27
27
|
from graphiti_core.graphiti_types import GraphitiClients
|
|
28
28
|
from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS, semaphore_gather
|
|
29
29
|
from graphiti_core.llm_client import LLMClient
|
|
30
|
+
from graphiti_core.llm_client.config import ModelSize
|
|
30
31
|
from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode
|
|
31
32
|
from graphiti_core.prompts import prompt_library
|
|
32
33
|
from graphiti_core.prompts.dedupe_edges import EdgeDuplicate, UniqueFacts
|
|
@@ -377,7 +378,9 @@ async def dedupe_extracted_edge(
|
|
|
377
378
|
}
|
|
378
379
|
|
|
379
380
|
llm_response = await llm_client.generate_response(
|
|
380
|
-
prompt_library.dedupe_edges.edge(context),
|
|
381
|
+
prompt_library.dedupe_edges.edge(context),
|
|
382
|
+
response_model=EdgeDuplicate,
|
|
383
|
+
model_size=ModelSize.small,
|
|
381
384
|
)
|
|
382
385
|
|
|
383
386
|
duplicate_fact_id: int = llm_response.get('duplicate_fact_id', -1)
|
|
@@ -25,6 +25,7 @@ from pydantic import BaseModel, Field
|
|
|
25
25
|
from graphiti_core.graphiti_types import GraphitiClients
|
|
26
26
|
from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS, semaphore_gather
|
|
27
27
|
from graphiti_core.llm_client import LLMClient
|
|
28
|
+
from graphiti_core.llm_client.config import ModelSize
|
|
28
29
|
from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode, create_entity_node_embeddings
|
|
29
30
|
from graphiti_core.prompts import prompt_library
|
|
30
31
|
from graphiti_core.prompts.dedupe_nodes import NodeDuplicate
|
|
@@ -33,8 +34,10 @@ from graphiti_core.prompts.extract_nodes import (
|
|
|
33
34
|
ExtractedEntity,
|
|
34
35
|
MissedEntities,
|
|
35
36
|
)
|
|
37
|
+
from graphiti_core.search.search import search
|
|
38
|
+
from graphiti_core.search.search_config import SearchResults
|
|
39
|
+
from graphiti_core.search.search_config_recipes import NODE_HYBRID_SEARCH_RRF
|
|
36
40
|
from graphiti_core.search.search_filters import SearchFilters
|
|
37
|
-
from graphiti_core.search.search_utils import get_relevant_nodes
|
|
38
41
|
from graphiti_core.utils.datetime_utils import utc_now
|
|
39
42
|
|
|
40
43
|
logger = logging.getLogger(__name__)
|
|
@@ -226,13 +229,23 @@ async def resolve_extracted_nodes(
|
|
|
226
229
|
entity_types: dict[str, BaseModel] | None = None,
|
|
227
230
|
) -> tuple[list[EntityNode], dict[str, str]]:
|
|
228
231
|
llm_client = clients.llm_client
|
|
229
|
-
driver = clients.driver
|
|
230
232
|
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
|
|
233
|
+
search_results: list[SearchResults] = await semaphore_gather(
|
|
234
|
+
*[
|
|
235
|
+
search(
|
|
236
|
+
clients=clients,
|
|
237
|
+
query=node.name,
|
|
238
|
+
query_vector=node.name_embedding,
|
|
239
|
+
group_ids=[node.group_id],
|
|
240
|
+
search_filter=SearchFilters(),
|
|
241
|
+
config=NODE_HYBRID_SEARCH_RRF,
|
|
242
|
+
)
|
|
243
|
+
for node in extracted_nodes
|
|
244
|
+
]
|
|
234
245
|
)
|
|
235
246
|
|
|
247
|
+
existing_nodes_lists: list[list[EntityNode]] = [result.nodes for result in search_results]
|
|
248
|
+
|
|
236
249
|
resolved_nodes: list[EntityNode] = await semaphore_gather(
|
|
237
250
|
*[
|
|
238
251
|
resolve_extracted_node(
|
|
@@ -291,14 +304,14 @@ async def resolve_extracted_node(
|
|
|
291
304
|
extracted_node_context = {
|
|
292
305
|
'name': extracted_node.name,
|
|
293
306
|
'entity_type': entity_type.__name__ if entity_type is not None else 'Entity', # type: ignore
|
|
294
|
-
'entity_type_description': entity_type.__doc__
|
|
295
|
-
if entity_type is not None
|
|
296
|
-
else 'Default Entity Type',
|
|
297
307
|
}
|
|
298
308
|
|
|
299
309
|
context = {
|
|
300
310
|
'existing_nodes': existing_nodes_context,
|
|
301
311
|
'extracted_node': extracted_node_context,
|
|
312
|
+
'entity_type_description': entity_type.__doc__
|
|
313
|
+
if entity_type is not None
|
|
314
|
+
else 'Default Entity Type',
|
|
302
315
|
'episode_content': episode.content if episode is not None else '',
|
|
303
316
|
'previous_episodes': [ep.content for ep in previous_episodes]
|
|
304
317
|
if previous_episodes is not None
|
|
@@ -306,7 +319,9 @@ async def resolve_extracted_node(
|
|
|
306
319
|
}
|
|
307
320
|
|
|
308
321
|
llm_response = await llm_client.generate_response(
|
|
309
|
-
prompt_library.dedupe_nodes.node(context),
|
|
322
|
+
prompt_library.dedupe_nodes.node(context),
|
|
323
|
+
response_model=NodeDuplicate,
|
|
324
|
+
model_size=ModelSize.small,
|
|
310
325
|
)
|
|
311
326
|
|
|
312
327
|
duplicate_id: int = llm_response.get('duplicate_node_id', -1)
|
|
@@ -315,6 +330,8 @@ async def resolve_extracted_node(
|
|
|
315
330
|
existing_nodes[duplicate_id] if 0 <= duplicate_id < len(existing_nodes) else extracted_node
|
|
316
331
|
)
|
|
317
332
|
|
|
333
|
+
node.name = llm_response.get('name', '')
|
|
334
|
+
|
|
318
335
|
end = time()
|
|
319
336
|
logger.debug(
|
|
320
337
|
f'Resolved node: {extracted_node.name} is {node.name}, in {(end - start) * 1000} ms'
|
|
@@ -371,13 +388,9 @@ async def extract_attributes_from_node(
|
|
|
371
388
|
'summary': (
|
|
372
389
|
str,
|
|
373
390
|
Field(
|
|
374
|
-
description='Summary containing the important information about the entity. Under
|
|
391
|
+
description='Summary containing the important information about the entity. Under 500 words',
|
|
375
392
|
),
|
|
376
|
-
)
|
|
377
|
-
'name': (
|
|
378
|
-
str,
|
|
379
|
-
Field(description='Name of the ENTITY'),
|
|
380
|
-
),
|
|
393
|
+
)
|
|
381
394
|
}
|
|
382
395
|
|
|
383
396
|
if entity_type is not None:
|
|
@@ -20,6 +20,7 @@ from time import time
|
|
|
20
20
|
|
|
21
21
|
from graphiti_core.edges import EntityEdge
|
|
22
22
|
from graphiti_core.llm_client import LLMClient
|
|
23
|
+
from graphiti_core.llm_client.config import ModelSize
|
|
23
24
|
from graphiti_core.nodes import EpisodicNode
|
|
24
25
|
from graphiti_core.prompts import prompt_library
|
|
25
26
|
from graphiti_core.prompts.extract_edge_dates import EdgeDates
|
|
@@ -81,7 +82,9 @@ async def get_edge_contradictions(
|
|
|
81
82
|
context = {'new_edge': new_edge_context, 'existing_edges': existing_edge_context}
|
|
82
83
|
|
|
83
84
|
llm_response = await llm_client.generate_response(
|
|
84
|
-
prompt_library.invalidate_edges.v2(context),
|
|
85
|
+
prompt_library.invalidate_edges.v2(context),
|
|
86
|
+
response_model=InvalidatedEdges,
|
|
87
|
+
model_size=ModelSize.small,
|
|
85
88
|
)
|
|
86
89
|
|
|
87
90
|
contradicted_facts: list[int] = llm_response.get('contradicted_facts', [])
|
|
@@ -10,18 +10,18 @@ graphiti_core/embedder/gemini.py,sha256=Dh80q21auMvDBjwqHsI_wFrJtgWwCXRHzwg31-BS
|
|
|
10
10
|
graphiti_core/embedder/openai.py,sha256=bIThUoLMeGlHG2-3VikzK6JZfOHKn4PKvUMx5sHxJy8,2192
|
|
11
11
|
graphiti_core/embedder/voyage.py,sha256=gQhdcz2IYPSyOcDn3w8aHToVS3KQhyZrUBm4vqr3WcE,2224
|
|
12
12
|
graphiti_core/errors.py,sha256=Nib1uQx2cO_VOizupmRjpFfmuRg-hFAVqTtZAuBehR8,2405
|
|
13
|
-
graphiti_core/graphiti.py,sha256=
|
|
13
|
+
graphiti_core/graphiti.py,sha256=14i2hArJSJtaWSdEC3kluMxRu5WVn4lXb59a5g6oMNs,26984
|
|
14
14
|
graphiti_core/graphiti_types.py,sha256=46ueysKPwUCpxkMePHdCJLspfTImoZN7JiRwpz7cqd0,1013
|
|
15
15
|
graphiti_core/helpers.py,sha256=Rj_ayM7hphlyY4v0ln_9vjFtr76aTjLVenYny0Wntwo,3086
|
|
16
16
|
graphiti_core/llm_client/__init__.py,sha256=PA80TSMeX-sUXITXEAxMDEt3gtfZgcJrGJUcyds1mSo,207
|
|
17
|
-
graphiti_core/llm_client/anthropic_client.py,sha256=
|
|
18
|
-
graphiti_core/llm_client/client.py,sha256=
|
|
19
|
-
graphiti_core/llm_client/config.py,sha256=
|
|
17
|
+
graphiti_core/llm_client/anthropic_client.py,sha256=Fo__9YUUXdo_czvF5MPai8y0bpMgHWWnJZ1RWYpRCcg,13153
|
|
18
|
+
graphiti_core/llm_client/client.py,sha256=v_w5TBbDJYYADCXSs2r287g5Ami2Urma-GGEbHSI_Jg,5826
|
|
19
|
+
graphiti_core/llm_client/config.py,sha256=90IgSBxZE_3nWdaEONVLUznI8lytPA7ZyexQz-_c55U,2560
|
|
20
20
|
graphiti_core/llm_client/errors.py,sha256=pn6brRiLW60DAUIXJYKBT6MInrS4ueuH1hNLbn_JbQo,1243
|
|
21
|
-
graphiti_core/llm_client/gemini_client.py,sha256=
|
|
22
|
-
graphiti_core/llm_client/groq_client.py,sha256=
|
|
23
|
-
graphiti_core/llm_client/openai_client.py,sha256=
|
|
24
|
-
graphiti_core/llm_client/openai_generic_client.py,sha256=
|
|
21
|
+
graphiti_core/llm_client/gemini_client.py,sha256=JdcQTvwbaqko0alodUW3WP328i6Pu_GLUQ9yBAFBXwY,7558
|
|
22
|
+
graphiti_core/llm_client/groq_client.py,sha256=k7zbXHfOpb4jhvvKFsccVYTq4yGGpxmY7xzNA02N2zk,2559
|
|
23
|
+
graphiti_core/llm_client/openai_client.py,sha256=lLTZkd-PxEicTBmQefGoWLGTCb4QSU2Cq3x5W4kRYXg,7412
|
|
24
|
+
graphiti_core/llm_client/openai_generic_client.py,sha256=WElMnPqdb1CxzYH4p2-m_9rVMr5M93-eXnc3yVxBgFg,7001
|
|
25
25
|
graphiti_core/llm_client/utils.py,sha256=zKpxXEbKa369m4W7RDEf-m56kH46V1Mx3RowcWZEWWs,1000
|
|
26
26
|
graphiti_core/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
27
27
|
graphiti_core/models/edges/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
@@ -31,11 +31,11 @@ graphiti_core/models/nodes/node_db_queries.py,sha256=AQgRGVO-GgFWfLq1G6k8s86WItw
|
|
|
31
31
|
graphiti_core/nodes.py,sha256=_V1KkM31MEVv688ovXlXmBVeibm20d8FEps5jMDPlp0,17774
|
|
32
32
|
graphiti_core/prompts/__init__.py,sha256=EA-x9xUki9l8wnu2l8ek_oNf75-do5tq5hVq7Zbv8Kw,101
|
|
33
33
|
graphiti_core/prompts/dedupe_edges.py,sha256=q60fqIjFQlOzOeL7Y35gwABWQBqKkMarBQBok1pj1C4,3409
|
|
34
|
-
graphiti_core/prompts/dedupe_nodes.py,sha256=
|
|
34
|
+
graphiti_core/prompts/dedupe_nodes.py,sha256=hdeBMkSL1Ixt-Pu89PIDCoLvyc0zy2uX3kp4MzY2kzI,4138
|
|
35
35
|
graphiti_core/prompts/eval.py,sha256=gnBQTmwsCl3Qvwpcm7aieVszzo6y1sMCUT8jQiKTvvE,5317
|
|
36
36
|
graphiti_core/prompts/extract_edge_dates.py,sha256=3Drs3CmvP0gJN5BidWSxrNvLet3HPoTybU3BUIAoc0Y,4218
|
|
37
37
|
graphiti_core/prompts/extract_edges.py,sha256=uSoQS32rpUzQJGEhayErMdj72rocGdqVq54Macjf7po,5102
|
|
38
|
-
graphiti_core/prompts/extract_nodes.py,sha256=
|
|
38
|
+
graphiti_core/prompts/extract_nodes.py,sha256=7aQTZaIjcEBHOS6GKcFw3FEpQ4rbXSs6CuH6GmoPUGE,9624
|
|
39
39
|
graphiti_core/prompts/invalidate_edges.py,sha256=3KZQ-Hyop2hae1jK_8GIdUg4ltvFVEDQpvvzt98KvwY,3547
|
|
40
40
|
graphiti_core/prompts/lib.py,sha256=DCyHePM4_q-CptTpEXGO_dBv9k7xDtclEaB1dGu7EcI,4092
|
|
41
41
|
graphiti_core/prompts/models.py,sha256=NgxdbPHJpBEcpbXovKyScgpBc73Q-GIW-CBDlBtDjto,894
|
|
@@ -54,13 +54,13 @@ graphiti_core/utils/bulk_utils.py,sha256=Z1QzpcopVHxttRmVeYdVDHRjjY7z0cW1wwHB-0y
|
|
|
54
54
|
graphiti_core/utils/datetime_utils.py,sha256=Ti-2tnrDFRzBsbfblzsHybsM3jaDLP4-VT2t0VhpIzU,1357
|
|
55
55
|
graphiti_core/utils/maintenance/__init__.py,sha256=vW4H1KyapTl-OOz578uZABYcpND4wPx3Vt6aAPaXh78,301
|
|
56
56
|
graphiti_core/utils/maintenance/community_operations.py,sha256=aLMfQ5l8PkNJY3ugqV6UOMG74aI6L9Zz2zFiLBsnKvM,10155
|
|
57
|
-
graphiti_core/utils/maintenance/edge_operations.py,sha256=
|
|
57
|
+
graphiti_core/utils/maintenance/edge_operations.py,sha256=xfYo49ufpJmGG6sinMxwRwGx83swADoE40m8bp3gC30,14418
|
|
58
58
|
graphiti_core/utils/maintenance/graph_data_operations.py,sha256=BIJKc8tbvU4IjWxLgeotw57b1eE3Iw8YtV74j6eo4RQ,7493
|
|
59
|
-
graphiti_core/utils/maintenance/node_operations.py,sha256=
|
|
60
|
-
graphiti_core/utils/maintenance/temporal_operations.py,sha256=
|
|
59
|
+
graphiti_core/utils/maintenance/node_operations.py,sha256=9pvAEoVIiT_RpRJYvZpvSobP2_oq-37IdHRHRcmVO1Q,15559
|
|
60
|
+
graphiti_core/utils/maintenance/temporal_operations.py,sha256=mJkw9xLB4W2BsLfC5POr0r-PHWL9SIfNj_l_xu0B5ug,3410
|
|
61
61
|
graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
62
62
|
graphiti_core/utils/ontology_utils/entity_types_utils.py,sha256=QJX5cG0GSSNF_Mm_yrldr69wjVAbN_MxLhOSznz85Hk,1279
|
|
63
|
-
graphiti_core-0.11.
|
|
64
|
-
graphiti_core-0.11.
|
|
65
|
-
graphiti_core-0.11.
|
|
66
|
-
graphiti_core-0.11.
|
|
63
|
+
graphiti_core-0.11.5.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
|
|
64
|
+
graphiti_core-0.11.5.dist-info/METADATA,sha256=VW8Ca8H_RpAoQ232EdjNLl6AlMzjWCUkqUY8bXp_k-A,15348
|
|
65
|
+
graphiti_core-0.11.5.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
|
|
66
|
+
graphiti_core-0.11.5.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|