From 207af40f54943ccee47256d25156215e01d4456c Mon Sep 17 00:00:00 2001
From: Ghazi-raad <raad.m.ghazi@gmail.com>
Date: Wed, 26 Nov 2025 21:56:25 +0000
Subject: [PATCH] Optimize for OpenAI Prompt Caching: Restructure entity
 extraction prompts

- Remove input_text from entity_extraction_system_prompt to enable caching
- Move input_text to entity_extraction_user_prompt for per-chunk variability
- Update operate.py to format system prompt once without input_text
- Format user prompts with input_text for each chunk

This enables OpenAI's automatic prompt caching (50% discount on cached tokens):
- ~1300 token system message cached and reused for ALL chunks
- Only ~150 token user message varies per chunk
- Expected 45% cost reduction on prompt tokens during indexing
- 2-3x faster response times from cached prompts

Fixes #2355
---
 lightrag/operate.py |  4 +++-
 lightrag/prompt.py  | 10 ++++++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/lightrag/operate.py b/lightrag/operate.py
index 5f824af0..521e8b03 100644
--- a/lightrag/operate.py
+++ b/lightrag/operate.py
@@ -2832,9 +2832,11 @@ async def extract_entities(
         cache_keys_collector = []
 
         # Get initial extraction
+        # Format system prompt once without input_text for OpenAI prompt caching
         entity_extraction_system_prompt = PROMPTS[
             "entity_extraction_system_prompt"
-        ].format(**{**context_base, "input_text": content})
+        ].format(**context_base)
+        # Format user prompts with input_text for each chunk
         entity_extraction_user_prompt = PROMPTS["entity_extraction_user_prompt"].format(
             **{**context_base, "input_text": content}
         )
diff --git a/lightrag/prompt.py b/lightrag/prompt.py
index bf514fe8..cfa5b76f 100644
--- a/lightrag/prompt.py
+++ b/lightrag/prompt.py
@@ -62,14 +62,16 @@ You are a Knowledge Graph Specialist responsible for extracting entities and rel
 ---Real Data to be Processed---
 <Input>
 Entity_types: [{entity_types}]
+"""
+
+PROMPTS["entity_extraction_user_prompt"] = """---Task---
+Extract entities and relationships from the following input text.
+
+---Input Text---
 Text:
 ```
 {input_text}
 ```
-"""
-
-PROMPTS["entity_extraction_user_prompt"] = """---Task---
-Extract entities and relationships from the input text to be processed.
 
 ---Instructions---
 1.  **Strict Adherence to Format:** Strictly adhere to all format requirements for entity and relationship lists, including output order, field delimiters, and proper noun handling, as specified in the system prompt.