Add support for non-ASCII characters in LLM prompts (#805)
* Add support for non-ASCII characters in LLM prompts - Add ensure_ascii parameter to Graphiti class (default: True) - Create to_prompt_json helper function for consistent JSON serialization - Update all prompt files to use new helper function - Preserve Korean/Japanese/Chinese characters when ensure_ascii=False - Maintain backward compatibility with existing behavior Fixes issue where non-ASCII characters were escaped as unicode sequences in prompts, making them unreadable in LLM logs and potentially affecting model understanding. * Remove unused json imports after replacing with to_prompt_json helper - Fix ruff lint errors (F401) for unused json imports - All prompt files now use to_prompt_json helper instead of json.dumps - Maintains clean code style and passes lint checks * Fix ensure_ascii propagation to all LLM calls - Add ensure_ascii parameter to maintenance operation functions that were missing it - Update function signatures in node_operations, community_operations, temporal_operations, and edge_operations - Ensure all llm_client.generate_response calls receive proper ensure_ascii context - Fix hardcoded ensure_ascii: True values that prevented non-ASCII character preservation - Maintain backward compatibility with default ensure_ascii=True - Complete the fix for issue #804 ensuring Korean/Japanese/Chinese characters are properly handled in LLM prompts
This commit is contained in:
HUGO SON
GitHub
parent
0f186c59ca
commit
ce9ef3ca79
15 changed files with 144 additions and 58 deletions
|
|
@ -123,6 +123,7 @@ class Graphiti:
|
|||
store_raw_episode_content: bool = True,
|
||||
graph_driver: GraphDriver | None = None,
|
||||
max_coroutines: int | None = None,
|
||||
ensure_ascii: bool = True,
|
||||
):
|
||||
"""
|
||||
Initialize a Graphiti instance.
|
||||
|
|
@ -155,6 +156,10 @@ class Graphiti:
|
|||
max_coroutines : int | None, optional
|
||||
The maximum number of concurrent operations allowed. Overrides SEMAPHORE_LIMIT set in the environment.
|
||||
If not set, the Graphiti default is used.
|
||||
ensure_ascii : bool, optional
|
||||
Whether to escape non-ASCII characters in JSON serialization for prompts. Defaults to True.
|
||||
Set to False to preserve non-ASCII characters (e.g., Korean, Japanese, Chinese) in their
|
||||
original form, making them readable in LLM logs and improving model understanding.
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
|
@ -184,6 +189,7 @@ class Graphiti:
|
|||
|
||||
self.store_raw_episode_content = store_raw_episode_content
|
||||
self.max_coroutines = max_coroutines
|
||||
self.ensure_ascii = ensure_ascii
|
||||
if llm_client:
|
||||
self.llm_client = llm_client
|
||||
else:
|
||||
|
|
@ -202,6 +208,7 @@ class Graphiti:
|
|||
llm_client=self.llm_client,
|
||||
embedder=self.embedder,
|
||||
cross_encoder=self.cross_encoder,
|
||||
ensure_ascii=self.ensure_ascii,
|
||||
)
|
||||
|
||||
# Capture telemetry event
|
||||
|
|
@ -541,7 +548,9 @@ class Graphiti:
|
|||
if update_communities:
|
||||
communities, community_edges = await semaphore_gather(
|
||||
*[
|
||||
update_community(self.driver, self.llm_client, self.embedder, node)
|
||||
update_community(
|
||||
self.driver, self.llm_client, self.embedder, node, self.ensure_ascii
|
||||
)
|
||||
for node in nodes
|
||||
],
|
||||
max_coroutines=self.max_coroutines,
|
||||
|
|
@ -1021,6 +1030,8 @@ class Graphiti:
|
|||
entity_edges=[],
|
||||
group_id=edge.group_id,
|
||||
),
|
||||
None,
|
||||
self.ensure_ascii,
|
||||
)
|
||||
|
||||
edges: list[EntityEdge] = [resolved_edge] + invalidated_edges
|
||||
|
|
|
|||
|
|
@ -27,5 +27,6 @@ class GraphitiClients(BaseModel):
|
|||
llm_client: LLMClient
|
||||
embedder: EmbedderClient
|
||||
cross_encoder: CrossEncoderClient
|
||||
ensure_ascii: bool = True
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
|
|
|||
|
|
@ -14,12 +14,12 @@ See the License for the specific language governing permissions and
|
|||
limitations under the License.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Protocol, TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .models import Message, PromptFunction, PromptVersion
|
||||
from .prompt_helpers import to_prompt_json
|
||||
|
||||
|
||||
class EdgeDuplicate(BaseModel):
|
||||
|
|
@ -67,11 +67,11 @@ def edge(context: dict[str, Any]) -> list[Message]:
|
|||
Given the following context, determine whether the New Edge represents any of the edges in the list of Existing Edges.
|
||||
|
||||
<EXISTING EDGES>
|
||||
{json.dumps(context['related_edges'], indent=2)}
|
||||
{to_prompt_json(context['related_edges'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</EXISTING EDGES>
|
||||
|
||||
<NEW EDGE>
|
||||
{json.dumps(context['extracted_edges'], indent=2)}
|
||||
{to_prompt_json(context['extracted_edges'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</NEW EDGE>
|
||||
|
||||
Task:
|
||||
|
|
@ -98,7 +98,7 @@ def edge_list(context: dict[str, Any]) -> list[Message]:
|
|||
Given the following context, find all of the duplicates in a list of facts:
|
||||
|
||||
Facts:
|
||||
{json.dumps(context['edges'], indent=2)}
|
||||
{to_prompt_json(context['edges'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
|
||||
Task:
|
||||
If any facts in Facts is a duplicate of another fact, return a new fact with one of their uuid's.
|
||||
|
|
|
|||
|
|
@ -14,12 +14,12 @@ See the License for the specific language governing permissions and
|
|||
limitations under the License.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Protocol, TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .models import Message, PromptFunction, PromptVersion
|
||||
from .prompt_helpers import to_prompt_json
|
||||
|
||||
|
||||
class NodeDuplicate(BaseModel):
|
||||
|
|
@ -64,20 +64,20 @@ def node(context: dict[str, Any]) -> list[Message]:
|
|||
role='user',
|
||||
content=f"""
|
||||
<PREVIOUS MESSAGES>
|
||||
{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
|
||||
{to_prompt_json([ep for ep in context['previous_episodes']], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</PREVIOUS MESSAGES>
|
||||
<CURRENT MESSAGE>
|
||||
{context['episode_content']}
|
||||
</CURRENT MESSAGE>
|
||||
<NEW ENTITY>
|
||||
{json.dumps(context['extracted_node'], indent=2)}
|
||||
{to_prompt_json(context['extracted_node'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</NEW ENTITY>
|
||||
<ENTITY TYPE DESCRIPTION>
|
||||
{json.dumps(context['entity_type_description'], indent=2)}
|
||||
{to_prompt_json(context['entity_type_description'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</ENTITY TYPE DESCRIPTION>
|
||||
|
||||
<EXISTING ENTITIES>
|
||||
{json.dumps(context['existing_nodes'], indent=2)}
|
||||
{to_prompt_json(context['existing_nodes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</EXISTING ENTITIES>
|
||||
|
||||
Given the above EXISTING ENTITIES and their attributes, MESSAGE, and PREVIOUS MESSAGES; Determine if the NEW ENTITY extracted from the conversation
|
||||
|
|
@ -114,7 +114,7 @@ def nodes(context: dict[str, Any]) -> list[Message]:
|
|||
role='user',
|
||||
content=f"""
|
||||
<PREVIOUS MESSAGES>
|
||||
{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
|
||||
{to_prompt_json([ep for ep in context['previous_episodes']], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</PREVIOUS MESSAGES>
|
||||
<CURRENT MESSAGE>
|
||||
{context['episode_content']}
|
||||
|
|
@ -139,11 +139,11 @@ def nodes(context: dict[str, Any]) -> list[Message]:
|
|||
}}
|
||||
|
||||
<ENTITIES>
|
||||
{json.dumps(context['extracted_nodes'], indent=2)}
|
||||
{to_prompt_json(context['extracted_nodes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</ENTITIES>
|
||||
|
||||
<EXISTING ENTITIES>
|
||||
{json.dumps(context['existing_nodes'], indent=2)}
|
||||
{to_prompt_json(context['existing_nodes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</EXISTING ENTITIES>
|
||||
|
||||
For each of the above ENTITIES, determine if the entity is a duplicate of any of the EXISTING ENTITIES.
|
||||
|
|
@ -180,7 +180,7 @@ def node_list(context: dict[str, Any]) -> list[Message]:
|
|||
Given the following context, deduplicate a list of nodes:
|
||||
|
||||
Nodes:
|
||||
{json.dumps(context['nodes'], indent=2)}
|
||||
{to_prompt_json(context['nodes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
|
||||
Task:
|
||||
1. Group nodes together such that all duplicate nodes are in the same list of uuids
|
||||
|
|
|
|||
|
|
@ -14,12 +14,12 @@ See the License for the specific language governing permissions and
|
|||
limitations under the License.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Protocol, TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .models import Message, PromptFunction, PromptVersion
|
||||
from .prompt_helpers import to_prompt_json
|
||||
|
||||
|
||||
class QueryExpansion(BaseModel):
|
||||
|
|
@ -68,7 +68,7 @@ def query_expansion(context: dict[str, Any]) -> list[Message]:
|
|||
Bob is asking Alice a question, are you able to rephrase the question into a simpler one about Alice in the third person
|
||||
that maintains the relevant context?
|
||||
<QUESTION>
|
||||
{json.dumps(context['query'])}
|
||||
{to_prompt_json(context['query'], ensure_ascii=context.get('ensure_ascii', True))}
|
||||
</QUESTION>
|
||||
"""
|
||||
return [
|
||||
|
|
@ -84,10 +84,10 @@ def qa_prompt(context: dict[str, Any]) -> list[Message]:
|
|||
Your task is to briefly answer the question in the way that you think Alice would answer the question.
|
||||
You are given the following entity summaries and facts to help you determine the answer to your question.
|
||||
<ENTITY_SUMMARIES>
|
||||
{json.dumps(context['entity_summaries'])}
|
||||
{to_prompt_json(context['entity_summaries'], ensure_ascii=context.get('ensure_ascii', True))}
|
||||
</ENTITY_SUMMARIES>
|
||||
<FACTS>
|
||||
{json.dumps(context['facts'])}
|
||||
{to_prompt_json(context['facts'], ensure_ascii=context.get('ensure_ascii', True))}
|
||||
</FACTS>
|
||||
<QUESTION>
|
||||
{context['query']}
|
||||
|
|
|
|||
|
|
@ -14,12 +14,12 @@ See the License for the specific language governing permissions and
|
|||
limitations under the License.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Protocol, TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .models import Message, PromptFunction, PromptVersion
|
||||
from .prompt_helpers import to_prompt_json
|
||||
|
||||
|
||||
class Edge(BaseModel):
|
||||
|
|
@ -73,7 +73,7 @@ def edge(context: dict[str, Any]) -> list[Message]:
|
|||
</FACT TYPES>
|
||||
|
||||
<PREVIOUS_MESSAGES>
|
||||
{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
|
||||
{to_prompt_json([ep for ep in context['previous_episodes']], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</PREVIOUS_MESSAGES>
|
||||
|
||||
<CURRENT_MESSAGE>
|
||||
|
|
@ -132,7 +132,7 @@ def reflexion(context: dict[str, Any]) -> list[Message]:
|
|||
|
||||
user_prompt = f"""
|
||||
<PREVIOUS MESSAGES>
|
||||
{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
|
||||
{to_prompt_json([ep for ep in context['previous_episodes']], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</PREVIOUS MESSAGES>
|
||||
<CURRENT MESSAGE>
|
||||
{context['episode_content']}
|
||||
|
|
@ -166,7 +166,7 @@ def extract_attributes(context: dict[str, Any]) -> list[Message]:
|
|||
content=f"""
|
||||
|
||||
<MESSAGE>
|
||||
{json.dumps(context['episode_content'], indent=2)}
|
||||
{to_prompt_json(context['episode_content'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</MESSAGE>
|
||||
<REFERENCE TIME>
|
||||
{context['reference_time']}
|
||||
|
|
|
|||
|
|
@ -14,12 +14,12 @@ See the License for the specific language governing permissions and
|
|||
limitations under the License.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Protocol, TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .models import Message, PromptFunction, PromptVersion
|
||||
from .prompt_helpers import to_prompt_json
|
||||
|
||||
|
||||
class ExtractedEntity(BaseModel):
|
||||
|
|
@ -89,7 +89,7 @@ def extract_message(context: dict[str, Any]) -> list[Message]:
|
|||
</ENTITY TYPES>
|
||||
|
||||
<PREVIOUS MESSAGES>
|
||||
{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
|
||||
{to_prompt_json([ep for ep in context['previous_episodes']], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</PREVIOUS MESSAGES>
|
||||
|
||||
<CURRENT MESSAGE>
|
||||
|
|
@ -196,7 +196,7 @@ def reflexion(context: dict[str, Any]) -> list[Message]:
|
|||
|
||||
user_prompt = f"""
|
||||
<PREVIOUS MESSAGES>
|
||||
{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
|
||||
{to_prompt_json([ep for ep in context['previous_episodes']], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</PREVIOUS MESSAGES>
|
||||
<CURRENT MESSAGE>
|
||||
{context['episode_content']}
|
||||
|
|
@ -220,7 +220,7 @@ def classify_nodes(context: dict[str, Any]) -> list[Message]:
|
|||
|
||||
user_prompt = f"""
|
||||
<PREVIOUS MESSAGES>
|
||||
{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
|
||||
{to_prompt_json([ep for ep in context['previous_episodes']], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</PREVIOUS MESSAGES>
|
||||
<CURRENT MESSAGE>
|
||||
{context['episode_content']}
|
||||
|
|
@ -258,8 +258,8 @@ def extract_attributes(context: dict[str, Any]) -> list[Message]:
|
|||
content=f"""
|
||||
|
||||
<MESSAGES>
|
||||
{json.dumps(context['previous_episodes'], indent=2)}
|
||||
{json.dumps(context['episode_content'], indent=2)}
|
||||
{to_prompt_json(context['previous_episodes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
{to_prompt_json(context['episode_content'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</MESSAGES>
|
||||
|
||||
Given the above MESSAGES and the following ENTITY, update any of its attributes based on the information provided
|
||||
|
|
@ -288,8 +288,8 @@ def extract_summary(context: dict[str, Any]) -> list[Message]:
|
|||
content=f"""
|
||||
|
||||
<MESSAGES>
|
||||
{json.dumps(context['previous_episodes'], indent=2)}
|
||||
{json.dumps(context['episode_content'], indent=2)}
|
||||
{to_prompt_json(context['previous_episodes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
{to_prompt_json(context['episode_content'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</MESSAGES>
|
||||
|
||||
Given the above MESSAGES and the following ENTITY, update the summary that combines relevant information about the entity
|
||||
|
|
|
|||
|
|
@ -1 +1,24 @@
|
|||
import json
|
||||
from typing import Any
|
||||
|
||||
DO_NOT_ESCAPE_UNICODE = '\nDo not escape unicode characters.\n'
|
||||
|
||||
|
||||
def to_prompt_json(data: Any, ensure_ascii: bool = True, indent: int = 2) -> str:
|
||||
"""
|
||||
Serialize data to JSON for use in prompts.
|
||||
|
||||
Args:
|
||||
data: The data to serialize
|
||||
ensure_ascii: If True, escape non-ASCII characters. If False, preserve them.
|
||||
indent: Number of spaces for indentation
|
||||
|
||||
Returns:
|
||||
JSON string representation of the data
|
||||
|
||||
Notes:
|
||||
When ensure_ascii=False, non-ASCII characters (e.g., Korean, Japanese, Chinese)
|
||||
are preserved in their original form in the prompt, making them readable
|
||||
in LLM logs and improving model understanding.
|
||||
"""
|
||||
return json.dumps(data, ensure_ascii=ensure_ascii, indent=indent)
|
||||
|
|
|
|||
|
|
@ -14,12 +14,12 @@ See the License for the specific language governing permissions and
|
|||
limitations under the License.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Protocol, TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .models import Message, PromptFunction, PromptVersion
|
||||
from .prompt_helpers import to_prompt_json
|
||||
|
||||
|
||||
class Summary(BaseModel):
|
||||
|
|
@ -59,7 +59,7 @@ def summarize_pair(context: dict[str, Any]) -> list[Message]:
|
|||
Summaries must be under 250 words.
|
||||
|
||||
Summaries:
|
||||
{json.dumps(context['node_summaries'], indent=2)}
|
||||
{to_prompt_json(context['node_summaries'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
""",
|
||||
),
|
||||
]
|
||||
|
|
@ -76,8 +76,8 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
|
|||
content=f"""
|
||||
|
||||
<MESSAGES>
|
||||
{json.dumps(context['previous_episodes'], indent=2)}
|
||||
{json.dumps(context['episode_content'], indent=2)}
|
||||
{to_prompt_json(context['previous_episodes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
{to_prompt_json(context['episode_content'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</MESSAGES>
|
||||
|
||||
Given the above MESSAGES and the following ENTITY name, create a summary for the ENTITY. Your summary must only use
|
||||
|
|
@ -100,7 +100,7 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
|
|||
</ENTITY CONTEXT>
|
||||
|
||||
<ATTRIBUTES>
|
||||
{json.dumps(context['attributes'], indent=2)}
|
||||
{to_prompt_json(context['attributes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
</ATTRIBUTES>
|
||||
""",
|
||||
),
|
||||
|
|
@ -120,7 +120,7 @@ def summary_description(context: dict[str, Any]) -> list[Message]:
|
|||
Summaries must be under 250 words.
|
||||
|
||||
Summary:
|
||||
{json.dumps(context['summary'], indent=2)}
|
||||
{to_prompt_json(context['summary'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
|
||||
""",
|
||||
),
|
||||
]
|
||||
|
|
|
|||
|
|
@ -14,9 +14,8 @@ See the License for the specific language governing permissions and
|
|||
limitations under the License.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
from graphiti_core.edges import EntityEdge
|
||||
from graphiti_core.prompts.prompt_helpers import to_prompt_json
|
||||
from graphiti_core.search.search_config import SearchResults
|
||||
|
||||
|
||||
|
|
@ -25,7 +24,9 @@ def format_edge_date_range(edge: EntityEdge) -> str:
|
|||
return f'{edge.valid_at if edge.valid_at else "date unknown"} - {(edge.invalid_at if edge.invalid_at else "present")}'
|
||||
|
||||
|
||||
def search_results_to_context_string(search_results: SearchResults) -> str:
|
||||
def search_results_to_context_string(
|
||||
search_results: SearchResults, ensure_ascii: bool = True
|
||||
) -> str:
|
||||
"""Reformats a set of SearchResults into a single string to pass directly to an LLM as context"""
|
||||
fact_json = [
|
||||
{
|
||||
|
|
@ -57,16 +58,16 @@ def search_results_to_context_string(search_results: SearchResults) -> str:
|
|||
These are the most relevant facts and their valid and invalid dates. Facts are considered valid
|
||||
between their valid_at and invalid_at dates. Facts with an invalid_at date of "Present" are considered valid.
|
||||
<FACTS>
|
||||
{json.dumps(fact_json, indent=12)}
|
||||
{to_prompt_json(fact_json, ensure_ascii=ensure_ascii, indent=12)}
|
||||
</FACTS>
|
||||
<ENTITIES>
|
||||
{json.dumps(entity_json, indent=12)}
|
||||
{to_prompt_json(entity_json, ensure_ascii=ensure_ascii, indent=12)}
|
||||
</ENTITIES>
|
||||
<EPISODES>
|
||||
{json.dumps(episode_json, indent=12)}
|
||||
{to_prompt_json(episode_json, ensure_ascii=ensure_ascii, indent=12)}
|
||||
</EPISODES>
|
||||
<COMMUNITIES>
|
||||
{json.dumps(community_json, indent=12)}
|
||||
{to_prompt_json(community_json, ensure_ascii=ensure_ascii, indent=12)}
|
||||
</COMMUNITIES>
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -343,7 +343,13 @@ async def dedupe_edges_bulk(
|
|||
] = await semaphore_gather(
|
||||
*[
|
||||
resolve_extracted_edge(
|
||||
clients.llm_client, edge, candidates, candidates, episode, edge_types
|
||||
clients.llm_client,
|
||||
edge,
|
||||
candidates,
|
||||
candidates,
|
||||
episode,
|
||||
edge_types,
|
||||
clients.ensure_ascii,
|
||||
)
|
||||
for episode, edge, candidates in dedupe_tuples
|
||||
]
|
||||
|
|
|
|||
|
|
@ -122,9 +122,14 @@ def label_propagation(projection: dict[str, list[Neighbor]]) -> list[list[str]]:
|
|||
return clusters
|
||||
|
||||
|
||||
async def summarize_pair(llm_client: LLMClient, summary_pair: tuple[str, str]) -> str:
|
||||
async def summarize_pair(
|
||||
llm_client: LLMClient, summary_pair: tuple[str, str], ensure_ascii: bool = True
|
||||
) -> str:
|
||||
# Prepare context for LLM
|
||||
context = {'node_summaries': [{'summary': summary} for summary in summary_pair]}
|
||||
context = {
|
||||
'node_summaries': [{'summary': summary} for summary in summary_pair],
|
||||
'ensure_ascii': ensure_ascii,
|
||||
}
|
||||
|
||||
llm_response = await llm_client.generate_response(
|
||||
prompt_library.summarize_nodes.summarize_pair(context), response_model=Summary
|
||||
|
|
@ -135,8 +140,13 @@ async def summarize_pair(llm_client: LLMClient, summary_pair: tuple[str, str]) -
|
|||
return pair_summary
|
||||
|
||||
|
||||
async def generate_summary_description(llm_client: LLMClient, summary: str) -> str:
|
||||
context = {'summary': summary}
|
||||
async def generate_summary_description(
|
||||
llm_client: LLMClient, summary: str, ensure_ascii: bool = True
|
||||
) -> str:
|
||||
context = {
|
||||
'summary': summary,
|
||||
'ensure_ascii': ensure_ascii,
|
||||
}
|
||||
|
||||
llm_response = await llm_client.generate_response(
|
||||
prompt_library.summarize_nodes.summary_description(context),
|
||||
|
|
@ -149,7 +159,7 @@ async def generate_summary_description(llm_client: LLMClient, summary: str) -> s
|
|||
|
||||
|
||||
async def build_community(
|
||||
llm_client: LLMClient, community_cluster: list[EntityNode]
|
||||
llm_client: LLMClient, community_cluster: list[EntityNode], ensure_ascii: bool = True
|
||||
) -> tuple[CommunityNode, list[CommunityEdge]]:
|
||||
summaries = [entity.summary for entity in community_cluster]
|
||||
length = len(summaries)
|
||||
|
|
@ -161,7 +171,9 @@ async def build_community(
|
|||
new_summaries: list[str] = list(
|
||||
await semaphore_gather(
|
||||
*[
|
||||
summarize_pair(llm_client, (str(left_summary), str(right_summary)))
|
||||
summarize_pair(
|
||||
llm_client, (str(left_summary), str(right_summary)), ensure_ascii
|
||||
)
|
||||
for left_summary, right_summary in zip(
|
||||
summaries[: int(length / 2)], summaries[int(length / 2) :], strict=False
|
||||
)
|
||||
|
|
@ -174,7 +186,7 @@ async def build_community(
|
|||
length = len(summaries)
|
||||
|
||||
summary = summaries[0]
|
||||
name = await generate_summary_description(llm_client, summary)
|
||||
name = await generate_summary_description(llm_client, summary, ensure_ascii)
|
||||
now = utc_now()
|
||||
community_node = CommunityNode(
|
||||
name=name,
|
||||
|
|
@ -191,7 +203,10 @@ async def build_community(
|
|||
|
||||
|
||||
async def build_communities(
|
||||
driver: GraphDriver, llm_client: LLMClient, group_ids: list[str] | None
|
||||
driver: GraphDriver,
|
||||
llm_client: LLMClient,
|
||||
group_ids: list[str] | None,
|
||||
ensure_ascii: bool = True,
|
||||
) -> tuple[list[CommunityNode], list[CommunityEdge]]:
|
||||
community_clusters = await get_community_clusters(driver, group_ids)
|
||||
|
||||
|
|
@ -199,7 +214,7 @@ async def build_communities(
|
|||
|
||||
async def limited_build_community(cluster):
|
||||
async with semaphore:
|
||||
return await build_community(llm_client, cluster)
|
||||
return await build_community(llm_client, cluster, ensure_ascii)
|
||||
|
||||
communities: list[tuple[CommunityNode, list[CommunityEdge]]] = list(
|
||||
await semaphore_gather(
|
||||
|
|
@ -285,15 +300,21 @@ async def determine_entity_community(
|
|||
|
||||
|
||||
async def update_community(
|
||||
driver: GraphDriver, llm_client: LLMClient, embedder: EmbedderClient, entity: EntityNode
|
||||
driver: GraphDriver,
|
||||
llm_client: LLMClient,
|
||||
embedder: EmbedderClient,
|
||||
entity: EntityNode,
|
||||
ensure_ascii: bool = True,
|
||||
) -> tuple[list[CommunityNode], list[CommunityEdge]]:
|
||||
community, is_new = await determine_entity_community(driver, entity)
|
||||
|
||||
if community is None:
|
||||
return [], []
|
||||
|
||||
new_summary = await summarize_pair(llm_client, (entity.summary, community.summary))
|
||||
new_name = await generate_summary_description(llm_client, new_summary)
|
||||
new_summary = await summarize_pair(
|
||||
llm_client, (entity.summary, community.summary), ensure_ascii
|
||||
)
|
||||
new_name = await generate_summary_description(llm_client, new_summary, ensure_ascii)
|
||||
|
||||
community.summary = new_summary
|
||||
community.name = new_name
|
||||
|
|
|
|||
|
|
@ -151,6 +151,7 @@ async def extract_edges(
|
|||
'reference_time': episode.valid_at,
|
||||
'edge_types': edge_types_context,
|
||||
'custom_prompt': '',
|
||||
'ensure_ascii': clients.ensure_ascii,
|
||||
}
|
||||
|
||||
facts_missed = True
|
||||
|
|
@ -311,6 +312,7 @@ async def resolve_extracted_edges(
|
|||
existing_edges,
|
||||
episode,
|
||||
extracted_edge_types,
|
||||
clients.ensure_ascii,
|
||||
)
|
||||
for extracted_edge, related_edges, existing_edges, extracted_edge_types in zip(
|
||||
extracted_edges,
|
||||
|
|
@ -382,6 +384,7 @@ async def resolve_extracted_edge(
|
|||
existing_edges: list[EntityEdge],
|
||||
episode: EpisodicNode,
|
||||
edge_types: dict[str, type[BaseModel]] | None = None,
|
||||
ensure_ascii: bool = True,
|
||||
) -> tuple[EntityEdge, list[EntityEdge], list[EntityEdge]]:
|
||||
if len(related_edges) == 0 and len(existing_edges) == 0:
|
||||
return extracted_edge, [], []
|
||||
|
|
@ -415,6 +418,7 @@ async def resolve_extracted_edge(
|
|||
'new_edge': extracted_edge.fact,
|
||||
'edge_invalidation_candidates': invalidation_edge_candidates_context,
|
||||
'edge_types': edge_types_context,
|
||||
'ensure_ascii': ensure_ascii,
|
||||
}
|
||||
|
||||
llm_response = await llm_client.generate_response(
|
||||
|
|
@ -449,6 +453,7 @@ async def resolve_extracted_edge(
|
|||
'episode_content': episode.content,
|
||||
'reference_time': episode.valid_at,
|
||||
'fact': resolved_edge.fact,
|
||||
'ensure_ascii': ensure_ascii,
|
||||
}
|
||||
|
||||
edge_model = edge_types.get(fact_type)
|
||||
|
|
|
|||
|
|
@ -48,12 +48,14 @@ async def extract_nodes_reflexion(
|
|||
episode: EpisodicNode,
|
||||
previous_episodes: list[EpisodicNode],
|
||||
node_names: list[str],
|
||||
ensure_ascii: bool = True,
|
||||
) -> list[str]:
|
||||
# Prepare context for LLM
|
||||
context = {
|
||||
'episode_content': episode.content,
|
||||
'previous_episodes': [ep.content for ep in previous_episodes],
|
||||
'extracted_entities': node_names,
|
||||
'ensure_ascii': ensure_ascii,
|
||||
}
|
||||
|
||||
llm_response = await llm_client.generate_response(
|
||||
|
|
@ -106,6 +108,7 @@ async def extract_nodes(
|
|||
'custom_prompt': custom_prompt,
|
||||
'entity_types': entity_types_context,
|
||||
'source_description': episode.source_description,
|
||||
'ensure_ascii': clients.ensure_ascii,
|
||||
}
|
||||
|
||||
while entities_missed and reflexion_iterations <= MAX_REFLEXION_ITERATIONS:
|
||||
|
|
@ -134,6 +137,7 @@ async def extract_nodes(
|
|||
episode,
|
||||
previous_episodes,
|
||||
[entity.name for entity in extracted_entities],
|
||||
clients.ensure_ascii,
|
||||
)
|
||||
|
||||
entities_missed = len(missing_entities) != 0
|
||||
|
|
@ -244,6 +248,7 @@ async def resolve_extracted_nodes(
|
|||
'previous_episodes': [ep.content for ep in previous_episodes]
|
||||
if previous_episodes is not None
|
||||
else [],
|
||||
'ensure_ascii': clients.ensure_ascii,
|
||||
}
|
||||
|
||||
llm_response = await llm_client.generate_response(
|
||||
|
|
@ -309,6 +314,7 @@ async def extract_attributes_from_nodes(
|
|||
entity_types.get(next((item for item in node.labels if item != 'Entity'), ''))
|
||||
if entity_types is not None
|
||||
else None,
|
||||
clients.ensure_ascii,
|
||||
)
|
||||
for node in nodes
|
||||
]
|
||||
|
|
@ -325,6 +331,7 @@ async def extract_attributes_from_node(
|
|||
episode: EpisodicNode | None = None,
|
||||
previous_episodes: list[EpisodicNode] | None = None,
|
||||
entity_type: type[BaseModel] | None = None,
|
||||
ensure_ascii: bool = True,
|
||||
) -> EntityNode:
|
||||
node_context: dict[str, Any] = {
|
||||
'name': node.name,
|
||||
|
|
@ -339,6 +346,7 @@ async def extract_attributes_from_node(
|
|||
'previous_episodes': [ep.content for ep in previous_episodes]
|
||||
if previous_episodes is not None
|
||||
else [],
|
||||
'ensure_ascii': ensure_ascii,
|
||||
}
|
||||
|
||||
summary_context: dict[str, Any] = {
|
||||
|
|
@ -347,6 +355,7 @@ async def extract_attributes_from_node(
|
|||
'previous_episodes': [ep.content for ep in previous_episodes]
|
||||
if previous_episodes is not None
|
||||
else [],
|
||||
'ensure_ascii': ensure_ascii,
|
||||
}
|
||||
|
||||
llm_response = (
|
||||
|
|
|
|||
|
|
@ -35,12 +35,14 @@ async def extract_edge_dates(
|
|||
edge: EntityEdge,
|
||||
current_episode: EpisodicNode,
|
||||
previous_episodes: list[EpisodicNode],
|
||||
ensure_ascii: bool = True,
|
||||
) -> tuple[datetime | None, datetime | None]:
|
||||
context = {
|
||||
'edge_fact': edge.fact,
|
||||
'current_episode': current_episode.content,
|
||||
'previous_episodes': [ep.content for ep in previous_episodes],
|
||||
'reference_timestamp': current_episode.valid_at.isoformat(),
|
||||
'ensure_ascii': ensure_ascii,
|
||||
}
|
||||
llm_response = await llm_client.generate_response(
|
||||
prompt_library.extract_edge_dates.v1(context), response_model=EdgeDates
|
||||
|
|
@ -70,7 +72,10 @@ async def extract_edge_dates(
|
|||
|
||||
|
||||
async def get_edge_contradictions(
|
||||
llm_client: LLMClient, new_edge: EntityEdge, existing_edges: list[EntityEdge]
|
||||
llm_client: LLMClient,
|
||||
new_edge: EntityEdge,
|
||||
existing_edges: list[EntityEdge],
|
||||
ensure_ascii: bool = True,
|
||||
) -> list[EntityEdge]:
|
||||
start = time()
|
||||
|
||||
|
|
@ -79,7 +84,11 @@ async def get_edge_contradictions(
|
|||
{'id': i, 'fact': existing_edge.fact} for i, existing_edge in enumerate(existing_edges)
|
||||
]
|
||||
|
||||
context = {'new_edge': new_edge_context, 'existing_edges': existing_edge_context}
|
||||
context = {
|
||||
'new_edge': new_edge_context,
|
||||
'existing_edges': existing_edge_context,
|
||||
'ensure_ascii': ensure_ascii,
|
||||
}
|
||||
|
||||
llm_response = await llm_client.generate_response(
|
||||
prompt_library.invalidate_edges.v2(context),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue