From 1c27a3563bc7b3d637c66fd3b0add02be2c7490f Mon Sep 17 00:00:00 2001 From: Preston Rasmussen <109292228+prasmussen15@users.noreply.github.com> Date: Tue, 19 Aug 2025 12:31:50 -0400 Subject: [PATCH] update prompts and support thinking models (#846) * update prompts and support thinking models * update * type ignore --- examples/podcast/podcast_runner.py | 6 +++++- graphiti_core/llm_client/openai_base_client.py | 6 +++--- graphiti_core/llm_client/openai_client.py | 10 ++++++---- graphiti_core/prompts/extract_edges.py | 1 + graphiti_core/prompts/extract_nodes.py | 2 +- graphiti_core/utils/maintenance/node_operations.py | 10 +++++++--- pyproject.toml | 2 +- uv.lock | 2 +- 8 files changed, 25 insertions(+), 14 deletions(-) diff --git a/examples/podcast/podcast_runner.py b/examples/podcast/podcast_runner.py index e3cf7496..70201b9b 100644 --- a/examples/podcast/podcast_runner.py +++ b/examples/podcast/podcast_runner.py @@ -77,7 +77,11 @@ class IsPresidentOf(BaseModel): async def main(use_bulk: bool = False): setup_logging() - client = Graphiti(neo4j_uri, neo4j_user, neo4j_password) + client = Graphiti( + neo4j_uri, + neo4j_user, + neo4j_password, + ) await clear_data(client.driver) await client.build_indices_and_constraints() messages = parse_podcast_messages() diff --git a/graphiti_core/llm_client/openai_base_client.py b/graphiti_core/llm_client/openai_base_client.py index 86bcde10..3e0f075f 100644 --- a/graphiti_core/llm_client/openai_base_client.py +++ b/graphiti_core/llm_client/openai_base_client.py @@ -107,10 +107,10 @@ class BaseOpenAIClient(LLMClient): def _handle_structured_response(self, response: Any) -> dict[str, Any]: """Handle structured response parsing and validation.""" - response_object = response.choices[0].message + response_object = response.output_text - if response_object.parsed: - return response_object.parsed.model_dump() + if response_object: + return json.loads(response_object) elif response_object.refusal: raise RefusalError(response_object.refusal) else: diff --git a/graphiti_core/llm_client/openai_client.py b/graphiti_core/llm_client/openai_client.py index b47f8e9e..619dd1a9 100644 --- a/graphiti_core/llm_client/openai_client.py +++ b/graphiti_core/llm_client/openai_client.py @@ -69,14 +69,16 @@ class OpenAIClient(BaseOpenAIClient): response_model: type[BaseModel], ): """Create a structured completion using OpenAI's beta parse API.""" - return await self.client.beta.chat.completions.parse( + response = await self.client.responses.parse( model=model, - messages=messages, + input=messages, # type: ignore temperature=temperature, - max_tokens=max_tokens, - response_format=response_model, # type: ignore + max_output_tokens=max_tokens, + text_format=response_model, # type: ignore ) + return response + async def _create_completion( self, model: str, diff --git a/graphiti_core/prompts/extract_edges.py b/graphiti_core/prompts/extract_edges.py index f2a981a2..dc73cbf5 100644 --- a/graphiti_core/prompts/extract_edges.py +++ b/graphiti_core/prompts/extract_edges.py @@ -94,6 +94,7 @@ Only extract facts that: - involve two DISTINCT ENTITIES from the ENTITIES list, - are clearly stated or unambiguously implied in the CURRENT MESSAGE, and can be represented as edges in a knowledge graph. +- Facts should include entity names rather than pronouns whenever possible. - The FACT TYPES provide a list of the most important types of facts, make sure to extract facts of these types - The FACT TYPES are not an exhaustive list, extract all facts from the message even if they do not fit into one of the FACT TYPES diff --git a/graphiti_core/prompts/extract_nodes.py b/graphiti_core/prompts/extract_nodes.py index f3086b59..d5774213 100644 --- a/graphiti_core/prompts/extract_nodes.py +++ b/graphiti_core/prompts/extract_nodes.py @@ -100,7 +100,7 @@ Instructions: You are given a conversation context and a CURRENT MESSAGE. Your task is to extract **entity nodes** mentioned **explicitly or implicitly** in the CURRENT MESSAGE. Pronoun references such as he/she/they or this/that/those should be disambiguated to the names of the -reference entities. +reference entities. Only extract distinct entities from the CURRENT MESSAGE. Don't extract pronouns like you, me, he/she/they, we/us as entities. 1. **Speaker Extraction**: Always extract the speaker (the part before the colon `:` in each dialogue line) as the first entity node. - If the speaker is mentioned again in the message, treat both mentions as a **single entity**. diff --git a/graphiti_core/utils/maintenance/node_operations.py b/graphiti_core/utils/maintenance/node_operations.py index d911e569..f17706ab 100644 --- a/graphiti_core/utils/maintenance/node_operations.py +++ b/graphiti_core/utils/maintenance/node_operations.py @@ -152,9 +152,13 @@ async def extract_nodes( # Convert the extracted data into EntityNode objects extracted_nodes = [] for extracted_entity in filtered_extracted_entities: - entity_type_name = entity_types_context[extracted_entity.entity_type_id].get( - 'entity_type_name' - ) + type_id = extracted_entity.entity_type_id + if 0 <= type_id < len(entity_types_context): + entity_type_name = entity_types_context[extracted_entity.entity_type_id].get( + 'entity_type_name' + ) + else: + entity_type_name = 'Entity' # Check if this entity type should be excluded if excluded_entity_types and entity_type_name in excluded_entity_types: diff --git a/pyproject.toml b/pyproject.toml index 1a4ed522..cd02f6f6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "graphiti-core" description = "A temporal graph building library" -version = "0.18.8" +version = "0.18.9" authors = [ { name = "Paul Paliychuk", email = "paul@getzep.com" }, { name = "Preston Rasmussen", email = "preston@getzep.com" }, diff --git a/uv.lock b/uv.lock index 5f8407a6..58a2597a 100644 --- a/uv.lock +++ b/uv.lock @@ -746,7 +746,7 @@ wheels = [ [[package]] name = "graphiti-core" -version = "0.18.7" +version = "0.18.9" source = { editable = "." } dependencies = [ { name = "diskcache" },