Update generator

app.py

@@ -164,7 +164,6 @@ async def generate_saved_endpoint(
             )
             questions_list = []
             if isinstance(mcqs, dict):
-                # if dict values are lists or single objects, collect them
                 for v in mcqs.values():
                     if isinstance(v, list):
                         questions_list.extend(v)
@@ -173,13 +172,10 @@ async def generate_saved_endpoint(
             elif isinstance(mcqs, list):
                 questions_list = mcqs
             else:
-                # unexpected shape — skip or handle as needed
                 continue
 
-            # assign sequential ids (as strings). use `counter` for numeric keys if preferred.
             for qobj in questions_list:
                 if isinstance(qobj, dict):
-                    # optional: keep which difficulty the question came from
                     qobj["_difficulty"] = difficulty
                 all_mcqs[str(counter)] = qobj
                 counter += 1
@@ -208,7 +204,9 @@ async def generate_saved_endpoint(
 async def generate_endpoint(
     background_tasks: BackgroundTasks,
     file: UploadFile = File(...),
-    n_questions: int = Form(10),
+    n_easy_questions: int = Form(3),
+    n_medium_questions: int = Form(5), 
+    n_hard_questions: int = Form(2),
     qdrant_filename: str = Form("default_filename"),
     collection_name: str = Form("programming"),
     mode: str = Form("rag"),
@@ -244,19 +242,48 @@ async def generate_endpoint(
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Could not save file to Qdrant Cloud: {e}")
 
-    # generate
-    try:
-        mcqs = rag.generate_from_pdf(
-            tmp_path,
-            n_questions=n_questions,
-            mode=mode,
-            questions_per_page=questions_per_page,
-            top_k=top_k,
-            temperature=temperature,
-            enable_fiddler=enable_fiddler
-        )
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Generation failed: {e}")
+    difficulty_counts = {
+        "easy": n_easy_questions,
+        "medium": n_medium_questions,
+        "hard": n_hard_questions
+    }
+    
+    all_mcqs = {}
+    counter = 1
+
+    for difficulty, n_questions in difficulty_counts.items():
+        try:
+            mcqs = rag.generate_from_pdf(
+                tmp_path,
+                collection=collection_name,
+                n_questions=n_questions,
+                mode=mode,
+                questions_per_page=questions_per_page,
+                top_k=top_k,
+                temperature=temperature,
+                enable_fiddler=enable_fiddler,
+                target_difficulty=difficulty,
+            )
+            questions_list = []
+            if isinstance(mcqs, dict):
+                for v in mcqs.values():
+                    if isinstance(v, list):
+                        questions_list.extend(v)
+                    else:
+                        questions_list.append(v)
+            elif isinstance(mcqs, list):
+                questions_list = mcqs
+            else:
+                continue
+
+            for qobj in questions_list:
+                if isinstance(qobj, dict):
+                    qobj["_difficulty"] = difficulty
+                all_mcqs[str(counter)] = qobj
+                counter += 1
+
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Generation from saved file failed: {e}")
 
     validation_report = None