Security added for query reqs via LLM

app.py

@@ -294,6 +294,8 @@ def download_tdocs(req: DownloadRequest):
 async def gen_reqs(req: RequirementsRequest, background_tasks: BackgroundTasks):
     documents = req.documents
     n_docs = len(documents)
+    def prompt(doc_id, full):
+        return f"Here's the document whose ID is {doc_id} : {full}\n\nExtract all requirements and group them by context, returning a list of objects where each object includes a document ID, a concise description of the context where the requirements apply (not a chapter title or copied text), and a list of associated requirements; always return the result as a list, even if only one context is found. Remove the errors"
     
     async def process_document(doc):
         doc_id = doc.document
@@ -309,7 +311,7 @@ async def gen_reqs(req: RequirementsRequest, background_tasks: BackgroundTasks):
             async with limiter_mapping[model_used]:
                 resp_ai = await llm_router.acompletion(
                     model=model_used,
-                    messages=[{"role":"user","content": f"Here's the document whose ID is {doc_id} : {full}\n\nExtract all requirements and group them by context, returning a list of objects where each object includes a document ID, a concise description of the context where the requirements apply (not a chapter title or copied text), and a list of associated requirements; always return the result as a list, even if only one context is found. Remove the errors"}],
+                    messages=[{"role":"user","content": prompt(doc_id, full)}],
                     response_format=RequirementsResponse
                 )
             return RequirementsResponse.model_validate_json(resp_ai.choices[0].message.content).requirements
@@ -320,7 +322,7 @@ async def gen_reqs(req: RequirementsRequest, background_tasks: BackgroundTasks):
                     async with limiter_mapping[model_used]:
                         resp_ai = await llm_router.acompletion(
                             model=model_used,
-                            messages=[{"role":"user","content": f"Here's the document whose ID is {doc_id} : {full}\n\nExtract all requirements and group them by context, returning a list of objects where each object includes a document ID, a concise description of the context where the requirements apply (not a chapter title or copied text), and a list of associated requirements; always return the result as a list, even if only one context is found. Remove the errors"}],
+                            messages=[{"role":"user","content": prompt(doc_id, full)}],
                             response_format=RequirementsResponse
                         )
                     return RequirementsResponse.model_validate_json(resp_ai.choices[0].message.content).requirements
@@ -357,14 +359,19 @@ def find_requirements_from_problem_description(req: ReqSearchRequest):
     requirements = req.requirements
     query = req.query
 
-    requirements_text = "\n".join([f"[Document: {r.document} | Context: {r.context} | Requirement: {r.requirement}]" for r in requirements])
+    requirements_text = "\n".join([f"[Selection ID: {x} | Document: {r.document} | Context: {r.context} | Requirement: {r.requirement}]" for x, r in enumerate(requirements)])
     
     print("Called the LLM")
     resp_ai = llm_router.completion(
         model="gemini-v2",
-        messages=[{"role":"user","content": f"Given all the requirements : \n {requirements_text} \n and the problem description \"{query}\", return a list of objects each with document ID, context, and requirement for the most relevant requirements that reference or best cover the problem."}],
-        response_format=ReqSearchResponse
+        messages=[{"role":"user","content": f"Given all the requirements : \n {requirements_text} \n and the problem description \"{query}\", return a list of 'Selection ID' for the most relevant corresponding requirements that reference or best cover the problem. If none of the requirements covers the problem, simply return an empty list"}],
+        response_format=ReqSearchLLMResponse
     )
     print("Answered")
+    print(resp_ai.choices[0].message.content)
 
-    return ReqSearchResponse.model_validate_json(resp_ai.choices[0].message.content)
+    out_llm = ReqSearchLLMResponse.model_validate_json(resp_ai.choices[0].message.content).selected
+    if max(out_llm) > len(out_llm) - 1:
+        raise HTTPException(status_code=500, detail="LLM error : Generated a wrong index, please try again.")
+
+    return ReqSearchResponse(requirements=[requirements[i] for i in out_llm])

schemas.py

@@ -37,6 +37,9 @@ class SingleRequirement(BaseModel):
     context: str
     requirement: str
 
+class ReqSearchLLMResponse(BaseModel):
+    selected: List[int]
+
 class ReqSearchRequest(BaseModel):
     query: str
     requirements: List[SingleRequirement]