update app and utils

app.py

@@ -102,9 +102,9 @@ with gr.Blocks(theme=gr.themes.Monochrome(radius_size=gr.themes.sizes.radius_sm)
                 # clear_button = gr.ClearButton(value="Clear Inputs")
                 generate_button = gr.Button("Generate Workout", variant="primary", elem_id="gen-button")
         with gr.Column(scale=1, min_width=800, elem_id="#gen-output"):
-            generation = gr.Markdown(value="workout")
+            generation = gr.Markdown(value="")
 
-    generate_button.click(run, inputs=[gender, level, muscle_group, equipment, duration], outputs=generation)
+    generate_button.click(run, inputs=[gender, level, muscle_group, equipment, duration, vectorstore, system_prompt], outputs=generation)
     # clear_button.click(fn=lambda: [None, None, None, None, None], outputs=[gender, level, muscle_group, equipment, duration])
 
 demo.launch(share=False)

utils.py

@@ -4,7 +4,7 @@ import time
 from langchain.schema.messages import HumanMessage, SystemMessage
 
 
-def retrieve_knowledge(query, k=10, randomize=True):
+def retrieve_knowledge(query, vectorstore, k=10, randomize=True):
     knowledge = [d.page_content.strip() for d in vectorstore.similarity_search(query, k=k)]
     
     if randomize:
@@ -15,7 +15,7 @@ def retrieve_knowledge(query, k=10, randomize=True):
     return knowledge
 
 
-def generate_workout(query, knowledge):
+def generate_workout(system_prompt, query, knowledge):
     messages = [
         SystemMessage(content=system_prompt.format(workout_context=knowledge)),
         HumanMessage(content=query)
@@ -26,10 +26,10 @@ def generate_workout(query, knowledge):
     return response
 
 
-def run(gender, muscle_group, equipment, level, duration, k=5, randomize=True):
+def run(gender, muscle_group, equipment, level, duration, vectorstore, system_prompt, k=5, randomize=True):
     query = f"{duration}-minute {muscle_group} workout for {gender} {level} level {equipment}"
-    knowledge = retrieve_knowledge(query, k, randomize)
-    response = generate_workout(query, knowledge)
+    knowledge = retrieve_knowledge(query, vectorstore, k, randomize)
+    response = generate_workout(system_prompt, query, knowledge)
 
     for i in range(len(response)):
         time.sleep(0.01)