Pull last commit app.py

app.py

@@ -1,3 +1,16 @@
+from bs4 import BeautifulSoup
+import warnings
+import io
+import zipfile
+from lxml import etree
+import os
+from dotenv import load_dotenv
+import requests
+import subprocess
+import string
+from nltk.tokenize import word_tokenize
+from nltk.corpus import stopwords
+from nltk.stem import WordNetLemmatizer
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import json
 import traceback
@@ -10,40 +23,54 @@ from litellm.router import Router
 from aiolimiter import AsyncLimiter
 import pandas as pd
 import asyncio
+import logging
 import re
 import nltk
 
+load_dotenv()
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='[%(asctime)s][%(levelname)s][%(filename)s:%(lineno)d]: %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+
 nltk.download('stopwords')
 nltk.download('punkt_tab')
 nltk.download('wordnet')
 
-from nltk.stem import WordNetLemmatizer
-from nltk.corpus import stopwords
-from nltk.tokenize import word_tokenize
-
-import string
-import subprocess
-import requests
-from dotenv import load_dotenv
-
-load_dotenv()
-
-import os
-from lxml import etree
-import zipfile
-import io
-import warnings
-
 warnings.filterwarnings("ignore")
 
-from bs4 import BeautifulSoup
-
 app = FastAPI(title="Requirements Extractor")
 app.mount("/static", StaticFiles(directory="static"), name="static")
-app.add_middleware(CORSMiddleware, allow_credentials=True, allow_headers=["*"], allow_methods=["*"], allow_origins=["*"])
-llm_router = Router(model_list=[{"model_name": "gemini-v1", "litellm_params": {"model": "gemini/gemini-2.0-flash", "api_key": os.environ.get("GEMINI"), "max_retries": 10, "rpm": 15}},
-                                {"model_name": "gemini-v2", "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": os.environ.get("GEMINI"), "max_retries": 10, "rpm": 10}}]
-                                , fallbacks=[{"gemini-v2": ["gemini-v1"]}], num_retries=10)
+app.add_middleware(CORSMiddleware, allow_credentials=True, allow_headers=[
+                   "*"], allow_methods=["*"], allow_origins=["*"])
+
+llm_router = Router(model_list=[
+    {
+        "model_name": "gemini-v1",
+        "litellm_params":
+        {
+            "model": "gemini/gemini-2.0-flash",
+            "api_key": os.environ.get("GEMINI"),
+            "max_retries": 10,
+            "rpm": 15,
+            "allowed_fails": 1,
+            "cooldown": 30,
+        }
+    },
+    {
+        "model_name": "gemini-v2",
+        "litellm_params":
+        {
+            "model": "gemini/gemini-2.5-flash",
+            "api_key": os.environ.get("GEMINI"),
+            "max_retries": 10,
+            "rpm": 10,
+            "allowed_fails": 1,
+            "cooldown": 30,
+        }
+    }], fallbacks=[{"gemini-v2": ["gemini-v1"]}], num_retries=10, retry_after=30)
 
 limiter_mapping = {
     model["model_name"]: AsyncLimiter(model["litellm_params"]["rpm"], 60)
@@ -56,15 +83,18 @@ NSMAP = {
     'v': 'urn:schemas-microsoft-com:vml'
 }
 
+
 def lemma(text: str):
     stop_words = set(stopwords.words('english'))
     txt = text.translate(str.maketrans('', '', string.punctuation)).strip()
-    tokens = [token for token in word_tokenize(txt.lower()) if token not in stop_words]
+    tokens = [token for token in word_tokenize(
+        txt.lower()) if token not in stop_words]
     return [lemmatizer.lemmatize(token) for token in tokens]
 
+
 def get_docx_archive(url: str) -> zipfile.ZipFile:
     """Récupère le docx depuis l'URL et le retourne comme objet ZipFile"""
-    if not url.endswith("zip"): 
+    if not url.endswith("zip"):
         raise ValueError("URL doit pointer vers un fichier ZIP")
     doc_id = os.path.splitext(os.path.basename(url))[0]
     resp = requests.get(url, verify=False, headers={
@@ -84,7 +114,7 @@ def get_docx_archive(url: str) -> zipfile.ZipFile:
 
                 with open(input_path, "wb") as f:
                     f.write(docx_bytes)
-                
+
                 subprocess.run([
                     "libreoffice",
                     "--headless",
@@ -98,17 +128,19 @@ def get_docx_archive(url: str) -> zipfile.ZipFile:
 
                 os.remove(input_path)
                 os.remove(output_path)
-                
+
                 return zipfile.ZipFile(io.BytesIO(docx_bytes))
 
     raise ValueError("Aucun fichier docx/doc trouvé dans l'archive")
 
+
 def parse_document_xml(docx_zip: zipfile.ZipFile) -> etree._ElementTree:
     """Parse le document.xml principal"""
     xml_bytes = docx_zip.read('word/document.xml')
     parser = etree.XMLParser(remove_blank_text=True)
     return etree.fromstring(xml_bytes, parser=parser)
 
+
 def clean_document_xml(root: etree._Element) -> None:
     """Nettoie le XML en modifiant l'arbre directement"""
     # Suppression des balises <w:del> et leur contenu
@@ -116,7 +148,7 @@ def clean_document_xml(root: etree._Element) -> None:
         parent = del_elem.getparent()
         if parent is not None:
             parent.remove(del_elem)
-    
+
     # Désencapsulation des balises <w:ins>
     for ins_elem in root.xpath('//w:ins', namespaces=NSMAP):
         parent = ins_elem.getparent()
@@ -125,7 +157,7 @@ def clean_document_xml(root: etree._Element) -> None:
             parent.insert(index, child)
             index += 1
         parent.remove(ins_elem)
-    
+
     # Nettoyage des commentaires
     for tag in ['w:commentRangeStart', 'w:commentRangeEnd', 'w:commentReference']:
         for elem in root.xpath(f'//{tag}', namespaces=NSMAP):
@@ -133,16 +165,17 @@ def clean_document_xml(root: etree._Element) -> None:
             if parent is not None:
                 parent.remove(elem)
 
+
 def create_modified_docx(original_zip: zipfile.ZipFile, modified_root: etree._Element) -> bytes:
     """Crée un nouveau docx avec le XML modifié"""
     output = io.BytesIO()
-    
+
     with zipfile.ZipFile(output, 'w', compression=zipfile.ZIP_DEFLATED) as new_zip:
         # Copier tous les fichiers non modifiés
         for file in original_zip.infolist():
             if file.filename != 'word/document.xml':
                 new_zip.writestr(file, original_zip.read(file.filename))
-        
+
         # Ajouter le document.xml modifié
         xml_str = etree.tostring(
             modified_root,
@@ -151,10 +184,11 @@ def create_modified_docx(original_zip: zipfile.ZipFile, modified_root: etree._El
             pretty_print=True
         )
         new_zip.writestr('word/document.xml', xml_str)
-    
+
     output.seek(0)
     return output.getvalue()
 
+
 def docx_to_txt(doc_id: str, url: str):
     docx_zip = get_docx_archive(url)
     root = parse_document_xml(docx_zip)
@@ -165,7 +199,7 @@ def docx_to_txt(doc_id: str, url: str):
     output_path = f"/tmp/{doc_id}_cleaned.txt"
     with open(input_path, "wb") as f:
         f.write(modified_bytes)
-    
+
     subprocess.run([
         "libreoffice",
         "--headless",
@@ -181,18 +215,20 @@ def docx_to_txt(doc_id: str, url: str):
     os.remove(output_path)
     return txt_data
 
+
 @app.get("/")
 def render_page():
     return FileResponse("index.html")
 
+
 @app.post("/get_meetings", response_model=MeetingsResponse)
 def get_meetings(req: MeetingsRequest):
     working_group = req.working_group
     tsg = re.sub(r"\d+", "", working_group)
     wg_number = re.search(r"\d", working_group).group(0)
-    print(tsg, wg_number)
+    logging.debug(tsg, wg_number)
     url = "https://www.3gpp.org/ftp/tsg_" + tsg
-    print(url)
+    logging.debug(url)
     resp = requests.get(url, verify=False)
     soup = BeautifulSoup(resp.text, "html.parser")
     meeting_folders = []
@@ -205,22 +241,27 @@ def get_meetings(req: MeetingsRequest):
             break
 
     url += "/" + selected_folder
-    print(url)
+    logging.debug(url)
 
     if selected_folder:
         resp = requests.get(url, verify=False)
         soup = BeautifulSoup(resp.text, "html.parser")
-        meeting_folders = [item.get_text() for item in soup.select("tr td a") if item.get_text().startswith("TSG") or (item.get_text().startswith("CT") and "-" in item.get_text())]
-        all_meetings = [working_group + "#" + meeting.split("_", 1)[1].replace("_", " ").replace("-", " ") if meeting.startswith('TSG') else meeting.replace("-","#") for meeting in meeting_folders]
-    
+        meeting_folders = [item.get_text() for item in soup.select("tr td a") if item.get_text(
+        ).startswith("TSG") or (item.get_text().startswith("CT") and "-" in item.get_text())]
+        all_meetings = [working_group + "#" + meeting.split("_", 1)[1].replace("_", " ").replace(
+            "-", " ") if meeting.startswith('TSG') else meeting.replace("-", "#") for meeting in meeting_folders]
+
     return MeetingsResponse(meetings=dict(zip(all_meetings, meeting_folders)))
 
+
 @app.post("/get_dataframe", response_model=DataResponse)
 def get_change_request_dataframe(req: DataRequest):
     working_group = req.working_group
     tsg = re.sub(r"\d+", "", working_group)
     wg_number = re.search(r"\d", working_group).group(0)
     url = "https://www.3gpp.org/ftp/tsg_" + tsg
+    logging.info("Fetching TDocs dataframe")
+
     resp = requests.get(url, verify=False)
     soup = BeautifulSoup(resp.text, "html.parser")
     wg_folders = [item.get_text() for item in soup.select("tr td a")]
@@ -233,18 +274,21 @@ def get_change_request_dataframe(req: DataRequest):
     url += "/" + selected_folder + "/" + req.meeting + "/docs"
     resp = requests.get(url, verify=False)
     soup = BeautifulSoup(resp.text, "html.parser")
-    files = [item.get_text() for item in soup.select("tr td a") if item.get_text().endswith(".xlsx")]
+    files = [item.get_text() for item in soup.select("tr td a")
+             if item.get_text().endswith(".xlsx")]
 
     def gen_url(tdoc: str):
         return f"{url}/{tdoc}.zip"
 
     df = pd.read_excel(str(url + "/" + files[0]).replace("#", "%23"))
-    filtered_df = df[(((df["Type"] == "CR") & ((df["CR category"] == "B") | (df["CR category"] == "C"))) | (df["Type"] == "pCR")) & ~(df["Uploaded"].isna())][["TDoc", "Title", "CR category", "Source", "Type", "Agenda item", "Agenda item description", "TDoc Status"]]
+    filtered_df = df[(((df["Type"] == "CR") & ((df["CR category"] == "B") | (df["CR category"] == "C"))) | (df["Type"] == "pCR")) & ~(
+        df["Uploaded"].isna())][["TDoc", "Title", "CR category", "Source", "Type", "Agenda item", "Agenda item description", "TDoc Status"]]
     filtered_df["URL"] = filtered_df["TDoc"].apply(gen_url)
 
     df = filtered_df.fillna("")
     return DataResponse(data=df[["TDoc", "Title", "Type", "TDoc Status", "Agenda item description", "URL"]].to_dict(orient="records"))
 
+
 @app.post("/download_tdocs")
 def download_tdocs(req: DownloadRequest):
     documents = req.documents
@@ -290,13 +334,17 @@ def download_tdocs(req: DownloadRequest):
         media_type="application/zip"
     )
 
+
 @app.post("/generate_requirements", response_model=RequirementsResponse)
 async def gen_reqs(req: RequirementsRequest, background_tasks: BackgroundTasks):
     documents = req.documents
     n_docs = len(documents)
+
+    logging.info("Generating requirements for documents: {}".format([doc.document for doc in documents]))
+
     def prompt(doc_id, full):
         return f"Here's the document whose ID is {doc_id} : {full}\n\nExtract all requirements and group them by context, returning a list of objects where each object includes a document ID, a concise description of the context where the requirements apply (not a chapter title or copied text), and a list of associated requirements; always return the result as a list, even if only one context is found. Remove the errors"
-    
+
     async def process_document(doc):
         doc_id = doc.document
         url = doc.url
@@ -305,13 +353,14 @@ async def gen_reqs(req: RequirementsRequest, background_tasks: BackgroundTasks):
         except Exception as e:
             traceback.print_exception(e)
             return RequirementsResponse(requirements=[DocRequirements(document=doc_id, context="Error LLM", requirements=[])]).requirements
-        
+
         try:
             model_used = "gemini-v2"  # À adapter si fallback activé
             async with limiter_mapping[model_used]:
                 resp_ai = await llm_router.acompletion(
                     model=model_used,
-                    messages=[{"role":"user","content": prompt(doc_id, full)}],
+                    messages=[
+                        {"role": "user", "content": prompt(doc_id, full)}],
                     response_format=RequirementsResponse
                 )
             return RequirementsResponse.model_validate_json(resp_ai.choices[0].message.content).requirements
@@ -322,7 +371,8 @@ async def gen_reqs(req: RequirementsRequest, background_tasks: BackgroundTasks):
                     async with limiter_mapping[model_used]:
                         resp_ai = await llm_router.acompletion(
                             model=model_used,
-                            messages=[{"role":"user","content": prompt(doc_id, full)}],
+                            messages=[
+                                {"role": "user", "content": prompt(doc_id, full)}],
                             response_format=RequirementsResponse
                         )
                     return RequirementsResponse.model_validate_json(resp_ai.choices[0].message.content).requirements
@@ -332,46 +382,50 @@ async def gen_reqs(req: RequirementsRequest, background_tasks: BackgroundTasks):
             else:
                 traceback.print_exception(e)
                 return RequirementsResponse(requirements=[DocRequirements(document=doc_id, context="Error LLM", requirements=[])]).requirements
-    
+
     async def process_batch(batch):
         results = await asyncio.gather(*(process_document(doc) for doc in batch))
         return [item for sublist in results for item in sublist]
-    
+
     all_requirements = []
-    
+
     if n_docs <= 30:
         batch_results = await process_batch(documents)
         all_requirements.extend(batch_results)
     else:
         batch_size = 30
-        batches = [documents[i:i + batch_size] for i in range(0, n_docs, batch_size)]
-        
+        batches = [documents[i:i + batch_size]
+                   for i in range(0, n_docs, batch_size)]
+
         for i, batch in enumerate(batches):
             batch_results = await process_batch(batch)
             all_requirements.extend(batch_results)
-            
+
             if i < len(batches) - 1:
                 background_tasks.add_task(asyncio.sleep, 60)
     return RequirementsResponse(requirements=all_requirements)
 
+
 @app.post("/get_reqs_from_query", response_model=ReqSearchResponse)
 def find_requirements_from_problem_description(req: ReqSearchRequest):
     requirements = req.requirements
     query = req.query
 
-    requirements_text = "\n".join([f"[Selection ID: {r.req_id} | Document: {r.document} | Context: {r.context} | Requirement: {r.requirement}]" for r in requirements])
-    
+    requirements_text = "\n".join(
+        [f"[Selection ID: {r.req_id} | Document: {r.document} | Context: {r.context} | Requirement: {r.requirement}]" for r in requirements])
     print("Called the LLM")
     resp_ai = llm_router.completion(
         model="gemini-v2",
-        messages=[{"role":"user","content": f"Given all the requirements : \n {requirements_text} \n and the problem description \"{query}\", return a list of 'Selection ID' for the most relevant corresponding requirements that reference or best cover the problem. If none of the requirements covers the problem, simply return an empty list"}],
+        messages=[{"role": "user", "content": f"Given all the requirements : \n {requirements_text} \n and the problem description \"{query}\", return a list of 'Selection ID' for the most relevant corresponding requirements that reference or best cover the problem. If none of the requirements covers the problem, simply return an empty list"}],
         response_format=ReqSearchLLMResponse
     )
     print("Answered")
     print(resp_ai.choices[0].message.content)
-    
-    out_llm = ReqSearchLLMResponse.model_validate_json(resp_ai.choices[0].message.content).selected
+
+    out_llm = ReqSearchLLMResponse.model_validate_json(
+        resp_ai.choices[0].message.content).selected
     if max(out_llm) > len(requirements) - 1:
-        raise HTTPException(status_code=500, detail="LLM error : Generated a wrong index, please try again.")
+        raise HTTPException(
+            status_code=500, detail="LLM error : Generated a wrong index, please try again.")
 
-    return ReqSearchResponse(requirements=[requirements[i] for i in out_llm])
+    return ReqSearchResponse(requirements=[requirements[i] for i in out_llm])