File size: 6,842 Bytes
0fc97a4 |
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 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 |
import gradio as gr
from core.rag_agent import RAGAgent
from core.document_manager import DocumentManager
import os
# Initialize components
doc_manager = DocumentManager()
rag_agent = None
def initialize_agent():
"""Initialize RAG agent lazily"""
global rag_agent
if rag_agent is None:
rag_agent = RAGAgent()
return rag_agent
def upload_files(files):
"""Handle file uploads"""
if not files:
return "No files selected", get_file_list()
results = []
for file in files:
try:
result = doc_manager.add_document(file.name)
results.append(result)
except Exception as e:
results.append(f"Error processing {os.path.basename(file.name)}: {str(e)}")
return "\n".join(results), get_file_list()
def get_file_list():
"""Get list of documents in the knowledge base"""
try:
files = doc_manager.list_documents()
if not files:
return "No documents in knowledge base"
return "\n".join([f"• {f}" for f in files])
except Exception as e:
return f"Error listing files: {str(e)}"
def clear_database():
"""Clear all documents from the knowledge base"""
try:
result = doc_manager.clear_all()
return result, get_file_list()
except Exception as e:
return f"Error clearing database: {str(e)}", get_file_list()
def chat_with_agent(message, history):
"""Handle chat interactions with the RAG agent"""
if not message.strip():
return history
try:
agent = initialize_agent()
# Stream the agent's response
response_text = ""
for event in agent.agent_graph.stream(
{"messages": [("user", message)]},
agent.get_config(),
stream_mode="values"
):
if "messages" in event and len(event["messages"]) > 0:
last_message = event["messages"][-1]
if hasattr(last_message, "content"):
response_text = last_message.content
if not response_text:
response_text = "I apologize, but I couldn't generate a response. Please try again."
return response_text
except Exception as e:
return f"Error: {str(e)}"
def reset_conversation():
"""Reset the conversation thread"""
global rag_agent
if rag_agent:
rag_agent.reset_thread()
return None # Clear chat history
def create_gradio_ui():
"""Create the complete Gradio interface"""
with gr.Blocks(title="RAG Agent with Agentic Memory", theme=gr.themes.Soft()) as demo:
gr.Markdown("""
# 🤖 RAG Agent with Agentic Memory
Upload documents and chat with an intelligent agent that uses:
- 📚 **Local Knowledge Base** (ChromaDB)
- 🔍 **Web Search** (Tavily)
- 📖 **Wikipedia**
- 🎓 **ArXiv** (Academic Papers)
""")
with gr.Tabs():
# Documents Tab
with gr.Tab("📄 Documents"):
gr.Markdown("### Upload and Manage Documents")
gr.Markdown("Upload PDF or Markdown files to add them to the knowledge base.")
with gr.Row():
with gr.Column(scale=2):
file_upload = gr.File(
label="Upload Documents",
file_count="multiple",
file_types=[".pdf", ".md"]
)
upload_btn = gr.Button("📤 Add to Knowledge Base", variant="primary")
upload_status = gr.Textbox(label="Upload Status", lines=3)
with gr.Column(scale=1):
file_list = gr.Textbox(
label="Documents in Knowledge Base",
lines=10,
value=get_file_list()
)
refresh_btn = gr.Button("🔄 Refresh List")
clear_btn = gr.Button("🗑️ Clear All Documents", variant="stop")
# Connect document management buttons
upload_btn.click(
fn=upload_files,
inputs=[file_upload],
outputs=[upload_status, file_list]
)
refresh_btn.click(
fn=get_file_list,
outputs=[file_list]
)
clear_btn.click(
fn=clear_database,
outputs=[upload_status, file_list]
)
# Chat Tab
with gr.Tab("💬 Chat"):
gr.Markdown("### Chat with Your Documents")
gr.Markdown("Ask questions about your documents or any topic. The agent will search multiple sources.")
chatbot = gr.Chatbot(
label="Conversation",
height=500,
show_label=True,
avatar_images=(None, "🤖")
)
with gr.Row():
msg = gr.Textbox(
label="Your Message",
placeholder="Ask me anything about your documents or general knowledge...",
scale=4
)
submit_btn = gr.Button("Send", variant="primary", scale=1)
with gr.Row():
clear_chat_btn = gr.Button("🔄 Reset Conversation")
gr.Markdown("*Note: Resetting clears the conversation history*")
# Chat interface
chat_interface = gr.ChatInterface(
fn=chat_with_agent,
chatbot=chatbot,
textbox=msg,
submit_btn=submit_btn,
retry_btn=None,
undo_btn=None,
clear_btn=None
)
clear_chat_btn.click(
fn=reset_conversation,
outputs=[chatbot]
)
gr.Markdown("""
---
### 🔧 How it works:
1. **Upload documents** in the Documents tab
2. **Ask questions** in the Chat tab
3. The agent will:
- Analyze your query
- Search relevant sources
- Provide comprehensive answers with citations
""")
return demo
if __name__ == "__main__":
demo = create_gradio_ui()
demo.launch(share=False, server_name="127.0.0.1", server_port=7860) |