Update summarize.py

summarize.py

@@ -2,6 +2,7 @@ import os
 from langchain.chains.summarize import load_summarize_chain
 from utils import openai_llm
 from dotenv import load_dotenv
+from langchain_core.prompts import PromptTemplate
 
 from langfuse.callback import CallbackHandler
 load_dotenv()
@@ -14,13 +15,35 @@ os.environ["LANGFUSE_HOST"] = os.getenv("LANGFUSE_HOST")
  
 langfuse_handler = CallbackHandler()
 
+map_prompt = """
+Write a detailed summary of the following text.  
+The summary should be in the language of the text.
+Do not miss any important details.
+Mention all the important entities at the end of the summary with a heading 'IMPORTANT ENTITIES' and bullets underneath the heading
 
 
+"{text}"
+DETAILED SUMMARY:
+"""
+map_prompt_template = PromptTemplate(template=map_prompt, input_variables=["text"])
+
+combine_prompt = """
+Write a detailed summary of the following text delimited by triple backquotes.
+Return your response as a detailed paragraph.
+The Response should be in the language of the text
+Mention all the important entities at the end of the summary with a heading 'IMPORTANT ENTITIES' and bullets underneath the heading
+
+```{text}```
+PARAGRAPH SUMMARY:
+"""
+combine_prompt_template = PromptTemplate(template=combine_prompt, input_variables=["text"])
 
 def setup_summary_chain(api_key):
     """Set up a summary chain with a specified LLM."""
     llm = openai_llm(api_key=api_key)
-    return load_summarize_chain(llm=llm, chain_type='map_reduce')
+    return load_summarize_chain(llm=llm, chain_type='map_reduce',
+                                map_prompt=map_prompt_template,
+                                combine_prompt=combine_prompt_template)
 
 
 def summarize_documents(documents, api_key):