Finish criticize_solutions endpoint

.gitignore

@@ -1,2 +1,3 @@
 .venv
-__pycache__
+__pycache__
+.env

app.py

@@ -4,9 +4,8 @@ import os
 import sys
 import uvicorn
 from fastapi import FastAPI
-from schemas import CriticizeSolutionsRequest, RequirementInfo, ReqGroupingCategory, ReqGroupingResponse, ReqGroupingRequest, _ReqGroupingCategory, _ReqGroupingOutput
+from schemas import _SolutionCriticismOutput, CriticizeSolutionsRequest, CritiqueResponse, RequirementInfo, ReqGroupingCategory, ReqGroupingResponse, ReqGroupingRequest, _ReqGroupingCategory, _ReqGroupingOutput, SolutionCriticism
 from jinja2 import Environment, FileSystemLoader
-from litellm import acompletion
 from litellm.router import Router
 from dotenv import load_dotenv
 
@@ -57,8 +56,6 @@ async def categorize_reqs(params: ReqGroupingRequest) -> ReqGroupingResponse:
         "max_n_categories": params.max_n_categories,
         "response_schema": _ReqGroupingOutput.model_json_schema()})
 
-    logging.info(req_prompt)
-
     # add system prompt with requirements
     messages.append({"role": "user", "content": req_prompt})
 
@@ -68,10 +65,16 @@ async def categorize_reqs(params: ReqGroupingRequest) -> ReqGroupingResponse:
         output = _ReqGroupingOutput.model_validate_json(
             req_completion.choices[0].message.content)
 
-        # quick check to ensure no requirement was left out by the LLM by checking all IDs are contained in at least a single category
+        # # quick check to ensure no requirement was left out by the LLM by checking all IDs are contained in at least a single category
+        valid_ids_universe = set(range(1, len(params.requirements)))
         assigned_ids = {
             req_id for cat in output.categories for req_id in cat.items}
-        unassigned_ids = set(range(1, len(params.requirements))) - assigned_ids
+
+        # keep only non-hallucinated, valid assigned ids
+        valid_assigned_ids = assigned_ids.intersection(valid_ids_universe)
+
+        # check for remaining requirements assigned to none of the categories
+        unassigned_ids = valid_ids_universe - valid_assigned_ids
 
         if len(unassigned_ids) == 0:
             categories.extend(output.categories)
@@ -99,11 +102,18 @@ async def categorize_reqs(params: ReqGroupingRequest) -> ReqGroupingResponse:
 
 
 @api.post("/criticize_solution")
-async def criticize_solution(params: CriticizeSolutionsRequest) -> str:
-    req_prompt = await prompt_env.get_template("criticize.txt").render_async(solutions=[sol.model_dump() for sol in params.solutions])
-    req_completion = await llm_router.acompletion(model="chat", messages=[{"role": "user", "content": req_prompt}])
-
-    return req_completion.choices[0].message.content
-
+async def criticize_solution(params: CriticizeSolutionsRequest) -> CritiqueResponse:
+    """Criticize the challenges, weaknesses and limitations of the provided solutions."""
+    req_prompt = await prompt_env.get_template("criticize.txt").render_async(**{
+        "solutions": [sol.model_dump() for sol in params.solutions],
+        "response_schema": _SolutionCriticismOutput.model_json_schema()
+    })
+    req_completion = await llm_router.acompletion(model="chat", messages=[{"role": "user", "content": req_prompt}], response_format=_SolutionCriticismOutput)
+    criticism_out = _SolutionCriticismOutput.model_validate_json(
+        req_completion.choices[0].message.content)
+
+    return CritiqueResponse(critiques=[
+        SolutionCriticism(solution=sol, criticism=crit) for (sol, crit) in zip(params.solutions, criticism_out.criticisms)
+    ])
 
 uvicorn.run(api, host="0.0.0.0", port=8000)

prompts/criticize.txt

@@ -13,4 +13,9 @@ Here are the solutions:
 - Solution description: {{solution["Solution Description"]}}
 ---
 {% endfor -%}
-</solutions>
+</solutions>
+
+<response_format>
+Reply in JSON using the following format:
+{{response_schema}}
+</response_format> 

schemas.py

@@ -20,7 +20,7 @@ class ReqGroupingCategory(BaseModel):
         ..., description="List of grouped requirements")
 
 
-class SolutionSearchResult(BaseModel):
+class SolutionModel(BaseModel):
     Context: str
     Requirements: List[str]
     ProblemDescription: str
@@ -57,4 +57,28 @@ class _ReqGroupingOutput(BaseModel):
 # Criticize solution endpoint
 
 class CriticizeSolutionsRequest(BaseModel):
-    solutions: list[SolutionSearchResult]
+    solutions: list[SolutionModel]
+
+
+class _SolutionCriticism(BaseModel):
+    technical_challenges: List[str] = Field(
+        ..., description="Technical challenges encountered by the solution")
+    weaknesses: List[str] = Field(...,
+                                  description="Identified weaknesses of the solution")
+    limitations: List[str] = Field(...,
+                                   description="Identified limitations of the solution")
+
+
+class _SolutionCriticismOutput(BaseModel):
+    criticisms: List[_SolutionCriticism]
+
+# response format
+
+
+class SolutionCriticism(BaseModel):
+    solution: SolutionModel
+    criticism: _SolutionCriticism
+
+
+class CritiqueResponse(BaseModel):
+    critiques: List[SolutionCriticism]