Update main.py

main.py

@@ -10,6 +10,8 @@ from fastapi import Query
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
 from helper_functions_api import md_to_html
+from duckduckgo_search import DDGS
+import time
 
 # Retrieve environment variables
 TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY")
@@ -66,7 +68,12 @@ def generate_report(topic, description):
     md_report = together_response(prompt, model = "meta-llama/Llama-3-70b-chat-hf", SysPrompt = SysPromptMdOffline)
     return md_to_html(md_report)
 
-
+def get_images(query, num_results):
+    time.sleep(0.5)
+    ddgs = DDGS() 
+    imgs = ddgs.images(keywords=query, safesearch="on", max_results=num_results)
+    return imgs
+    
 # Define the app
 app = FastAPI()
 
@@ -101,4 +108,8 @@ async def create_topics(input: TopicInput):
 async def create_report(input: ReportInput):
     report = generate_report(input.topic, input.description)  # You'll need to implement this function
     return {"report": report}
-    
+
+@app.post("/get_images")
+async def get_images(input: ReportInput):
+    images = get_images(input.user_input, input.num_topics)
+    return {"images": images}