Update base.py

base.py

@@ -1,12 +1,9 @@
 import os
-from groq import Groq
 import xml.etree.ElementTree as ET
-from fpdf import FPDF
+from groq import Groq
 
 # Initialize Groq Client
-client = Groq(
-    api_key=os.environ.get("GROQ_API_KEY"),
-)
+client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
 def extract_uavariable_chunks(xml_file):
     """
@@ -41,51 +38,28 @@ def send_to_llm(client, prompt, chunk):
     )
     return response.choices[0].message.content.strip()
 
-def generate_pdf(output_object, pdf_buffer):
-    """
-    Generates a PDF from the output object containing the natural language descriptions.
-    Writes the PDF into a BytesIO buffer instead of a file.
-    """
-    pdf = FPDF()
-    pdf.set_auto_page_break(auto=True, margin=15)
-    pdf.add_page()
-    pdf.set_font("Arial", size=12)
-
-    for idx, (original, natural_language) in enumerate(output_object, 1):
-        pdf.set_font("Arial", style='B', size=12)
-        pdf.cell(0, 10, f"UAVariable {idx}", ln=True)
-        pdf.set_font("Arial", size=10)
-        pdf.multi_cell(0, 10, f"Original XML:\n{original}", align="L")
-        pdf.ln(5)
-        pdf.set_font("Arial", style='B', size=12)
-        pdf.cell(0, 10, "Natural Language Conversion:", ln=True)
-        pdf.set_font("Arial", size=10)
-        pdf.multi_cell(0, 10, natural_language, align="L")
-        pdf.ln(10)
-
-    # Write the PDF to the buffer
-    pdf.output(pdf_buffer)
-
-
-def process_xml_to_pdf(xml_file, output_pdf):
+def process_xml_and_generate_natural_language(xml_file):
     """
-    Main function to process the XML file, extract UAVariable chunks, 
-    send them to the LLaMA model, and save the results in a PDF.
+    Processes the XML file, extracts UAVariable chunks, sends them to LLaMA model, 
+    and returns the natural language description.
+    Limits the number of UAVariable chunks to 15.
     """
     prompt = ("Convert this xml chunk to natural language. The hierarchy should be taken into account "
               "and no values should be missed. Do not provide extra text. Just the conversion of xml to natural language.")
     
     uavariables = extract_uavariable_chunks(xml_file)
     output_object = []
-    
-    for chunk in uavariables:
+    limit = 15  # Set a limit for the number of variables to be processed
+
+    for idx, chunk in enumerate(uavariables):
+        if idx >= limit:
+            break  # Stop processing after 15 UAVariables
+
         try:
             natural_language = send_to_llm(client, prompt, chunk)
-            output_object.append((chunk, natural_language))
+            output_object.append(natural_language)
         except Exception as e:
             print(f"Error processing chunk: {e}")
             continue
     
-    # Generate PDF
-    generate_pdf(output_object, output_pdf)
-
+    return output_object