Update function.py

function.py

@@ -3,7 +3,7 @@ from langchain.chains import LLMChain, SimpleSequentialChain
 from langchain.memory import SimpleMemory
 from langchain.prompts import PromptTemplate
 from langchain.llms import OpenAI
-from langchain.prompts import ChatPromptTemplate
+
 
 
 
@@ -16,12 +16,10 @@ def ai_generated_content(topic, audience, benefit, date_time):
     memory = SimpleMemory(memory={"topic": topic, "audience": audience, "benefit": benefit, "date_time": date_time})
 
     # Define prompt templates using stored memory variables
-    headline_template = ChatPromptTemplate.from_template("Create a headline for {audience} interested in {topic} that highlights the benefit of {benefit}."
+    headline_template = PromptTemplate(
+        input_variables=["topic", "audience", "benefit"],
+        template="Create a headline for {audience} interested in {topic} that highlights the benefit of {benefit}."
     )
-        
-    #     input_variables=["topic", "audience", "benefit"],
-    #     template="Create a headline for {audience} interested in {topic} that highlights the benefit of {benefit}."
-    # )
 
     subheadline_template = PromptTemplate(
         input_variables=["audience", "benefit"],
@@ -29,8 +27,8 @@ def ai_generated_content(topic, audience, benefit, date_time):
     )
 
     # Create LLMChains for each section
-    headline_chain = LLMChain(llm=OpenAI(model="gpt-4"), prompt=headline_template, memory=memory)
-    subheadline_chain = LLMChain(llm=OpenAI(model="gpt-4"), prompt=subheadline_template, memory=memory)
+    headline_chain = LLMChain(llm=llm, prompt=headline_template, memory=memory)
+    subheadline_chain = LLMChain(llm=llm), prompt=subheadline_template, memory=memory)
 
     # Run chains and retrieve results
     headline = headline_chain.run({    "topic": topic,