Add search_solution/v2 endpoint

app.py

@@ -4,7 +4,7 @@ import os
 import sys
 import uvicorn
 from fastapi import APIRouter, FastAPI
-from schemas import _RefinedSolutionModel, _SearchedSolutionModel, _SolutionCriticismOutput, CriticizeSolutionsRequest, CritiqueResponse, RequirementInfo, ReqGroupingCategory, ReqGroupingResponse, ReqGroupingRequest, _ReqGroupingCategory, _ReqGroupingOutput, SolutionCriticism, SolutionModel, SolutionSearchResponse
+from schemas import _RefinedSolutionModel, _SearchedSolutionModel, _SolutionCriticismOutput, CriticizeSolutionsRequest, CritiqueResponse, RequirementInfo, ReqGroupingCategory, ReqGroupingResponse, ReqGroupingRequest, _ReqGroupingCategory, _ReqGroupingOutput, SolutionCriticism, SolutionModel, SolutionSearchResponse, SolutionSearchV2Request
 from jinja2 import Environment, FileSystemLoader, StrictUndefined
 from litellm.router import Router
 from dotenv import load_dotenv
@@ -178,6 +178,72 @@ async def search_solutions(params: ReqGroupingResponse) -> SolutionSearchRespons
     return SolutionSearchResponse(solutions=final_solutions)
 
 
+@solution_router.post("/search_solutions_gemini/v2", response_model=SolutionSearchResponse)
+async def search_solutions(params: SolutionSearchV2Request) -> SolutionSearchResponse:
+    """Searches solutions solving the given grouping params and respecting the user constraints using Gemini and grounded on google search"""
+
+    logging.info(f"Searching solutions for categories: {params.categories}")
+
+    async def _search_inner(cat: ReqGroupingCategory) -> SolutionModel:
+        # ================== generate the solution with web grounding
+        req_prompt = await prompt_env.get_template("search_solution_v2.txt").render_async(**{
+            "category": cat.model_dump(),
+            "user_constraints": params.user_constraints,
+        })
+
+        # generate the completion in non-structured mode.
+        # the googleSearch tool enables grounding gemini with google search
+        # this also forces gemini to perform a tool call
+        req_completion = await llm_router.acompletion(model="chat", messages=[
+            {"role": "user", "content": req_prompt}
+        ], tools=[{"googleSearch": {}}], tool_choice="required")
+
+        # ==================== 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 ================================
+
+        sources_metadata = []
+        # extract the source metadata from the search items, if gemini actually called the tools to search .... and didn't hallucinated
+        try:
+            sources_metadata.extend([{"name": a["web"]["title"], "url": a["web"]["uri"]}
+                                     for a in req_completion["vertex_ai_grounding_metadata"][0]['groundingChunks']])
+        except KeyError as ke:
+            pass
+
+        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(*[retry_until(_search_inner, cat, lambda v: len(v.References) > 0, 2) 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)
+
+
+# =================================================================================================================
+
+
 @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."""

prompts/search_solution.txt

@@ -4,6 +4,12 @@ Your task is to create a solution that should cover the maximum possible of requ
 The solution should be composed of multiple mechanisms, that are defined as processes, methods, or sequences of steps using a technology that explain how something works or achieves its requirements.
 You may for that search mechanisms for the different requirements.
 **Please actually make searches and do not simulate them.**
+
+You must aim for the following goals:
+    - 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 and how they work in the final solution description.
+    - Please describe the solution description using markdown in a consistent format.
 </task>
 
 Here is the category item and the associated requirements:
@@ -14,11 +20,4 @@ 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 and how they work in the final solution description.
-- Please describe the solution description using markdown in a consistent format.
-</additional_instructions>
+</requirements>

prompts/search_solution_v2.txt

@@ -0,0 +1,31 @@
+<role>You are an expert system designer</role>
+<task>
+Your task is to create a solution that should cover the maximum possible of requirements at once.
+The solution should be composed of multiple mechanisms, that are defined as processes, methods, or sequences of steps using a technology that explain how something works or achieves its requirements.
+You may for that search mechanisms for the different requirements.
+
+**Please actually make searches and do not simulate them.**
+
+You must aim for the following goals:
+- 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 and how they work in the final solution description.
+- Please describe the solution description using markdown in a consistent format.
+</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>
+
+{% if user_constraints is not none %}
+Here are additional user constraints the solution must respect:
+<user_constraints>
+    {{user_constraints}}
+</user_constraints> 
+{% endif %}

schemas.py

@@ -114,6 +114,15 @@ class SolutionSearchResponse(BaseModel):
     """Response model for solution search"""
     solutions: list[SolutionModel]
 
+# ================================================================ search solution endpoint v2
+
+
+class SolutionSearchV2Request(BaseModel):
+    """Response of a requirement grouping call."""
+    categories: List[ReqGroupingCategory]
+    user_constraints: Optional[str] = Field(
+        default=None, description="Additional user constraints to respect when generating the solutions.")
+
 
 # ================================================================= refine solution endpoints
 
@@ -124,4 +133,4 @@ class _RefinedSolutionModel(BaseModel):
     problem_description: str = Field(...,
                                      description="New description of the problem being solved.")
     solution_description: str = Field(...,
-                                      description="New detailed description of the solution.")
+                                      description="New detailed description of the solution.")