Reorganize endpoints + add refine_solution endpoint

app.py

@@ -3,8 +3,8 @@ import logging
 import os
 import sys
 import uvicorn
-from fastapi import FastAPI
-from schemas import _SearchedSolutionModel, _SolutionCriticismOutput, CriticizeSolutionsRequest, CritiqueResponse, RequirementInfo, ReqGroupingCategory, ReqGroupingResponse, ReqGroupingRequest, _ReqGroupingCategory, _ReqGroupingOutput, SolutionCriticism, SolutionModel, SolutionSearchResponse
+from fastapi import APIRouter, FastAPI
+from schemas import _RefinedSolutionModel, _SearchedSolutionModel, _SolutionCriticismOutput, CriticizeSolutionsRequest, CritiqueResponse, RequirementInfo, ReqGroupingCategory, ReqGroupingResponse, ReqGroupingRequest, _ReqGroupingCategory, _ReqGroupingOutput, SolutionCriticism, SolutionModel, SolutionSearchResponse
 from jinja2 import Environment, FileSystemLoader, StrictUndefined
 from litellm.router import Router
 from dotenv import load_dotenv
@@ -44,9 +44,13 @@ prompt_env = Environment(loader=FileSystemLoader(
     'prompts'), enable_async=True, undefined=StrictUndefined)
 
 api = FastAPI(docs_url="/")
+# requirements routes
+requirements_router = APIRouter(prefix="/reqs", tags=["requirements"])
+# solution routes
+solution_router = APIRouter(prefix="/solution", tags=["solution"])
 
 
-@api.post("/categorize_requirements")
+@requirements_router.post("/categorize_requirements")
 async def categorize_reqs(params: ReqGroupingRequest) -> ReqGroupingResponse:
     """Categorize the given service requirements into categories"""
 
@@ -105,34 +109,9 @@ async def categorize_reqs(params: ReqGroupingRequest) -> ReqGroupingResponse:
 
     return ReqGroupingResponse(categories=final_categories)
 
+# ========================================================= Solution Endpoints ===========================================================
 
-@api.post("/criticize_solution", response_model=CritiqueResponse)
-async def criticize_solution(params: CriticizeSolutionsRequest) -> CritiqueResponse:
-    """Criticize the challenges, weaknesses and limitations of the provided solutions."""
-
-    async def __criticize_single(solution: SolutionModel):
-        req_prompt = await prompt_env.get_template("criticize.txt").render_async(**{
-            "solutions": [solution.model_dump()],
-            "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 SolutionCriticism(solution=solution, criticism=criticism_out.criticisms[0])
-
-    critiques = await asyncio.gather(*[__criticize_single(sol) for sol in params.solutions], return_exceptions=False)
-    return CritiqueResponse(critiques=critiques)
-
-
-@api.post("/search_solutions_gemini", response_model=SolutionSearchResponse)
+@solution_router.post("/search_solutions_gemini", response_model=SolutionSearchResponse)
 async def search_solutions(params: ReqGroupingResponse) -> SolutionSearchResponse:
     """Searches solutions using Gemini and grounded on google search"""
 
@@ -187,5 +166,63 @@ async def search_solutions(params: ReqGroupingResponse) -> SolutionSearchRespons
 
     return SolutionSearchResponse(solutions=final_solutions)
 
+@solution_router.post("/criticize_solution", response_model=CritiqueResponse)
+async def criticize_solution(params: CriticizeSolutionsRequest) -> CritiqueResponse:
+    """Criticize the challenges, weaknesses and limitations of the provided solutions."""
+
+    async def __criticize_single(solution: SolutionModel):
+        req_prompt = await prompt_env.get_template("criticize.txt").render_async(**{
+            "solutions": [solution.model_dump()],
+            "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 SolutionCriticism(solution=solution, criticism=criticism_out.criticisms[0])
+
+    critiques = await asyncio.gather(*[__criticize_single(sol) for sol in params.solutions], return_exceptions=False)
+    return CritiqueResponse(critiques=critiques)
+
+
+@solution_router.post("/refine_solutions", response_model=SolutionSearchResponse)
+async def refine_solutions(params: CritiqueResponse) -> SolutionSearchResponse:
+    """Refines the previously critiqued solutions."""
+
+    async def __refine_solution(crit: SolutionCriticism):
+        req_prompt = await prompt_env.get_template("refine_solution.txt").render_async(**{
+            "solution": crit.solution.model_dump(),
+            "criticism": crit.criticism,
+            "response_schema": _RefinedSolutionModel.model_json_schema(),
+        })
+
+        req_completion = await llm_router.acompletion(model="chat", messages=[
+            {"role": "user", "content": req_prompt}
+        ], response_format=_RefinedSolutionModel)
+
+        req_model = _RefinedSolutionModel.model_validate_json(
+            req_completion.choices[0].message.content)
+
+        # copy previous solution model
+        refined_solution = crit.solution.model_copy(deep=True)
+        refined_solution.Problem_Description = req_model.problem_description
+        refined_solution.Solution_Description = req_model.solution_description
+
+        return refined_solution
+
+    refined_solutions = await asyncio.gather(*[__refine_solution(crit) for crit in params.critiques], return_exceptions=False)
+
+    return SolutionSearchResponse(solutions=refined_solutions)
+
+
+api.include_router(requirements_router)
+api.include_router(solution_router)
 
 uvicorn.run(api, host="0.0.0.0", port=8000)

prompts/classify.txt

@@ -16,6 +16,6 @@ Here are the requirements:
 </requirements>
 
 <response_format>
-Reply in JSON using the following format:
+Reply in JSON using the following schema:
 {{response_schema}}
 </response_format> 

prompts/criticize.txt

@@ -16,6 +16,6 @@ Here are the solutions:
 </solutions>
 
 <response_format>
-Reply in JSON using the following format:
+Reply in JSON using the following schema:
 {{response_schema}}
 </response_format> 

prompts/refine_solution.txt

@@ -0,0 +1,33 @@
+<role>You are an expert system designer</role>
+<task>
+Your task is to refine a solution to account for the technical challenges, weaknesses and limitations that were critiqued. 
+No need to include that the solution is refined.
+</task>
+
+<solution>
+Here is the solution:
+
+# Solution Context: 
+{{solution['Context']}}
+
+# Requirements solved by the solution
+{% for req in solution['Requirements'] -%}
+- {{req}}
+{% endfor %}
+
+# Problem description associated to the solution
+{{solution['Problem_Description']}}
+
+# Description of the solution
+{{solution['Solution_Description']}}
+</solution>
+
+<criticisim>
+Here is the criticism:
+{{criticism}}
+</criticism>
+
+<response_format>
+Reply in JSON using the following response schema:
+{{response_schema}}
+</response_format>

prompts/structure_solution.txt

@@ -7,6 +7,6 @@ Here is the solution
 </solution>
 
 <response_format>
-Reply in JSON using the following format:
+Reply in JSON using the following schema:
 {{response_schema}}
 </response_format> 

schemas.py

@@ -113,3 +113,15 @@ class _SearchedSolutionModel(BaseModel):
 class SolutionSearchResponse(BaseModel):
     """Response model for solution search"""
     solutions: list[SolutionModel]
+
+
+# ================================================================= refine solution endpoints
+
+
+class _RefinedSolutionModel(BaseModel):
+    """Internal model used for solution refining"""
+
+    problem_description: str = Field(...,
+                                     description="New description of the problem being solved.")
+    solution_description: str = Field(...,
+                                      description="New detailed description of the solution.")