Spaces:
Running
Running
File size: 5,622 Bytes
deb090d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
import gradio as gr
from rag import RAG # Assuming your rag.py is in the same directory
import os
import time
# --- Configuration ---
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
raise ValueError("GOOGLE_API_KEY environment variable not set. Please set it before running the app.")
COLLECTION_NAME = os.getenv("COLLECTION_NAME", "ca_documents")
# --- Initialize RAG System ---
try:
rag = RAG(GOOGLE_API_KEY, COLLECTION_NAME)
except Exception as e:
print(f"Fatal Error initializing RAG system: {e}")
raise
# --- Initial Chat State ---
welcome_message = {
"role": "assistant",
"content": """
👋 **Welcome to your CA Study Assistant!**
"""
}
initial_chat_history = [welcome_message]
# --- Core Functions ---
def upload_file(file):
if file is None:
return """<div class="message-bubble error">
<p><strong>⚠️ Please select a file.</strong></p>
<p>Supported formats: PDF, DOCX, TXT.</p>
</div>"""
try:
start_time = time.time()
file_path = file.name
success = rag.upload_document(file_path)
duration = time.time() - start_time
if success:
return f"""<div class=\"message-bubble success\">
<p><strong>✅ Success!</strong></p>
<p><em>{os.path.basename(file_path)}</em> uploaded.</p>
<p><small>Processed in {duration:.2f}s</small></p>
</div>"""
else:
return f"""<div class=\"message-bubble error\"><p><strong>❌ Upload Failed.</strong></p><p>Please ensure <em>{os.path.basename(file_path)}</em> is a valid file.</p></div>"""
except Exception as e:
return f"""<div class=\"message-bubble error\"><p><strong>❌ An Error Occurred</strong></p><p><small>{str(e)}</small></p></div>"""
def chat_with_docs(message: str, history: list):
if not message or not message.strip():
return "", history
history.append({"role": "user", "content": message})
answer = rag.ask_question(message)
history.append({"role": "assistant", "content": answer})
return "", history
def clear_chat():
return initial_chat_history, ""
# --- Load Custom CSS ---
with open("style.css") as f:
css = f.read()
# --- Gradio App Layout ---
with gr.Blocks(css=css, title="CA Study Assistant", theme=gr.themes.Base(), mode="auto") as app:
with gr.Column(elem_id="app-container"):
gr.HTML("""
<div id="header">
<h1>CA Study Assistant</h1>
<p>The smartest way to study your materials.</p>
</div>
""")
with gr.Tabs(elem_classes="tab-nav"):
with gr.TabItem("💬 Ask Questions", id="chat_tab"):
chatbot = gr.Chatbot(
value=initial_chat_history,
elem_id="chat-history",
show_label=False,
type="messages",
bubble_full_width=False,
)
with gr.Row(elem_id="chat-input-container"):
with gr.Column(scale=10):
chat_input = gr.Textbox(
placeholder="Ask anything about your documents...",
show_label=False, elem_id="chat-input", lines=1, max_lines=5,
)
with gr.Column(scale=1, min_width=60):
send_btn = gr.Button('➤', elem_id="send-btn")
clear_btn = gr.Button("🗑️ Clear", elem_id="clear-btn")
clear_btn.click(fn=clear_chat, outputs=[chatbot, chat_input])
with gr.TabItem("📚 Upload Documents", id="upload_tab"):
with gr.Column(elem_id="upload-area"):
gr.HTML("""
<div id="upload-icon">
<svg xmlns="http://www.w3.org/2000/svg" width="80" height="80" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"></path><polyline points="17 8 12 3 7 8"></polyline><line x1="12" y1="3" x2="12" y2="15"></line></svg>
</div>
<h3>Upload Your Study Materials</h3>
<p>Drag & drop or click below to select a file.</p>
<div style="margin-top: 0.5rem;">
<span style="background:#c7d2fe;color:#5046e5;padding:0.3rem 0.75rem;border-radius:999px;margin-right:0.5rem;">PDF</span>
<span style="background:#c7d2fe;color:#5046e5;padding:0.3rem 0.75rem;border-radius:999px;margin-right:0.5rem;">DOCX</span>
<span style="background:#c7d2fe;color:#5046e5;padding:0.3rem 0.75rem;border-radius:999px;">TXT</span>
</div>
""")
upload_input = gr.File(
label="", file_count="single", file_types=[".pdf", ".docx", ".txt"],
type="filepath",
)
upload_output = gr.HTML(elem_id="upload-output")
upload_input.upload(
fn=upload_file, inputs=[upload_input], outputs=[upload_output],
api_name="upload_document"
)
send_btn.click(
fn=chat_with_docs, inputs=[chat_input, chatbot], outputs=[chat_input, chatbot],
)
chat_input.submit(
fn=chat_with_docs, inputs=[chat_input, chatbot], outputs=[chat_input, chatbot],
)
if __name__ == "__main__":
app.launch(share=True, debug=True)
|