Add retry with max bailout

app.py

@@ -8,6 +8,7 @@ from schemas import _RefinedSolutionModel, _SearchedSolutionModel, _SolutionCrit
 from jinja2 import Environment, FileSystemLoader, StrictUndefined
 from litellm.router import Router
 from dotenv import load_dotenv
+from util import retry_until
 
 logging.basicConfig(
     level=logging.INFO,
@@ -119,6 +120,8 @@ async def categorize_reqs(params: ReqGroupingRequest) -> ReqGroupingResponse:
 async def search_solutions(params: ReqGroupingResponse) -> SolutionSearchResponse:
     """Searches solutions solving the given grouping params 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.txt").render_async(**{
@@ -167,7 +170,7 @@ async def search_solutions(params: ReqGroupingResponse) -> SolutionSearchRespons
         )
         return final_sol
 
-    solutions = await asyncio.gather(*[_search_inner(cat) for cat in params.categories], return_exceptions=True)
+    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)]

util.py

@@ -0,0 +1,21 @@
+import asyncio
+from typing import Awaitable, Callable, TypeVar
+
+
+T = TypeVar("T")
+A = TypeVar("A")
+
+
+async def retry_until(
+    func: Callable[[A], Awaitable[T]],
+    arg: A,
+    predicate: Callable[[T], bool],
+    max_retries: int,
+) -> T:
+    """Retries the given async function until the passed in validation predicate returns true."""
+    last_value = await func(arg)
+    for _ in range(max_retries):
+        if predicate(last_value):
+            return last_value
+        last_value = await func(arg)
+    return last_value