Update app.py

app.py

@@ -1,7 +1,163 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel, Field
+from typing import List, Optional
+import os
+import json
+
+# ---- Requirements Models ----
+
+class RequirementInfo(BaseModel):
+    """Represents an extracted requirement info."""
+    context: str = Field(..., description="Context for the requirement.")
+    requirement: str = Field(..., description="The requirement itself.")
+    document: Optional[str] = Field('', description="The document the requirement is extracted from.")
+
+class ReqGroupingCategory(BaseModel):
+    """Represents the category of requirements grouped together."""
+    id: int = Field(..., description="ID of the grouping category")
+    title: str = Field(..., description="Title given to the grouping category")
+    requirements: List[RequirementInfo] = Field(
+        ..., description="List of grouped requirements")
+
+class ReqGroupingResponse(BaseModel):
+    categories: List[ReqGroupingCategory]
+
+
+# ---- Solution Models ----
+
+class SolutionModel(BaseModel):
+    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: Optional[str] = Field('', description="References to documents used for the solution.")
+
+    class Config:
+        allow_population_by_field_name = True  # Enables alias handling on input/output
+
+class SolutionsResponse(BaseModel):
+    solutions: List[SolutionModel]
+
+
+# ---- FastAPI app ----
 
 app = FastAPI()
 
+
+# ---- LLM Integration ----
+
+def ask_llm(user_message, model='compound-beta', system_prompt="You are a helpful assistant"):
+    from groq import Groq  # Import here so the app starts without the module if needed
+    client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+
+    response = client.chat.completions.create(
+        model=model,
+        messages=[
+            {
+                "role": "system",
+                "content": system_prompt
+            },
+            {
+                "role": "user",
+                "content": user_message
+            }
+        ],
+        stream=False,
+    )
+    ai_reply = response.choices[0].message.content
+    return ai_reply
+
+solution_prompt = """You are an expert system designer. Your task is to find a solution that addresses as many of the provided requirements as possible, while carefully considering the given context. Browse internet for reliable sources.
+
+Respond strictly in the following JSON format:
+
+{
+"Context": "<Insert the full context provided for this category>",
+"Requirements": [
+"<List each requirement clearly as a string item>"
+],
+"Problem Description": "<Describe the problem the solution is solving without introducing it>",
+"Solution Description": "<Explain the proposed solution, detailing how it meets each of the specified requirements and aligns with the given context. Prioritize completeness and practicality.>",
+"References": "<The references to the documents used to write the solution>"
+}
+
+text
+⚠️ Rules:
+
+Do not omit any part of the JSON structure.
+
+Replace newline characters with \"\\n\" (double backslash-n for JSON)
+
+Ensure all fields are present, even if empty.
+
+The solution must aim to maximize requirement satisfaction while respecting the context.
+
+Provide a clear and well-reasoned description of how your solution addresses each requirement.
+"""
+
+# ---- Main Endpoint ----
+
+@app.post("/find_solutions", response_model=SolutionsResponse)
+async def find_solutions(requirements: ReqGroupingResponse):
+    solutions = []
+
+    for category in requirements.categories:
+        category_title = category.title
+        category_requirements = category.requirements
+
+        # Compose the LLM prompt
+        problem_description = solution_prompt
+        problem_description += f"Category title: {category_title}\n"
+
+        context_list = []
+        requirement_list = []
+        for req_item in category_requirements:
+            context_list.append(f"- Context: {req_item.context}")
+            requirement_list.append(f"- Requirement: {req_item.requirement}")
+
+        problem_description += "Contexts:\n" + "\n".join(context_list) + "\n\n"
+        problem_description += "Requirements:\n" + "\n".join(requirement_list)
+
+        llm_response = ask_llm(problem_description)
+        print(f"Solution for '{category_title}' category:")
+        print(llm_response)
+
+        # Clean and parse the LLM response
+        try:
+            # Remove code blocks if present
+            cleaned = llm_response.strip()
+            if cleaned.startswith('```
+                cleaned = cleaned[7:]
+            if cleaned.startswith('```'):
+                cleaned = cleaned[3:]
+            if cleaned.endswith('```
+                cleaned = cleaned[:-3]
+            cleaned = cleaned.strip()
+
+            # Replace double backslashes with single if needed for parsing
+            cleaned = cleaned.replace('\\\\n', '\\n')
+            parsed = json.loads(cleaned)
+
+            # Use alias-aware population for SolutionModel
+            solution_obj = SolutionModel.parse_obj(parsed)
+            solutions.append(solution_obj)
+        except Exception as e:
+            # Append error info as a solution model (helps debug)
+            error_solution = SolutionModel(
+                Context="",
+                Requirements=[],
+                Problem_Description=f"Failed to parse LLM response: {str(e)}",
+                Solution_Description=f"Original LLM output: {llm_response}",
+                References=""
+            )
+            solutions.append(error_solution)
+
+    return SolutionsResponse(solutions=solutions)
+
+
 @app.get("/")
 def greet_json():
-    return {"Hello": "World!"}
+    return {"Status": "OK!"}