44import traceback
55from flask import Flask , render_template , request , jsonify , send_file , Response
66import openai
7+ from together import Together
78from reportlab .lib .pagesizes import letter
89from reportlab .lib import colors
910from reportlab .platypus import SimpleDocTemplate , Paragraph , Spacer
2324progress_queue = Queue ()
2425
2526OPENAI_API_KEY = os .getenv ('OPENAI_API_KEY' )
27+ TOGETHER_API_KEY = os .getenv ('TOGETHER_API_KEY' )
28+
2629if not OPENAI_API_KEY :
2730 raise ValueError ("No OpenAI API key found. Please set the OPENAI_API_KEY environment variable." )
31+ if not TOGETHER_API_KEY :
32+ raise ValueError ("No Together API key found. Please set the TOGETHER_API_KEY environment variable." )
33+
34+ # Initialize Together client
35+ together_client = Together (api_key = TOGETHER_API_KEY )
2836
2937# Initialize SQLite database
3038def init_db ():
@@ -41,34 +49,71 @@ def init_db():
4149
4250init_db ()
4351
44- async def generate_chunk (model , topic , current_word_count , language , is_new_chapter = False ):
52+ async def generate_chunk (api , model , topic , current_word_count , language , is_new_chapter = False ):
4553 if is_new_chapter :
46- prompt = f"Write the beginning of a new chapter for a book about { topic } { language } { current_word_count }
54+ prompt = f"Write a detailed chapter for a book about { topic } { language } { current_word_count } , then write at least 500 words of content ."
4755 else :
48- prompt = f"Continue writing a book about { topic } { language } { current_word_count } Make sure the narrative flows smoothly from the previous section."
56+ prompt = f"Continue writing a detailed book about { topic } { language } { current_word_count } Write at least 500 words, ensuring the narrative flows smoothly from the previous section."
4957
5058 try :
51- async with aiohttp .ClientSession () as session :
52- async with session .post (
53- "https://api.openai.com/v1/chat/completions" ,
54- headers = {"Authorization" : f"Bearer { OPENAI_API_KEY } },
55- json = {
56- "model" : model ,
57- "messages" : [
58- {"role" : "system" , "content" : f"You are an author writing a book in { language } },
59- {"role" : "user" , "content" : prompt }
59+ if api == 'openai' :
60+ async with aiohttp .ClientSession () as session :
61+ async with session .post (
62+ "https://api.openai.com/v1/chat/completions" ,
63+ headers = {"Authorization" : f"Bearer { OPENAI_API_KEY } },
64+ json = {
65+ "model" : model ,
66+ "messages" : [
67+ {"role" : "system" , "content" : f"You are an author writing a book in { language } },
68+ {"role" : "user" , "content" : prompt }
69+ ],
70+ "max_tokens" : 3000
71+ }
72+ ) as response :
73+ result = await response .json ()
74+ if 'choices' not in result or len (result ['choices' ]) == 0 :
75+ raise ValueError (f"Unexpected API response: { result } )
76+ return result ['choices' ][0 ]['message' ]['content' ].strip ()
77+ elif api == 'together' :
78+ response = together_client .chat .completions .create (
79+ model = model ,
80+ messages = [
81+ {"role" : "system" , "content" : f"You are an author writing a detailed book in { language } },
82+ {"role" : "user" , "content" : prompt }
83+ ],
84+ max_tokens = 2000 ,
85+ temperature = 0.7 ,
86+ top_p = 0.9 ,
87+ top_k = 50 ,
88+ repetition_penalty = 1.03 ,
89+ stop = None
90+ )
91+ generated_text = response .choices [0 ].message .content .strip ()
92+
93+ # Ensure minimum word count
94+ while len (generated_text .split ()) < 500 :
95+ additional_prompt = f"Continue the previous text, adding more details and expanding the narrative. Write at least { 500 - len (generated_text .split ())}
96+ additional_response = together_client .chat .completions .create (
97+ model = model ,
98+ messages = [
99+ {"role" : "system" , "content" : f"You are an author writing a detailed book in { language } },
100+ {"role" : "user" , "content" : generated_text },
101+ {"role" : "user" , "content" : additional_prompt }
60102 ],
61- "max_tokens" : 3000
62- }
63- ) as response :
64- result = await response .json ()
65- if 'choices' not in result or len (result ['choices' ]) == 0 :
66- raise ValueError (f"Unexpected API response: { result } )
67- return result ['choices' ][0 ]['message' ]['content' ].strip ()
103+ max_tokens = 1000 ,
104+ temperature = 0.7 ,
105+ top_p = 0.9 ,
106+ top_k = 50 ,
107+ repetition_penalty = 1.03 ,
108+ stop = None
109+ )
110+ generated_text += "\n " + additional_response .choices [0 ].message .content .strip ()
111+
112+ return generated_text
68113 except Exception as e :
69114 print (f"An error occurred: { e } )
70115 await asyncio .sleep (60 )
71- return await generate_chunk (model , topic , current_word_count , language , is_new_chapter )
116+ return await generate_chunk (api , model , topic , current_word_count , language , is_new_chapter )
72117
73118def create_pdf (content , title , language ):
74119 buffer = io .BytesIO ()
@@ -130,6 +175,7 @@ def playground():
130175
131176@app .route ('/generate' , methods = ['POST' ])
132177async def generate_book ():
178+ api = request .form ['api' ] # 'openai' or 'together'
133179 model = request .form ['model' ]
134180 topic = request .form ['topic' ]
135181 language = request .form ['language' ]
@@ -144,9 +190,9 @@ async def generate_book():
144190
145191 if is_new_chapter :
146192 chapter_count += 1
147- task = asyncio .create_task (generate_chunk (model , topic , current_word_count , language , is_new_chapter = True ))
193+ task = asyncio .create_task (generate_chunk (api , model , topic , current_word_count , language , is_new_chapter = True ))
148194 else :
149- task = asyncio .create_task (generate_chunk (model , topic , current_word_count , language ))
195+ task = asyncio .create_task (generate_chunk (api , model , topic , current_word_count , language ))
150196
151197 tasks .append (task )
152198 current_word_count += 500 # Approximate word count per chunk
0 commit comments