smart-rag / app.py
Dashmov's picture
Update app.py
823fab1 verified
import os
from llama_index.core import VectorStoreIndex
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.groq import Groq
from llama_index.core import Settings
from llama_index.readers.file import PyMuPDFReader
from llama_index.core.memory import ChatMemoryBuffer
import gradio as gr
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.llm = Groq(model="llama-3.1-8b-instant")
chat_engine = None
INDUSTRY_PROMPTS = {
"🏠 Real Estate": "You are a professional real estate assistant. Answer questions about properties, contracts, valuations and listings accurately.",
"βš–οΈ Legal": "You are a legal document assistant. Answer questions about contracts and legal terms precisely. Always remind users to consult a qualified lawyer.",
"πŸ’° Finance & DeFi": "You are a financial and DeFi expert. Answer questions about financial reports, tokenomics and investments clearly.",
"πŸ₯ Medical": "You are a medical document assistant. Answer clearly and always remind users to consult a doctor.",
"πŸš— Car Dealership": "You are a car dealership assistant. Answer questions about vehicles, pricing and warranties helpfully.",
"πŸ“„ General": "You are a helpful assistant. Answer questions from the uploaded document clearly."
}
css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
* { box-sizing: border-box; margin: 0; padding: 0; }
body {
font-family: 'Inter', sans-serif !important;
background: linear-gradient(135deg, #0f0c29, #302b63, #24243e) !important;
min-height: 100vh;
}
.gradio-container {
max-width: 860px !important;
margin: 0 auto !important;
padding: 16px !important;
background: transparent !important;
}
.header-box {
text-align: center;
padding: 24px 16px 16px;
color: white;
}
.mascot {
font-size: 72px;
display: block;
margin-bottom: 8px;
animation: float 3s ease-in-out infinite;
}
@keyframes float {
0% { transform: translateY(0px); }
50% { transform: translateY(-10px); }
100% { transform: translateY(0px); }
}
.app-title {
font-size: clamp(22px, 5vw, 40px) !important;
font-weight: 700 !important;
background: linear-gradient(90deg, #a78bfa, #60a5fa, #34d399);
-webkit-background-clip: text;
-webkit-text-fill-color: transparent;
margin-bottom: 6px;
}
.app-subtitle {
font-size: clamp(13px, 2vw, 15px);
color: #94a3b8;
}
.card {
background: rgba(255,255,255,0.06) !important;
border: 1px solid rgba(255,255,255,0.12) !important;
border-radius: 16px !important;
padding: 16px !important;
margin-bottom: 14px !important;
backdrop-filter: blur(12px);
}
label {
color: #e2e8f0 !important;
font-weight: 500 !important;
font-size: 13px !important;
}
footer { display: none !important; }
@media (max-width: 640px) {
.gradio-container { padding: 8px !important; }
.mascot { font-size: 52px; }
}
"""
def process_pdf(pdf_file, industry):
global chat_engine
if pdf_file is None:
return "⚠️ Please upload a PDF first!"
loader = PyMuPDFReader()
documents = loader.load(file_path=pdf_file.name)
index = VectorStoreIndex.from_documents(documents)
memory = ChatMemoryBuffer.from_defaults(token_limit=3000)
system_prompt = INDUSTRY_PROMPTS.get(industry, list(INDUSTRY_PROMPTS.values())[-1])
chat_engine = index.as_chat_engine(
chat_mode="context",
memory=memory,
system_prompt=system_prompt
)
return f"βœ… Document loaded! Ask me anything."
def respond(message, chat_history):
if not message.strip():
return "", chat_history
if chat_engine is None:
chat_history.append({"role": "user", "content": message})
chat_history.append({"role": "assistant", "content": "⚠️ Please upload a PDF and select an industry first!"})
return "", chat_history
response = chat_engine.chat(message)
chat_history.append({"role": "user", "content": message})
chat_history.append({"role": "assistant", "content": str(response)})
return "", chat_history
with gr.Blocks(css=css) as demo:
gr.HTML("""
<div class="header-box">
<span class="mascot">πŸ€–</span>
<div class="app-title">Smart RAG Assistant</div>
<div class="app-subtitle">Select your industry β€’ Upload your document β€’ Ask anything!</div>
</div>
""")
with gr.Group(elem_classes="card"):
industry_select = gr.Radio(
choices=list(INDUSTRY_PROMPTS.keys()),
value="πŸ“„ General",
label="🏒 Select Your Industry",
)
with gr.Group(elem_classes="card"):
pdf_input = gr.File(
label="πŸ“Ž Upload Your PDF Document",
file_types=[".pdf"]
)
status = gr.Textbox(
label="Status",
interactive=False,
placeholder="Waiting for document..."
)
pdf_input.change(
fn=process_pdf,
inputs=[pdf_input, industry_select],
outputs=status
)
with gr.Group(elem_classes="card"):
chatbot = gr.Chatbot(
height=380,
label="πŸ’¬ Chat"
)
msg = gr.Textbox(
placeholder="Ask a question about your document...",
label="Your Message",
lines=1,
max_lines=4
)
with gr.Row():
submit = gr.Button("Send πŸš€", variant="primary")
clear = gr.Button("Clear πŸ—‘οΈ")
submit.click(fn=respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
msg.submit(fn=respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
clear.click(lambda: [], outputs=chatbot)
demo.launch(share=True)