Create gemini_functions.py

helpers/gemini_functions.py

@@ -0,0 +1,81 @@
+# Gemini Functions
+
+import gradio as gr
+import os
+import google.generativeai as genai
+import json
+
+def initialize_gemini():
+    try:
+        genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
+        generation_config = {
+            "temperature": 1,
+            "top_p": 0.95,
+            "top_k": 40,
+            "max_output_tokens": 8192,
+            "response_mime_type": "text/plain",
+        }
+        model = genai.GenerativeModel(
+            model_name="gemini-1.5-pro",
+            generation_config=generation_config,
+        )
+        return model
+    except Exception as e:
+        raise gr.Error(f"Error initializing Gemini: {str(e)}")
+
+def create_prompt(extracted_text: str, path_to_data_to_extract: str) -> str:
+
+    # load data to extract
+    with open(path_to_data_to_extract, 'r', encoding='utf-8') as file:
+        data_to_extract = json.load(file)
+
+    prompt = f"""Tu es un assistant juridique expert en analyse de documents judiciaires français.
+    Je vais te fournir le contenu d'un document judiciaire extrait d'un PDF.
+    Ta tâche est d'analyser ce texte et d'en extraire les informations suivantes de manière précise :
+
+    {json.dumps(data_to_extract, indent=2, ensure_ascii=False)}
+
+    Voici quelques règles à suivre :
+    - Si une information n'est pas présente dans le texte, indique "Non spécifié" pour cette catégorie.
+    - Pour les noms des parties (demandeurs et défendeurs, et leurs avocats), liste tous ceux que tu trouves
+    - Assure-toi de différencier correctement les demandeurs des défendeurs.
+    - Si tu n'es pas sûr d'une information, indique-le clairement.
+
+    Présente tes résultats sous forme de JSON, en utilisant les catégories mentionnées ci-dessus.
+
+    Voici le contenu du document :
+
+    {extracted_text.strip()}
+
+    Analyse ce texte et fournis-moi les informations demandées au format JSON uniquement.""".strip()
+
+    return prompt
+
+def extract_data_with_gemini(text_file_path: str, path_to_data_to_extract: str) -> dict:
+    try:
+        # Initialize Gemini
+        model = initialize_gemini()
+
+        # Read the extracted text
+        with open(text_file_path, 'r', encoding='utf-8') as f:
+            extracted_text = f.read()
+
+        # Create prompt and get response
+        prompt = create_prompt(extracted_text, path_to_data_to_extract)
+        response = model.generate_content(prompt)
+
+        # Parse the JSON response
+        try:
+            # Extract JSON from the response text
+            json_str = response.text
+            if "json" in json_str.lower():
+                json_str = json_str.split("json")[1].split("```")[0]
+            elif "```" in json_str:
+                json_str = json_str.split("```")[1]
+            result = json.loads(json_str)
+        except:
+            result = {"error": "Failed to parse JSON response", "raw_response": response.text}
+
+        return result
+    except Exception as e:
+        raise gr.Error(f"Error in Gemini processing: {str(e)}")