Update app.py

app.py

@@ -1,63 +1,177 @@
+import os
+from groq import Groq
 import gradio as gr
-from huggingface_hub import InferenceClient
+import fitz  # PyMuPDF for PDF text extraction
+from PIL import Image
+import pytesseract
+from fpdf import FPDF  # Library for creating PDFs
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# Set your API key
 
+# Initialize the Groq client
+client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+# Initialize the conversation history
+conversation_history = []
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+def extract_text_from_file(uploaded_file):
+    text = ""
+    try:
+        if uploaded_file.name.endswith(".pdf"):
+            doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
+            for page in doc:
+                text += page.get_text()
+        elif uploaded_file.name.endswith((".png", ".jpg", ".jpeg")):
+            image = Image.open(uploaded_file)
+            text = pytesseract.image_to_string(image)
+        else:
+            text = "Unsupported file format. Please upload a PDF or image file."
+    except Exception as e:
+        text = f"Error extracting text: {str(e)}"
+    return text
 
-    messages.append({"role": "user", "content": message})
+def legal_chatbot(user_query, uploaded_file):
+    role_context = (
+        "As a seasoned legal expert with over 30 years of experience specializing in Pakistani law and justice, "
+        "your role is to provide precise, actionable, and lawful advice. "
+        "You have an in-depth understanding of the Pakistan Penal Code, along with other significant legal frameworks, "
+        "and you are well-versed in recent amendments, landmark rulings, "
+        "and their implications. You are expected to respond with clear, comprehensive, and contextually accurate legal guidance. "
+        "Your response should focus on assisting legal professionals, such as lawyers and judges, in understanding and applying the relevant sections, rules, and case laws. "
+        "Ensure that your advice is practical, within the boundaries of law and ethics, and addresses the query with the aim of promoting justice and aiding in the resolution of legal matters. "
+        "If the query is unclear or lacks sufficient detail, provide general guidance or ask for additional information to deliver a more tailored response. "
+        "Your ultimate goal is to support the legal process and help people achieve justice."
+    )
 
-    response = ""
+    if uploaded_file:
+        file_text = extract_text_from_file(uploaded_file)
+        user_query = f"{user_query}\n\nRefer to the following case file content:\n{file_text}"
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+    llm_input = f"{role_context} {user_query}"
+    conversation_history.append({"role": "user", "content": user_query})
 
-        response += token
-        yield response
+    try:
+        chat_completion = client.chat.completions.create(
+            messages=[{"role": "user", "content": llm_input}],
+            model="llama3-8b-8192",
+        )
+        response = chat_completion.choices[0].message.content
+        conversation_history.append({"role": "assistant", "content": response})
+    except Exception as e:
+        response = f"An error occurred while generating the response: {str(e)}"
+        print(response)  # Log the error for debugging
 
+    return response
+
+def save_conversation():
+    pdf = FPDF()
+    pdf.set_auto_page_break(auto=True, margin=15)
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+
+    for message in conversation_history:
+        role = "You" if message["role"] == "user" else "Legal AI Assistant"
+        pdf.multi_cell(0, 10, f"{role}: {message['content']}\n\n")
+
+    pdf_file = "conversation.pdf"
+    pdf.output(pdf_file)
+
+    return pdf_file
+
+def clear_chat():
+    global conversation_history
+    conversation_history = []
+    return ""
+
+css = """
+#generate-button {
+    background-color: #388E3C;  /* Darker green */
+    color: black;
+    border-radius: 5px;
+    padding: 2px 8px;
+    font-size: 16px;  /* Slightly larger font size */
+    margin-top: 5px;
+    height: 35px;
+}
+
+#clear-button {
+    background-color: #D32F2F;  /* Darker red */
+    color: red;
+    border-radius: 5px;
+    padding: 2px 8px;
+    font-size: 16px;
+    margin-top: 5px;
+    height: 35px;
+}
+
+#save-button {
+    background-color: #1976D2;  /* Darker blue */
+    color: green;
+    border-radius: 5px;
+    padding: 2px 8px;
+    font-size: 16px;
+    margin-top: 5px;
+    height: 35px;
+}
+
+#chat-output {
+    overflow-y: auto;
+    height: 400px;
+    width: 100%;
+}
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-
-if __name__ == "__main__":
-    demo.launch()
+
+def gradio_interface():
+    with gr.Blocks(css=css) as demo:
+        with gr.Row():
+            gr.Markdown(f"<h2 style='text-align: center;'>⚖️ Legal AI Assistant Chatbot - Engr. Hamesh Raj</h2>")
+        with gr.Row():
+            gr.Markdown(f"<h4 style='text-align: center;'>Ask any legal questions related to Pakistani law. [Connect on LinkedIn](https://www.linkedin.com/in/hameshraj)</h4>")
+
+        with gr.Row():
+            with gr.Column(scale=0.5, min_width=30):
+                uploaded_file = gr.File(label="Upload a case file (PDF or Image)", elem_id="upload-box")
+                start_new_button = gr.Button("Start New", elem_id="clear-button", scale=0.25)
+                save_button = gr.Button("Save Conversation", elem_id="save-button", scale=0.25)
+
+            with gr.Column(scale=2):
+                user_query = gr.Textbox(label="Enter your legal query here", placeholder="Ask about any law, section, or legal issue...", lines=3, elem_id="query-box")
+                generate_button = gr.Button("Get Legal Advice", elem_id="generate-button")
+                download_link = gr.File(label="Download PDF", elem_id="download-link", visible=False)
+
+        with gr.Row():
+            chat_output = gr.Markdown(elem_id="chat-output", wrap_lines=True)
+
+        def update_chat(user_query, uploaded_file):
+            response = legal_chatbot(user_query, uploaded_file)
+            chat_history = ""
+            for message in conversation_history:
+                role = "You" if message["role"] == "user" else "Legal AI Assistant"
+                chat_history += f"**{role}:** {message['content']}\n\n---\n\n"
+            return chat_history
+
+        def save_and_update():
+            pdf_file = save_conversation()
+            return gr.update(visible=True), pdf_file
+
+        generate_button.click(
+            fn=update_chat,
+            inputs=[user_query, uploaded_file],
+            outputs=[chat_output]
+        )
+
+        start_new_button.click(
+            fn=clear_chat,
+            inputs=None,
+            outputs=[chat_output]
+        )
+
+        save_button.click(
+            fn=save_and_update,
+            inputs=None,
+            outputs=[download_link, download_link]
+        )
+
+    demo.launch(share=True)
+
+gradio_interface()