WIP

app.py

@@ -1,11 +1,13 @@
+import asyncio
 import logging
 import os
 import sys
 import uvicorn
 from fastapi import FastAPI
-from schemas import RequirementInfo, ReqGroupingCategory, ReqGroupingResponse, ReqGroupingRequest, _ReqGroupingCategory, _ReqGroupingOutput
+from schemas import CriticizeSolutionsRequest, RequirementInfo, ReqGroupingCategory, ReqGroupingResponse, ReqGroupingRequest, _ReqGroupingCategory, _ReqGroupingOutput
 from jinja2 import Environment, FileSystemLoader
 from litellm import acompletion
+from litellm.router import Router
 from dotenv import load_dotenv
 
 logging.basicConfig(
@@ -17,21 +19,30 @@ logging.basicConfig(
 # Load .env files
 load_dotenv()
 
-if "LLM_API_MODEL" not in os.environ or "LLM_API_KEY" not in os.environ:
+if "LLM_MODEL" not in os.environ or "LLM_API_KEY" not in os.environ:
     logging.error(
-        "No LLM token (`LLM_API_TOKEN`) and/or LLM model (`LLM_API_KEY`) were provided in the env vars. Exiting")
+        "No LLM token (`LLM_API_KEY`) and/or LLM model (`LLM_MODEL`) were provided in the env vars. Exiting")
     sys.exit(-1)
 
-LLM_API_MODEL = os.environ.get("LLM_API_MODEL")
-LLM_API_KEY = os.environ.get("LLM_API_KEY")
+# LiteLLM router
+llm_router = Router(model_list=[
+    {
+        "model_name": "chat",
+        "litellm_params": {
+            "model": os.environ.get("LLM_MODEL"),
+            "api_key": os.environ.get("LLM_API_KEY"),
+            "max_retries": 5
+        }
+    }
+])
 
 # Jinja2 environment to load prompt templates
 prompt_env = Environment(loader=FileSystemLoader('prompts'), enable_async=True)
 
-fastapi = FastAPI()
+api = FastAPI()
 
 
-@fastapi.post("/categorize_requirements")
+@api.post("/categorize_requirements")
 async def categorize_reqs(params: ReqGroupingRequest) -> ReqGroupingResponse:
     """Categorize the given service requirements into categories"""
 
@@ -46,12 +57,14 @@ 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})
 
     # ensure all requirements items are processed
     for attempt in range(MAX_ATTEMPTS):
-        req_completion = await acompletion(model=LLM_API_MODEL, api_key=LLM_API_KEY, messages=messages, response_format=_ReqGroupingOutput)
+        req_completion = await llm_router.acompletion(model="chat", messages=messages, response_format=_ReqGroupingOutput)
         output = _ReqGroupingOutput.model_validate_json(
             req_completion.choices[0].message.content)
 
@@ -85,4 +98,12 @@ async def categorize_reqs(params: ReqGroupingRequest) -> ReqGroupingResponse:
     return ReqGroupingResponse(categories=final_categories)
 
 
-uvicorn.run(fastapi, host="0.0.0.0", port=8000)
+@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
+
+
+uvicorn.run(api, host="0.0.0.0", port=8000)

prompts/classify.txt

@@ -3,20 +3,19 @@
 For each category indicate which requirements belong in that category using their ID. An item may appear in one category at a time. 
 Please make each category title indicative of whats in it.
 </task>
-{% if max_n_categories is none %}
+{% if max_n_categories is none -%}
 <number_of_categories>You may have at most as much categories as you think is needed</number_of_categories>
-{% else %}
+{% else -%}
 <number_of_categories>You may have at most {{max_n_categories}} categories</number_of_categories>
-{%endif%}
+{% endif-%}
 
 Here are the requirements:
 <requirements>
-{% for req in requirements %}
+{% for req in requirements -%}
 - {{ loop.index }}. {{ req["requirement"] }}
-{% endfor %}
+{% endfor -%}
 </requirements>
 
-
 <response_format>
 Reply in JSON using the following format:
 {{response_schema}}

prompts/criticize.txt

@@ -0,0 +1,16 @@
+<role>You are an useful engineering assistant for innovation</role>
+<task>
+You are tasked with criticizing multiple solutions solving a set of requirements on different points, 
+namely the technical challenges of the solutions, the weaknesses and limitations of the proposed solutions.
+</task>
+
+Here are the solutions:
+<solutions>
+{% for solution in solutions %}
+## Solution
+- Context: {{solution["Context"]}}
+- Problem description: {{solution["Problem Description"]}}
+- Solution description: {{solution["Solution Description"]}}
+---
+{% endfor -%}
+</solutions>

schemas.py

@@ -3,6 +3,7 @@ from typing import List, Optional
 
 # Shared model schemas
 
+
 class RequirementInfo(BaseModel):
     """Represents an extracted requirement info"""
     context: str = Field(..., description="Context for the requirement.")
@@ -18,7 +19,16 @@ class ReqGroupingCategory(BaseModel):
     requirements: List[RequirementInfo] = Field(
         ..., description="List of grouped requirements")
 
-# Endpoint model schemas
+
+class SolutionSearchResult(BaseModel):
+    Context: str
+    Requirements: List[str]
+    ProblemDescription: str
+    SolutionDescription: str
+    References: Optional[str] = ""
+
+# Categorize requirements endpoint
+
 
 class ReqGroupingRequest(BaseModel):
     """Request schema of a requirement grouping call."""
@@ -42,3 +52,9 @@ class _ReqGroupingCategory(BaseModel):
 class _ReqGroupingOutput(BaseModel):
     categories: list[_ReqGroupingCategory] = Field(
         ..., description="List of grouping categories")
+
+
+# Criticize solution endpoint
+
+class CriticizeSolutionsRequest(BaseModel):
+    solutions: list[SolutionSearchResult]