import shlex import re import os import subprocess import sys import argparse from pathlib import Path from pydantic import BaseModel, Field from typing import List, Union from google import genai from google.genai import types from concurrent.futures import ThreadPoolExecutor from utils import compile_to_pdf def get_lcp(s1: str, s2: str) -> str: i = 0 while i < len(s1) and i < len(s2) and s1[i] == s2[i]: i += 1 return s1[:i] # MODEL_ID = "gemini-3-flash-preview" MODEL_ID = "gemini-3.1-flash-lite" api_key = os.environ.get("GEMINI_API_KEY") # --- Modèles pour la Requête 1 --- class QuestionOnlyItem(BaseModel): label: str = Field(description="The unique label of the question (e.g., '1.a', 'Exercice 1')") question_content: str = Field(description="The source text of the question, strictly extracted from the enonce file, EXCLUDING the label itself.") class ExamQuestions(BaseModel): questions: List[QuestionOnlyItem] # --- Modèles pour la Requête 2 --- class SolutionOnlyItem(BaseModel): label: str = Field(description="The exact unique label of the question provided in the input.") solution_content: str = Field(description="The source text of the solution, strictly extracted from the correction file.") class ExamSolutions(BaseModel): solutions: List[SolutionOnlyItem] # --- Modèles pour la Requête 3 --- class ExtractedContext(BaseModel): target_question_label: str = Field(description="The exact label of the FIRST question that comes immediately AFTER this information in the exam.") context_content: str = Field(description="The source text of the definitions, notations, or hypotheses, extracted from the enonce.") class ExamContext(BaseModel): contexts: List[ExtractedContext] # --- Modèle fusionné (pour le reste du script) --- class QuestionItem(BaseModel): label: str question_content: str solution_content: str class ContextItem(BaseModel): content: str # Juste une string encapsulée pour le différencier facilement class ExamExtraction(BaseModel): items: List[Union[QuestionItem, ContextItem]] # Liste mixte class GroupedExamExtraction(BaseModel): groups: List[List[Union[QuestionItem, ContextItem]]] PROMPT_1 = """I am providing: 1. A PDF of an exam (`enonce.pdf`) 2. The source code of the exam questions (`enonce` file) Your task: 1. Identify all distinct question labels using the PDF document. These labels should be unique : use `Ex 1 : 1)a)` or `I)1)b)`. 2. For each label, extract its exact corresponding question text from the `enonce` source file. Do not include the label itself in this extracted text (nor LaTeX like `item` nor org-mode list labelling like `2.`). Return the result as a JSON list in the exact reading order of the document. """ PROMPT_2 = """I am providing: 1. A JSON list of question labels and their texts extracted from an exam. 2. The source code of the exam solutions (`correction` file). Your task: For each question label provided in the JSON, extract its exact corresponding solution textual content from the `correction` source file. Return the result as a JSON list in the exact same order. """ PROMPT_3 = """I am providing: 1. A JSON list of question labels and their texts extracted from an exam. 2. The source code of the exam questions (`enonce` file). Your task: Extract important information necessary to understand the questions (e.g., definitions of objects, global notations, hypotheses, context) that are NOT part of the question texts themselves. For each extracted piece of information, identify the label of the FIRST question that comes immediately AFTER this information in the exam. Return the result as a JSON list. """ def find_file(folder: Path, base_name: str) -> Path: for ext in [".org", ".tex"]: path = folder / f"{base_name}{ext}" if path.is_file(): return path return None def process_exam(folder_path: str): folder = Path(folder_path) # 1. Resolve files pdf_path = folder / "enonce.pdf" enonce_path = find_file(folder, "enonce") correction_path = find_file(folder, "correction") missing = [] if not pdf_path.is_file(): missing.append("enonce.pdf") if not enonce_path: missing.append("enonce.org or enonce.tex") if not correction_path: missing.append("correction.org or correction.tex") if missing: print(f"Error: Missing files in {folder}: {', '.join(missing)}") sys.exit(1) print("Reading files...") pdf_bytes = pdf_path.read_bytes() enonce_text = enonce_path.read_text(encoding="utf-8") correction_text = correction_path.read_text(encoding="utf-8") client = genai.Client(api_key=api_key) # ========================================== # REQUÊTE 1 : Extraction des Énoncés # ========================================== contents_1 = [ types.Content( role="user", parts=[ types.Part.from_text(text=PROMPT_1), types.Part.from_bytes(data=pdf_bytes, mime_type="application/pdf"), types.Part.from_text(text=f"--- ENONCE SOURCE ({enonce_path.name}) ---\n{enonce_text}"), ], ) ] config_1 = types.GenerateContentConfig( temperature=0.1, response_mime_type="application/json", response_json_schema=ExamQuestions.model_json_schema(), ) cache_q_file = folder / "gemini_questions.json" if cache_q_file.is_file(): print("Loading cached questions from gemini_questions.json...") response_q_text = cache_q_file.read_text(encoding="utf-8") else: print("Sending request 1 (Questions) to Gemini...") response_q = client.models.generate_content( model=MODEL_ID, contents=contents_1, config=config_1 ) response_q_text = response_q.text print("Saving questions to cache...") cache_q_file.write_text(response_q_text, encoding="utf-8") questions_data = ExamQuestions.model_validate_json(response_q_text) # ========================================== # REQUÊTE 2 : Extraction des Corrections # ========================================== extracted_questions_json = questions_data.model_dump_json(indent=2) contents_2 = [ types.Content( role="user", parts=[ types.Part.from_text(text=PROMPT_2), types.Part.from_text(text=f"--- EXTRACTED QUESTIONS ---\n{extracted_questions_json}"), types.Part.from_text(text=f"--- CORRECTION SOURCE ({correction_path.name}) ---\n{correction_text}"), ], ) ] config_2 = types.GenerateContentConfig( temperature=0.1, response_mime_type="application/json", response_json_schema=ExamSolutions.model_json_schema(), ) cache_s_file = folder / "gemini_solutions.json" if cache_s_file.is_file(): print("Loading cached solutions from gemini_solutions.json...") response_s_text = cache_s_file.read_text(encoding="utf-8") else: print("Sending request 2 (Solutions) to Gemini...") response_s = client.models.generate_content( model=MODEL_ID, contents=contents_2, config=config_2 ) response_s_text = response_s.text print("Saving solutions to cache...") cache_s_file.write_text(response_s_text, encoding="utf-8") solutions_data = ExamSolutions.model_validate_json(response_s_text) # ========================================== # REQUÊTE 3 : Extraction du Contexte (Notations, etc.) # ========================================== contents_3 = [ types.Content( role="user", parts=[ types.Part.from_text(text=PROMPT_3), types.Part.from_text(text=f"--- EXTRACTED QUESTIONS ---\n{extracted_questions_json}"), types.Part.from_text(text=f"--- ENONCE SOURCE ({enonce_path.name}) ---\n{enonce_text}"), ], ) ] config_3 = types.GenerateContentConfig( temperature=0.1, response_mime_type="application/json", response_json_schema=ExamContext.model_json_schema(), ) cache_c_file = folder / "gemini_context.json" if cache_c_file.is_file(): print("Loading cached context from gemini_context.json...") response_c_text = cache_c_file.read_text(encoding="utf-8") else: print("Sending request 3 (Context) to Gemini...") response_c = client.models.generate_content( model=MODEL_ID, contents=contents_3, config=config_3 ) response_c_text = response_c.text print("Saving context to cache...") cache_c_file.write_text(response_c_text, encoding="utf-8") context_data = ExamContext.model_validate_json(response_c_text) # ========================================== # FUSION des trois résultats # ========================================== sol_map = {s.label: s.solution_content for s in solutions_data.solutions} # Grouper les contextes par label cible ctx_map = {} for c in context_data.contexts: if c.target_question_label in ctx_map: ctx_map[c.target_question_label] += "\n\n" + c.context_content else: ctx_map[c.target_question_label] = c.context_content merged_items = [] for q in questions_data.questions: # 1. S'il y a un contexte pour cette question, on l'insère d'abord dans la liste if q.label in ctx_map: merged_items.append(ContextItem(content=ctx_map[q.label])) # 2. Puis on ajoute la question sol_content = sol_map.get(q.label, "") merged_items.append(QuestionItem( label=q.label, question_content=q.question_content, solution_content=sol_content )) extracted_data = ExamExtraction(items=merged_items) # ========================================== # INITIAL GROUPING COMPUTATION # ========================================== items_file = folder / "exam_items.txt" full_items_file = folder / "exam_items_full.txt" trunc_map = {} # --- INITIAL GROUPING COMPUTATION --- questions_only = [item for item in extracted_data.items if isinstance(item, QuestionItem)] q_group_indices = [] if questions_only: current_g = [0] for i in range(1, len(questions_only)): p = get_lcp(questions_only[current_g[0]].label, questions_only[i].label) proposed = current_g + [i] valid = True for k in range(len(proposed) - 1): if get_lcp(questions_only[proposed[k]].label, questions_only[proposed[k+1]].label) != p: valid = False break if valid: current_g.append(i) else: q_group_indices.append(current_g) current_g = [i] q_group_indices.append(current_g) group_starter_labels = {questions_only[g[0]].label for g in q_group_indices[1:]} if q_group_indices else set() initial_groups = [] current_group = [] for i, item in enumerate(extracted_data.items): is_new_group = False if isinstance(item, QuestionItem): if item.label in group_starter_labels: if not (len(current_group) > 0 and isinstance(current_group[-1], ContextItem)): is_new_group = True elif isinstance(item, ContextItem): if i + 1 < len(extracted_data.items): next_item = extracted_data.items[i+1] if isinstance(next_item, QuestionItem) and next_item.label in group_starter_labels: is_new_group = True if is_new_group and current_group: initial_groups.append(current_group) current_group = [] current_group.append(item) if current_group: initial_groups.append(current_group) # ---- Transform labels, and check uniqueness seen_labels = set() for group in initial_groups: for item in group: if isinstance(item, QuestionItem): # 1. Transform label item.label = item.label.replace("Exercice", "Ex") item.label = item.label.replace(".", ")") # 2. Ensure uniqueness (prefix with XX) while item.label in seen_labels: item.label = f"XX{item.label}" seen_labels.add(item.label) # --- WRITE TEXT FILES --- print(f"Writing items files to {items_file.name} and {full_items_file.name}...") with open(items_file, "w", encoding="utf-8") as f, \ open(full_items_file, "w", encoding="utf-8") as f_full: header = "# Edit labels. Modify groups (---). Ensure label uniqueness (XX). Duplicate CONTEXT.\n\n" f.write(header) f_full.write(header) for g_idx, group in enumerate(initial_groups): if g_idx > 0: f.write("\n---\n\n") f_full.write("\n---\n\n") for item in group: if isinstance(item, QuestionItem): safe_content = item.question_content.replace('\n', ' \\n ') f_full.write(f"{item.label} ### {safe_content}\n") if len(safe_content) > 65: trunc_content = safe_content[:64] + "…" trunc_map[trunc_content] = safe_content else: trunc_content = safe_content f.write(f"{item.label} ### {trunc_content}\n") elif isinstance(item, ContextItem): safe_content = item.content.replace('\n', ' \\n ') f.write(f"CONTEXT ### {safe_content}\n") f_full.write(f"CONTEXT ### {safe_content}\n") # --- OPEN EDITOR AND PARSE --- while True: print("Opening items file for editing...") editor = os.environ.get("EDITOR") try: if editor: subprocess.run(shlex.split(editor) + [str(items_file)]) else: if sys.platform.startswith("linux"): subprocess.run(["xdg-open", str(items_file)]) elif sys.platform == "darwin": subprocess.run(["open", str(items_file)]) else: os.startfile(str(items_file)) input("Press ENTER here once you have saved and closed the text file...") except Exception as e: print(f"Error running editor: {e}") print(f"Parsing edited items from {items_file.name}...") with open(items_file, "r", encoding="utf-8") as f: edited_lines = [line.strip() for line in f if line.strip() and not line.startswith("#")] # 1. Validation for XX labels has_xx = False for line in edited_lines: if " ### " in line: lbl = line.split(" ### ", 1)[0].strip() if lbl.startswith("XX"): has_xx = True break if has_xx: print("\n!!! ERROR: Some labels still start with 'XX'. Please remove the 'XX' prefixes to ensure unique, valid labels.") input("Press ENTER to return to the editor...") continue # 2. Actual Parsing grouped_items = [] current_group = [] labels_list = [] orig_idx = 0 for line in edited_lines: if line == "---": if current_group: grouped_items.append(current_group) current_group = [] continue if " ### " not in line: continue new_label, edited_content_raw = line.split(" ### ", 1) new_label = new_label.strip() if new_label != "CONTEXT": labels_list.append(new_label) if "…" in edited_content_raw and edited_content_raw in trunc_map: edited_content_raw = trunc_map[edited_content_raw] edited_content = edited_content_raw.replace(' \\n ', '\n') if orig_idx < len(extracted_data.items): orig_item = extracted_data.items[orig_idx] if isinstance(orig_item, QuestionItem): current_group.append(QuestionItem( label=new_label, question_content=edited_content, solution_content=orig_item.solution_content )) elif isinstance(orig_item, ContextItem): current_group.append(ContextItem(content=edited_content)) orig_idx += 1 if current_group: grouped_items.append(current_group) # If we reached here without 'continue', the data is valid break # Save labels and proceed with open(folder / "labels", 'w', encoding='utf-8') as f_labels: for label in labels_list: f_labels.write(f"{label}\n") grouped_extraction = GroupedExamExtraction(groups=grouped_items) # 2. Setup output directories text_dir = folder / "Text" sol_dir = folder / "Sol" text2_dir = folder / "Text2" sol2_dir = folder / "Sol2" dirs = [text_dir, sol_dir, text2_dir, sol2_dir] import shutil # Ask only if any directory already exists if any(d.exists() for d in dirs): answer = input( "Output directories already exist. Delete their contents? [y/N] " ).strip().lower() if answer not in ("y", "yes"): print("Aborted.") sys.exit(1) # Empty each directory for d in dirs: if d.exists(): shutil.rmtree(d) d.mkdir(parents=True) else: # Create them if they don't exist for d in dirs: d.mkdir(parents=True) text_dir.mkdir(exist_ok=True) sol_dir.mkdir(exist_ok=True) print("Writing grouped question and solution files...") for group in grouped_extraction.groups: q_items = [item for item in group if isinstance(item, QuestionItem)] labels = [q.label for q in q_items] if not labels: continue # Skip if a group has no questions (only contexts) # 1. Compute the common prefix for the group prefix = labels[0] for lbl in labels[1:]: prefix = get_lcp(prefix, lbl) # 2. Format the Text filename: prefix [label1, label2] labels_str = ",".join([label[len(prefix):] for label in labels]) group_filename = f"{prefix}[{labels_str}]" safe_group_filename = group_filename.replace("/", "_") text_content_lines = [] # 3. Process each item in the group for item in group: if isinstance(item, QuestionItem): # 1. Prepare tabulated content: # Start with a tab, then replace every newline+whitespace with newline+tab raw_content = item.question_content.strip() tabulated = "\t" + re.sub(r'\n\s*', '\n\t', raw_content) # 2. Build Text entry text_content_lines.append(f"{item.label} :") text_content_lines.append(tabulated) # Write individual Sol file (remains unchanged) safe_label = item.label.replace("/", "_") with open(sol_dir / safe_label, "w", encoding="utf-8") as f_sol: f_sol.write(f"{item.label}\n{item.solution_content}") with open(text2_dir / f"{safe_label}.tex", "w", encoding="utf-8") as f_t2: f_t2.write(f"\\textbf{{{item.label}}} {item.question_content}") with open(sol2_dir / f"{safe_label}.tex", "w", encoding="utf-8") as f_s2: f_s2.write(f"\\textbf{{{item.label}}} {item.solution_content}") elif isinstance(item, ContextItem): raw_ctx = item.content.strip() tabulated_ctx = "\t" + re.sub(r'\n\s*', '\n\t', raw_ctx) text_content_lines.append(f"CONTEXT :") text_content_lines.append(tabulated_ctx) # 4. Write the grouped Text file with open(text_dir / safe_group_filename, "w", encoding="utf-8") as f_text: f_text.write("\n".join(text_content_lines)) print(f"Success! Processed {len(grouped_extraction.groups)} groups.") # ========================================== # PDF COMPILATION (4 Threads) # ========================================== all_tex_files = list(text2_dir.glob("*.tex")) + list(sol2_dir.glob("*.tex")) def compile_worker(tex_file: Path): """Helper to read content and call the utility function.""" try: content = tex_file.read_text(encoding="utf-8") pdf_path = tex_file.with_suffix(".pdf") compile_to_pdf(content, pdf_path) except Exception as e: print(f"Error compiling {tex_file.name}: {e}") print(f"Compiling {len(all_tex_files)} files to PDF using 4 threads...") with ThreadPoolExecutor(max_workers=4) as executor: executor.map(compile_worker, all_tex_files) if __name__ == "__main__": if not api_key: print("Error: GEMINI_API_KEY environment variable is not set.") sys.exit(1) parser = argparse.ArgumentParser(description="Extract exam and solution code via Gemini.") parser.add_argument("folder", help="Directory containing the exam files") args = parser.parse_args() process_exam(args.folder)