Spaces:
Sleeping
Sleeping
File size: 5,429 Bytes
3bcf448 |
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 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import gradio as gr
from agent import Agent
from config import Config
from memory.chat_memory import MemoryManager
from utils.html_template import HtmlTemplates
from vector_db.qdrant_db import QdrantDBClient
from langchain_core.messages import HumanMessage, AIMessage
class WebApp:
def __init__(self):
self.title = "RAGent Chatbot"
self.uploaded_files = None
self.upload_btn = None
self.progress_output = None
self.status_output = None
self.css = HtmlTemplates.css()
self.agent = Agent()
self.memory = MemoryManager()
self.qdrant_client = QdrantDBClient()
def build_ui(self):
with gr.Blocks(theme=gr.themes.Default(), css=self.css) as demo:
self.build_header()
with gr.Row():
self.build_upload_section()
self.build_chat_section()
return demo
def build_header(self):
with gr.Row():
with gr.Column():
gr.HTML(f"<h1 id='title'>💬 {self.title}</h1>")
def clear_outputs(self):
return "", ""
def build_upload_section(self):
with gr.Column(scale=3):
gr.Markdown("### 📂 Drag & Drop Files Below")
self.uploaded_files = gr.File(
file_types=Config.FILE_EXTENSIONS,
file_count="multiple",
label="pdf, docx, xlsx, pptx, csv, txt, json"
)
self.upload_btn = gr.Button(value="Upload Files", elem_id="upload-btn", icon=Config.UPLOAD_ICON)
self.progress_output = gr.HTML()
self.status_output = gr.Markdown()
self.upload_btn.click(
fn=self.clear_outputs,
inputs=[],
outputs=[self.progress_output, self.status_output]
).then(
fn=self.upload_and_process,
inputs=self.uploaded_files,
outputs=[self.progress_output, self.status_output],
show_progress="hidden"
)
def build_chat_section(self):
with gr.Column(scale=7):
gr.Markdown("### 🤖 Ask Your Question")
gr.ChatInterface(
fn=self.run_agent,
type="messages",
show_progress="full",
save_history=False,
)
def run_agent(self, query, history):
session_id = Config.SESSION_ID
# Get history
past_messages = self.memory.get(session_id)
# Run agent (it appends the user query internally)
response = self.agent.run(query, past_messages)
#print("##### response : ", response)
# convert response to string. If response is a dict like {'input': ..., 'output': ...}
if isinstance(response, dict) and "output" in response:
answer = response["output"]
else:
answer = str(response)
# Save user + assistant message to memory
self.memory.add(session_id, HumanMessage(content=query))
self.memory.add(session_id, AIMessage(content=answer))
return f"🤖 {answer}"
def upload_and_process(self, files):
if not files or len(files) == 0:
yield HtmlTemplates.error_bar(), ""
return
total = len(files)
failed_files = []
for i, file in enumerate(files):
file_path = file.name # path to temp file
try:
# Load, chunk, and insert to vector DB
file_chunks = self.qdrant_client.load_and_chunk_docs(file_path)
self.qdrant_client.insert_chunks(file_chunks)
except Exception as e:
failed_files.append(file_path)
yield HtmlTemplates.progress_bar(int((i + 1) / total * 100), i + 1, total), (
f"⚠️ Skipped file {i + 1}/{total}: {os.path.basename(file_path)} - {str(e)}"
)
continue
percent = int((i + 1) / total * 100)
yield HtmlTemplates.progress_bar(percent, i + 1, total), f"📄 Processed {i + 1}/{total} file(s)..."
success_count = total - len(failed_files)
final_msg = f"✅ {success_count}/{total} file(s) processed and stored in DB!"
if failed_files:
failed_list = "\n".join(f"❌ {os.path.basename(f)}" for f in failed_files)
final_msg += f"\n\n⚠️ Failed to process:\n{failed_list}"
yield HtmlTemplates.progress_bar(100, total, total), final_msg
def upload_and_process1(self, files):
if not files or len(files) == 0:
yield HtmlTemplates.error_bar(), ""
return
total = len(files)
for i, file in enumerate(files):
file_path = file.name # get file path of temporary folder
# Load, chunk, and insert to vector DB
file_chunks = self.qdrant_client.load_and_chunk_docs(file_path)
self.qdrant_client.insert_chunks(file_chunks)
percent = int((i + 1) / total * 100)
yield HtmlTemplates.progress_bar(percent, i + 1, total), f"📄 Processed {i + 1}/{total} file(s)..."
yield HtmlTemplates.progress_bar(100, total, total), f"✅ {total} file(s) processed and stored in DB!"
if __name__ == "__main__":
app = WebApp()
demo = app.build_ui()
demo.launch()
|