Update main.py

main.py

@@ -91,6 +91,10 @@ class TopicInput(BaseModel):
     num_topics: int = Query(default=5, description="Number of subtopics to generate (default: 5)")
     previous_queries: list[str] = Query(default=[], description="List of previous queries for context")
 
+class ImageInput(BaseModel):
+    user_input: str = Query(default="market research", description="input query to output images")
+    num_images: int = Query(default=5, description="Number of image links to output (default: 5)")
+    
 class ReportInput(BaseModel):
     topic: str = Query(default="market research",description="The main topic for the report")
     description: str = Query(default="",description="A brief description of the topic")
@@ -110,6 +114,6 @@ async def create_report(input: ReportInput):
     return {"report": report}
 
 @app.post("/get_images")
-async def get_images(input: TopicInput):
-    images = get_images(input.user_input, input.num_topics)
+async def get_images(input: ImageInput):
+    images = get_images(input.user_input, input.num_images)
     return {"images": images}