Fix off-by-one issue with requirement IDs + add /search_solutions_gemini endpoint

app.py

@@ -4,8 +4,8 @@ import os
 import sys
 import uvicorn
 from fastapi import FastAPI
-from schemas import _SolutionCriticismOutput, CriticizeSolutionsRequest, CritiqueResponse, RequirementInfo, ReqGroupingCategory, ReqGroupingResponse, ReqGroupingRequest, _ReqGroupingCategory, _ReqGroupingOutput, SolutionCriticism
-from jinja2 import Environment, FileSystemLoader
+from schemas import _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
 
@@ -33,12 +33,13 @@ llm_router = Router(model_list=[
             "max_retries": 5
         }
     }
-])
+], cooldown_time=30)
 
 # Jinja2 environment to load prompt templates
-prompt_env = Environment(loader=FileSystemLoader('prompts'), enable_async=True)
+prompt_env = Environment(loader=FileSystemLoader(
+    'prompts'), enable_async=True, undefined=StrictUndefined)
 
-api = FastAPI()
+api = FastAPI(docs_url="/")
 
 
 @api.post("/categorize_requirements")
@@ -65,8 +66,8 @@ 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
-        valid_ids_universe = set(range(1, len(params.requirements)))
+        # 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(0, len(params.requirements)))
         assigned_ids = {
             req_id for cat in output.categories for req_id in cat.items}
 
@@ -101,19 +102,85 @@ async def categorize_reqs(params: ReqGroupingRequest) -> ReqGroupingResponse:
     return ReqGroupingResponse(categories=final_categories)
 
 
-@api.post("/criticize_solution")
+@api.post("/criticize_solution", response_model=CritiqueResponse)
 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)
-    ])
+
+    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)
+async def search_solutions(params: ReqGroupingResponse) -> SolutionSearchResponse:
+    """Searches solutions using Gemini and grounded on google search"""
+
+    async def _search_inner(cat: ReqGroupingCategory) -> SolutionModel:
+        # ================== generate the solution with web grounding
+        req_prompt = await prompt_env.get_template("search_solution.txt").render_async(**{
+            "category": cat.model_dump(),
+        })
+
+        # generate the completion in non-structured mode.
+        # the googleSearch tool enables grounding gemini with google search
+        req_completion = await llm_router.acompletion(model="chat", messages=[
+            {"role": "user", "content": req_prompt}
+        ], tools=[{"googleSearch": {}}])
+
+        # ==================== structure the solution as a json ===================================
+
+        structured_prompt = await prompt_env.get_template("structure_solution.txt").render_async(**{
+            "solution": req_completion.choices[0].message.content,
+            "response_schema": _SearchedSolutionModel.model_json_schema()
+        })
+
+        structured_completion = await llm_router.acompletion(model="chat", messages=[
+            {"role": "user", "content": structured_prompt}
+        ], response_format=_SearchedSolutionModel)
+        solution_model = _SearchedSolutionModel.model_validate_json(
+            structured_completion.choices[0].message.content)
+
+        # ======================== build the final solution object ================================
+
+        # extract the source metadata from the search items
+        sources_metadata = [
+            f'{a["web"]["title"]} - {a["web"]["uri"]}' for a in req_completion["vertex_ai_grounding_metadata"][0]['groundingChunks']]
+
+        final_sol = SolutionModel(
+            Context="",
+            Requirements=[
+                cat.requirements[i].requirement for i in solution_model.requirement_ids
+            ],
+            Problem_Description=solution_model.problem_description,
+            Solution_Description=solution_model.solution_description,
+            References=sources_metadata,
+            Category_Id=cat.id,
+        )
+        return final_sol
+
+    solutions = await asyncio.gather(*[_search_inner(cat) for cat in params.categories], return_exceptions=True)
+    logging.info(solutions)
+    final_solutions = [
+        sol for sol in solutions if not isinstance(sol, Exception)]
+
+    return SolutionSearchResponse(solutions=final_solutions)
+
 
 uvicorn.run(api, host="0.0.0.0", port=8000)

prompts/classify.txt

@@ -11,7 +11,7 @@ For each category indicate which requirements belong in that category using thei
 Here are the requirements:
 <requirements>
 {% for req in requirements -%}
-- {{ loop.index }}. {{ req["requirement"] }}
+- {{ loop.index0 }}. {{ req["requirement"] }}
 {% endfor -%}
 </requirements>
 

prompts/criticize.txt

@@ -9,8 +9,8 @@ Here are the solutions:
 {% for solution in solutions %}
 ## Solution
 - Context: {{solution["Context"]}}
-- Problem description: {{solution["Problem Description"]}}
-- Solution description: {{solution["Solution Description"]}}
+- Problem description: {{solution["Problem_Description"]}}
+- Solution description: {{solution["Solution_Description"]}}
 ---
 {% endfor -%}
 </solutions>

prompts/search_solution.txt

@@ -0,0 +1,21 @@
+<role>You are an expert system designer</role>
+<task>
+Your task is to create a solution which is a combination of mechanisms that addresses as many of the provided requirements of a category as possible and that by searching the web, while carefully considering the given context.
+Please actually make searches and do not simulate them.
+</task>
+
+Here is the category item and the associated requirements:
+<requirements>
+Category Title: {{category["title"]}}
+Context: {{category["requirements"][0]["context"]}}
+Requirements:
+{% for req in category["requirements"] -%}
+- {{loop.index0}} {{req["requirement"]}}
+{% endfor -%}
+</requirements>
+
+<additional_instructions>
+- The solution must aim to maximize requirement satisfaction while respecting the context.
+- Provide a list of requirements addressed by the solution (provide only the requirement IDs)
+- Please also detail each mechanism used in final solution
+</additional_instructions>

prompts/structure_solution.txt

@@ -0,0 +1,12 @@
+<role>You are an expert system designer</role>
+<task>Your task is to take a solution you've created previously and structure it into a JSON object.</task>
+
+Here is the solution
+<solution>
+{{solution}}
+</solution>
+
+<response_format>
+Reply in JSON using the following format:
+{{response_schema}}
+</response_format> 

schemas.py

@@ -21,13 +21,23 @@ class ReqGroupingCategory(BaseModel):
 
 
 class SolutionModel(BaseModel):
-    Context: str
-    Requirements: List[str]
-    ProblemDescription: str
-    SolutionDescription: str
-    References: Optional[str] = ""
+    Context: str = Field(...,
+                         description="Full context provided for this category.")
+    Requirements: List[str] = Field(...,
+                                    description="List of each requirement as string.")
+    Problem_Description: str = Field(..., alias="Problem Description",
+                                     description="Description of the problem being solved.")
+    Solution_Description: str = Field(..., alias="Solution Description",
+                                      description="Detailed description of the solution.")
+    References: list[str] = Field(
+        ..., description="References to documents used for the solution.")
+    Category_Id: int = Field(
+        ..., description="ID of the requirements category the solution is based on")
 
-# Categorize requirements endpoint
+    class Config:
+        validate_by_name = True  # Enables alias handling on input/output
+
+# ============================================================= Categorize requirements endpoint
 
 
 class ReqGroupingRequest(BaseModel):
@@ -54,7 +64,7 @@ class _ReqGroupingOutput(BaseModel):
         ..., description="List of grouping categories")
 
 
-# Criticize solution endpoint
+# =========================================================== Criticize solution endpoint
 
 class CriticizeSolutionsRequest(BaseModel):
     solutions: list[SolutionModel]
@@ -82,3 +92,24 @@ class SolutionCriticism(BaseModel):
 
 class CritiqueResponse(BaseModel):
     critiques: List[SolutionCriticism]
+
+
+# =================================================================== search solution response endpoint
+
+class _SolutionSearchOutput(BaseModel):
+    solution: SolutionModel
+
+
+class _SearchedSolutionModel(BaseModel):
+    """"Internal model used for solutions searched using gemini"""
+    requirement_ids: List[int] = Field(...,
+                                       description="List of each requirement ID addressed by the solution")
+    problem_description: str = Field(...,
+                                     description="Description of the problem being solved.")
+    solution_description: str = Field(...,
+                                      description="Detailed description of the solution.")
+
+
+class SolutionSearchResponse(BaseModel):
+    """Response model for solution search"""
+    solutions: list[SolutionModel]