Create nutrition.py

playground/demos/nutrition/nutrition.py

@@ -3,32 +3,20 @@ import base64
 import requests
 from dotenv import load_dotenv
 from swarms.models import Anthropic, OpenAIChat
-from swarms.structs import Agent
+from swarms.structs import Flow
 
 # Load environment variables
 load_dotenv()
 openai_api_key = os.getenv("OPENAI_API_KEY")
 
 # Define prompts for various tasks
-MEAL_PLAN_PROMPT = (
-    "Based on the following user preferences: dietary restrictions as"
-    " vegetarian, preferred cuisines as Italian and Indian, a total"
-    " caloric intake of around 2000 calories per day, and an"
-    " exclusion of legumes, create a detailed weekly meal plan."
-    " Include a variety of meals for breakfast, lunch, dinner, and"
-    " optional snacks."
-)
-IMAGE_ANALYSIS_PROMPT = (
-    "Identify the items in this fridge, including their quantities"
-    " and condition."
-)
-
+MEAL_PLAN_PROMPT = "Based on the following user preferences: dietary restrictions as vegetarian, preferred cuisines as Italian and Indian, a total caloric intake of around 2000 calories per day, and an exclusion of legumes, create a detailed weekly meal plan. Include a variety of meals for breakfast, lunch, dinner, and optional snacks."
+IMAGE_ANALYSIS_PROMPT = "Identify the items in this fridge, including their quantities and condition."
 
 # Function to encode image to base64
 def encode_image(image_path):
     with open(image_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode("utf-8")
-
+        return base64.b64encode(image_file.read()).decode('utf-8')
 
 # Initialize Language Model (LLM)
 llm = OpenAIChat(
@@ -36,13 +24,12 @@ llm = OpenAIChat(
     max_tokens=3000,
 )
 
-
 # Function to handle vision tasks
 def create_vision_agent(image_path):
     base64_image = encode_image(image_path)
     headers = {
         "Content-Type": "application/json",
-        "Authorization": f"Bearer {openai_api_key}",
+        "Authorization": f"Bearer {openai_api_key}"
     }
     payload = {
         "model": "gpt-4-vision-preview",
@@ -51,49 +38,30 @@ def create_vision_agent(image_path):
                 "role": "user",
                 "content": [
                     {"type": "text", "text": IMAGE_ANALYSIS_PROMPT},
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": f"data:image/jpeg;base64,{base64_image}"
-                        },
-                    },
-                ],
+                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
+                ]
             }
         ],
-        "max_tokens": 300,
+        "max_tokens": 300
     }
-    response = requests.post(
-        "https://api.openai.com/v1/chat/completions",
-        headers=headers,
-        json=payload,
-    )
+    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
     return response.json()
 
-
 # Function to generate an integrated shopping list considering meal plan and fridge contents
-def generate_integrated_shopping_list(
-    meal_plan_output, image_analysis, user_preferences
-):
+def generate_integrated_shopping_list(meal_plan_output, image_analysis, user_preferences):
     # Prepare the prompt for the LLM
-    fridge_contents = image_analysis["choices"][0]["message"][
-        "content"
-    ]
-    prompt = (
-        f"Based on this meal plan: {meal_plan_output}, and the"
-        f" following items in the fridge: {fridge_contents},"
-        " considering dietary preferences as vegetarian with a"
-        " preference for Italian and Indian cuisines, generate a"
-        " comprehensive shopping list that includes only the items"
-        " needed."
-    )
+    fridge_contents = image_analysis['choices'][0]['message']['content']
+    prompt = (f"Based on this meal plan: {meal_plan_output}, "
+              f"and the following items in the fridge: {fridge_contents}, "
+              f"considering dietary preferences as vegetarian with a preference for Italian and Indian cuisines, "
+              f"generate a comprehensive shopping list that includes only the items needed.")
 
     # Send the prompt to the LLM and return the response
     response = llm(prompt)
     return response  # assuming the response is a string
 
-
 # Define agent for meal planning
-meal_plan_agent = Agent(
+meal_plan_agent = Flow(
     llm=llm,
     sop=MEAL_PLAN_PROMPT,
     max_loops=1,
@@ -106,7 +74,7 @@ user_preferences = {
     "dietary_restrictions": "vegetarian",
     "preferred_cuisines": ["Italian", "Indian"],
     "caloric_intake": 2000,
-    "other notes": "Doesn't eat legumes",
+    "other notes": "Doesn't eat legumes"
 }
 
 # Generate Meal Plan
@@ -118,9 +86,7 @@ meal_plan_output = meal_plan_agent.run(
 image_analysis_output = create_vision_agent("full_fridge.jpg")
 
 # Generate Integrated Shopping List
-integrated_shopping_list = generate_integrated_shopping_list(
-    meal_plan_output, image_analysis_output, user_preferences
-)
+integrated_shopping_list = generate_integrated_shopping_list(meal_plan_output, image_analysis_output, user_preferences)
 
 # Print and save the outputs
 print("Meal Plan:", meal_plan_output)
@@ -128,10 +94,6 @@ print("Integrated Shopping List:", integrated_shopping_list)
 
 with open("nutrition_output.txt", "w") as file:
     file.write("Meal Plan:\n" + meal_plan_output + "\n\n")
-    file.write(
-        "Integrated Shopping List:\n"
-        + integrated_shopping_list
-        + "\n"
-    )
+    file.write("Integrated Shopping List:\n" + integrated_shopping_list + "\n")
 
 print("Outputs have been saved to nutrition_output.txt")