Update app.py

app.py

@@ -11,7 +11,7 @@ app = FastAPI(
     description="Find relevant technologies from a problem",
 )
 
-class InputData(BaseModel):
+class InputProblem(BaseModel):
     problem: str
 
 class InputConstraints(BaseModel):
@@ -27,37 +27,47 @@ class Technology(BaseModel):
     limitations: str
     id: int
 
-class PriorArt(BaseModel):
+class OutputPriorArt(BaseModel):
     """Represents the search of prior art using the technology combinations"""
     content: str
     uris: List
 
+class InputPriorArtConstraints(BaseModel):
+    technologies: List[Technology]
+    constriants: List[InputConstraints]
+    data_type: str
+
+class InputPriorArtProblems(BaseModel):
+    technologies: List[Technology]
+    constriants: List[InputProblem]
+    data_type: str
+
+
 # This schema defines the root structure of the JSON
 class TechnologyData(BaseModel):
     """Represents the top-level object containing a list of technologies."""
     technologies: List[Technology]
 
 @app.post("/process", response_model=TechnologyData)
-async def process(data: InputData):
-    result = process_input(data, global_tech, global_tech_embeddings, "problem")
-    return {"technologies": result}
+async def process(data: InputProblem):
+    result, prior_art = process_input(data, global_tech, global_tech_embeddings, "problem")
+    return {"technologies": result, "prior_art": prior_art}
 
 @app.post("/process-constraints", response_model=TechnologyData)
 async def process_constraints(constraints: InputConstraints):
-    result = process_input(constraints.constraints, global_tech, global_tech_embeddings, "constraints")
-    return {"technologies": result}
+    result, prior_art = process_input(constraints.constraints, global_tech, global_tech_embeddings, "constraints")
+    return {"technologies": result, "prior_art": prior_art}
 
-@app.post("/prior-art-constraints", response_model=PriorArt)
-async def prior_art_search(TechnologyData, InputConstraints, type: str):
-    prior_art = process_prior_art(TechnologyData, InputConstraints, type)
+@app.post("/prior-art-constraints", response_model=OutputPriorArt)
+async def prior_art_search(data: InputPriorArtConstraints):
+    prior_art = process_prior_art(InputPriorArtConstraints.technologies, InputPriorArtConstraints.constriants, InputPriorArtConstraints.data_type)
     return prior_art
 
-@app.post("/prior-art-problems", response_model=PriorArt)
-async def prior_art_search_problems(technologies: TechnologyData, problems: InputData, type: str):
-    prior_art = process_prior_art(technologies, problems, type)
+@app.post("/prior-art-problems", response_model=OutputPriorArt)
+async def prior_art_search(data: InputPriorArtProblems):
+    prior_art = process_prior_art(InputPriorArtProblems.technologies, InputPriorArtProblems.problems, InputPriorArtProblems.data_type)
     return prior_art
 
-
 def make_json_serializable(data):
     if isinstance(data, dict):
         return {k: make_json_serializable(v) for k, v in data.items()}