Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| from dataset_loader import load_dataset | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| if not os.path.isdir("low_partners_chatbot_data"): | |
| load_dataset() | |
| from low_partners_chatbot_data.agent import Agent | |
| import gradio as gr | |
| import fitz # PyMuPDF | |
| from PIL import Image | |
| import io | |
| import os | |
| from uuid import uuid4 | |
| agent = Agent() | |
| def extract_pages_from_multiple_pdfs(pdf_info): | |
| images = [] | |
| for pdf_path, page_numbers in pdf_info.items(): | |
| if not os.path.exists(pdf_path): | |
| print(f"PDF file not found: {pdf_path}") | |
| continue | |
| try: | |
| doc = fitz.open(pdf_path) | |
| for page_num in page_numbers: | |
| if 1 <= page_num <= len(doc): | |
| page = doc.load_page(page_num - 1) | |
| pix = page.get_pixmap( | |
| matrix=fitz.Matrix(2, 2) | |
| ) # Increase resolution | |
| img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) | |
| images.append( | |
| (img, f"{os.path.basename(pdf_path)} - Page {page_num}") | |
| ) | |
| else: | |
| print(f"Page {page_num} is out of range for {pdf_path}") | |
| doc.close() | |
| except Exception as e: | |
| print(f"An error occurred with {pdf_path}: {str(e)}") | |
| return images | |
| def get_answer(query, chat_history, session_id): | |
| print(f"Session ID: {session_id}") | |
| answer, pdf_info = agent.ask(query, session_id) | |
| images = [] | |
| if pdf_info: | |
| images = extract_pages_from_multiple_pdfs(pdf_info) | |
| # Update chat history | |
| chat_history.append((query, answer)) | |
| return chat_history, images, session_id | |
| with gr.Blocks(title="Low & Partners Chatbot") as demo: | |
| # Initialize session ID | |
| session_id = gr.State(value=lambda: str(uuid4())) | |
| gr.Markdown("## Low & Partners Chatbot") | |
| gr.Markdown( | |
| "This is a chatbot for Low & Partners Law Firm. Ask me anything related to law! \n" | |
| "Note: The chat session persists until you refresh the page." | |
| ) | |
| with gr.Row(): | |
| # Left Column - Chat Interface | |
| with gr.Column(scale=1): | |
| chatbot = gr.Chatbot(label="Conversation History") | |
| query_input = gr.Textbox( | |
| label="Enter your legal query", | |
| placeholder="Type your question here and press Enter...", | |
| lines=3, | |
| ) | |
| submit_btn = gr.Button("Submit", variant="primary") | |
| # Right Column - PDF References | |
| with gr.Column(scale=1): | |
| gallery = gr.Gallery( | |
| label="Reference Documents from Latest Answer", | |
| allow_preview=True, | |
| preview=True, | |
| object_fit="contain", | |
| height=600, | |
| ) | |
| # Event handlers for both button click and Enter key | |
| submit_event = submit_btn.click( | |
| fn=get_answer, | |
| inputs=[query_input, chatbot, session_id], | |
| outputs=[chatbot, gallery, session_id], | |
| ).then(lambda: "", outputs=[query_input]) | |
| query_input.submit( | |
| fn=get_answer, | |
| inputs=[query_input, chatbot, session_id], | |
| outputs=[chatbot, gallery, session_id], | |
| ).then(lambda: "", outputs=[query_input]) | |
| if __name__ == "__main__": | |
| demo.launch() | |