Update app.py

app.py

@@ -1,93 +1,85 @@
-# This is a simple script generator that uses OpenAI to generate scripts based on a document summary and content outline.
-# It's a Streamlit app that allows the user to input the document summary, content outline, and system prompts for context and script generation.
-# The app generates a context that captures the essential information from the summary, and then generates a detailed script for each heading in the content outline.
-# The scripts are then displayed in the app.
-
 import os
 import asyncio
 import streamlit as st
-import aiohttp
+from openai import AsyncOpenAI
 
-# This is the OpenAI API key and URL.
-API_KEY = os.environ["OPENAI_API_KEY"]
-API_URL = "https://api.openai.com/v1/chat/completions"
+# OpenAI API key setup
+api_key = os.environ.get("OPENAI_API_KEY")
+if not api_key:
+    st.error("OpenAI API key not found. Please set the OPENAI_API_KEY environment variable.")
+    st.stop()
 
+client = AsyncOpenAI(api_key=api_key)
 
-# 1st Step: Generate Context
-# まず初めに、コンテンツの前提条件を作るためにcontextを生成します。
-async def generate_context(summaries, system_prompt_prefix, system_prompt_suffix):
-    prompt = f"{system_prompt_prefix}\n\nYou are an AI assistant. Based on the document summary below, please provide a concise context that captures the essential information:\n\n{summaries}\n\n{system_prompt_suffix}"
+async def generate_content(prompt, model="gpt-4-turbo-preview", max_tokens=2048):
+    try:
+        response = await client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=max_tokens
+        )
+        return response.choices[0].message.content.strip()
+    except Exception as e:
+        st.error(f"Error generating content: {str(e)}")
+        return None
 
-    async with aiohttp.ClientSession() as session:
-        payload = {
-            "model": "gpt-4-turbo-preview",
-            "messages": [
-                        {"role": "system", "content": "You are an AI assistant. Based on the document summary below, please provide a concise context that captures the essential information:"},
-                        {"role": "user", "content": prompt}
-                        ],
-            "temperature": 0.7,
-            "max_tokens": 2048,
-        }
-        headers = {"Authorization": f"Bearer {API_KEY}"}
-        try:
-            async with session.post(API_URL, json=payload, headers=headers) as response:
-                if response.status != 200:
-                    raise Exception(f"Error generating context: {response.status} - {await response.text()}")
-                result = await response.json()
-                return result["choices"][0]["message"]["content"].strip()
-        except Exception as e:
-            st.error(f"Error generating context: {str(e)}")
-            return None
+async def generate_context(summaries, system_prompt_prefix, system_prompt_suffix):
+    prompt = f"{system_prompt_prefix}\n\n{summaries}\n\n{system_prompt_suffix}"
+    return await generate_content(prompt)
 
 async def generate_script(context, heading, system_prompt_prefix, system_prompt_suffix):
-    prompt = f"{system_prompt_prefix}\n\n{context}\n\nYou are an AI scriptwriter. Based on the context above and the heading provided below, please generate a detailed script of more than 600 characters.\n\nHeading: {heading}\n\n{system_prompt_suffix}"
-
-    async with aiohttp.ClientSession() as session:
-        payload = {
-            "model": "gpt-4-turbo-preview",
-            "messages": [
-                        {"role": "system", "content": f"You are an AI scriptwriter. Your task is to generate a detailed script based on the given context and the provided heading."},
-                        {"role": "user", "content": prompt}
-                        ],
-            "temperature": 0.7,
-            "max_tokens": 4000,
-        }
-        headers = {"Authorization": f"Bearer {API_KEY}"}
-        try:
-            async with session.post(API_URL, json=payload, headers=headers) as response:
-                if response.status != 200:
-                    raise Exception(f"Error generating script: {response.status} - {await response.text()}")
-                result = await response.json()
-                return result["choices"][0]["message"]["content"].strip()
-        except Exception as e:
-            st.error(f"Error generating script: {str(e)}")
-            return None
+    prompt = f"{system_prompt_prefix}\n\n{context}\n\nHeading: {heading}\n\n{system_prompt_suffix}"
+    return await generate_content(prompt, max_tokens=4000)
 
-
-
-async def main():
+def main():
     st.title("Script Generator")
 
-    context_system_prompt_prefix = st.text_input("Enter the prompt prefix for context generation:", value="You are an AI assistant. Your task is to provide a concise context that captures the essential information from the given document summary.")
+    context_system_prompt_prefix = st.text_input(
+        "Enter the prompt prefix for context generation:",
+        value="You are an AI assistant. Your task is to provide a concise context that captures the essential information from the given document summary."
+    )
     summaries = st.text_area("Enter the document summary:", height=200)
-    context_system_prompt_suffix = st.text_input("Enter the prompt suffix for context generation:", value="Focus on the key points and main ideas, and present the context in a clear and organized manner.")
+    context_system_prompt_suffix = st.text_input(
+        "Enter the prompt suffix for context generation:",
+        value="Focus on the key points and main ideas, and present the context in a clear and organized manner."
+    )
 
     outline = st.text_area("Enter the content outline (one heading per line):", height=200)
-    script_system_prompt_prefix = st.text_input("Enter the system prompt prefix for script generation:", value="You are an AI scriptwriter. Your task is to generate a detailed script based on the given context and the provided heading.")
-    script_system_prompt_suffix = st.text_input("Enter the system prompt suffix for script generation:", value="The script should be relevant, informative, engaging, and contain more than 600 characters.")
+    script_system_prompt_prefix = st.text_input(
+        "Enter the system prompt prefix for script generation:",
+        value="You are an AI scriptwriter. Your task is to generate a detailed script based on the given context and the provided heading."
+    )
+    script_system_prompt_suffix = st.text_input(
+        "Enter the system prompt suffix for script generation:",
+        value="The script should be relevant, informative, engaging, and contain more than 600 characters."
+    )
 
     if st.button("Generate Context and Scripts"):
-        with st.spinner("Generating context..."):
-            context = await generate_context(summaries, context_system_prompt_prefix, context_system_prompt_suffix)
-        st.header("Generated Context")
-        st.write(context)
+        if not summaries or not outline:
+            st.error("Please enter both the document summary and content outline.")
+            return
+
+        context = asyncio.run(generate_context(summaries, context_system_prompt_prefix, context_system_prompt_suffix))
+        if context:
+            st.header("Generated Context")
+            st.write(context)
+
+            headings = [h.strip() for h in outline.split("\n") if h.strip()]
+            scripts = []
+            for heading in headings:
+                script = asyncio.run(generate_script(context, heading, script_system_prompt_prefix, script_system_prompt_suffix))
+                if script:
+                    scripts.append((heading, script))
 
-        headings = outline.split("\n")
-        with st.spinner("Generating scripts..."):
-            tasks = [generate_script(context, heading, script_system_prompt_prefix, script_system_prompt_suffix) for heading in headings]
-            scripts = await asyncio.gather(*tasks)
+            if scripts:
+                st.header("Generated Scripts")
+                for heading, script in scripts:
+                    st.subheader(heading)
+                    st.write(script)
+            else:
+                st.error("Failed to generate scripts. Please try again.")
+        else:
+            st.error("Failed to generate context. Please try again.")
 
-        st.header("Generated Scripts")
-        for heading, script in zip(headings, scripts):
-            st.subheader(heading)
-            st.write(script)
+if __name__ == "__main__":
+    main()