ImaGen / generate_script.py
fmmkii's picture
Updated app.py
52cdd17
from langchain_openai import ChatOpenAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_community.utilities import GoogleSerperAPIWrapper
import os
def generate_script(prompt, video_length, creativity, openai_api_key, serper_api_key):
"""
Generate a YouTube video title and script based on the provided prompt, duration, and creativity level.
"""
# Ensure the OpenAI API key is set
if not openai_api_key:
raise ValueError("OpenAI API Key is missing.")
os.environ["OPENAI_API_KEY"] = openai_api_key
# Ensure the Serper API key is set
if not serper_api_key:
raise ValueError("Serper API Key is missing.")
os.environ["SERPER_API_KEY"] = serper_api_key
# Initialize the LLM
llm = ChatOpenAI(
model='gpt-4o-mini',
temperature=creativity,
max_tokens=None,
timeout=None,
max_retries=2
)
# Define the title generation prompt
title_template = PromptTemplate(
input_variables=["topic"],
template="Generate a YouTube video title for the topic '{topic}'."
)
title_chain = LLMChain(llm=llm, prompt=title_template, verbose=False)
# Define the script generation prompt
script_template = PromptTemplate(
input_variables=["title", "duration", "search_data"],
template=(
"Create a YouTube script with the title '{title}' that lasts approximately {duration} minutes. "
"Incorporate the following research information: {search_data}."
)
)
script_chain = LLMChain(llm=llm, prompt=script_template, verbose=False)
# Fetch relevant search data using Google Serper
search_tool = GoogleSerperAPIWrapper()
search_data = search_tool.results(prompt)
# Generate the video title ONCE
title = title_chain.run({"topic": prompt})
# Generate the video script using the title
script = script_chain.run({
"title": title,
"duration": video_length,
"search_data": search_data,
})
return title, script, search_data