import gradio as gr import asyncio from rag.pipeline.language_model import LM, LMConfig from rag.retriever.langchain_retriever import LangChainRetriever from rag.inference.inferencer import InferencerConfig, Inferencer from rag import cs_agent, query_maker_agent def test_inference(): """Main function that sets up and runs the RAG chatbot interface""" # Initialize RAG components print("==== Start Inference Test ===") print("RAG system initialized successfully!") def chatbot_response(message, history): """Streaming response menggunakan RAG inferencer""" try: # Create new event loop for this thread loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) async def stream_response(): partial_response = "" # print("message = ", message) # formatted_query = await query_maker_agent.get_result(question = message) # print("Formatted Query = ", formatted_query) # formatted_query = formatted_query['responses'][0]['rag_response'] await cs_agent.load_documents() async for stream_data in cs_agent.get_result(question = message): if stream_data["type"] == "chunk": chunk = stream_data["data"]["chunk"] partial_response += chunk yield partial_response elif stream_data["type"] == "metadata": setup_time = stream_data['data']['setup_time'] print(f"\nSetup completed in {setup_time:.2f}s") elif stream_data["type"] == "complete": total_time = stream_data['data']['total_time'] print(f"\nTotal time: {total_time:.2f}s") async_gen = stream_response() try: while True: result = loop.run_until_complete(async_gen.__anext__()) yield result except StopAsyncIteration: pass finally: loop.close() except Exception as e: yield f"❌ Error: {str(e)}" def add_document_to_vectorstore(file_path): """Add document to vectorstore""" if not file_path: return "⚠️ Please select a file first." try: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) async def add_doc(): result = "" return result result = loop.run_until_complete(add_doc().__anext__()) loop.close() if result.success: return f"✅ Successfully added: {result.document_metadata.file_name} ({result.document_metadata.chunk_count} chunks)" else: return f"❌ Failed to add document: {result.error_message}" except Exception as e: return f"❌ Error adding document: {str(e)}" def clear_chat(): """Function untuk clear chat history""" return [], "" # CSS untuk styling css = """ .gradio-container { max-width: 900px !important; margin: auto !important; } .chat-message { padding: 10px; margin: 5px; border-radius: 10px; } #chatbot { height: 500px; } """ # Membuat interface Gradio with gr.Blocks(css=css, title="RAG Chatbot") as demo: gr.Markdown(""" # 🤖 SakuraAI, Virtual Assistant """) # Status indicator with gr.Row(): status_text = gr.Textbox( value="✅ RAG System Ready", label="System Status", interactive=False, container=True ) with gr.Row(): with gr.Column(scale=2): chatbot = gr.Chatbot( elem_id="chatbot", show_label=False, container=True, bubble_full_width=False, show_copy_button=True, layout="panel" ) with gr.Row(): msg = gr.Textbox( placeholder="Tanyakan sesuatu tentang dokumen Anda...", show_label=False, scale=4, container=False, lines=1, max_lines=3, autofocus=True ) send_btn = gr.Button("Kirim", variant="primary", scale=1) with gr.Row(): clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary") stop_btn = gr.Button("⏹️ Stop", variant="stop", visible=False) # Document management panel with gr.Column(scale=1): gr.Markdown("### 📚 Document Management") with gr.Group(): file_upload = gr.File( label="Upload Document", file_types=[".pdf", ".txt", ".docx"], type="filepath" ) upload_btn = gr.Button("Add to Knowledge Base", variant="secondary") upload_status = gr.Textbox( label="Upload Status", interactive=False, lines=3 ) gr.Markdown(""" ### ⚙️ RAG Settings - **K**: 3 (documents retrieved) - **Template**: Friendly - **Reranking**: Disabled - **Vectorstore**: ChromaDB """) # State untuk tracking is_generating = gr.State(False) # Event handlers untuk chat def user_message(message, history, generating): """Handle user message""" if message.strip() and not generating: history.append([message, None]) return "", history, True, gr.update(visible=True), gr.update(interactive=False) return message, history, generating, gr.update(visible=False), gr.update(interactive=True) def bot_message_stream(history, generating): """Handle streaming bot response""" if history and history[-1][1] is None and generating: user_msg = history[-1][0] for partial_response in chatbot_response(user_msg, history): history[-1][1] = partial_response yield history, True, gr.update(visible=True), gr.update(interactive=False) yield history, False, gr.update(visible=False), gr.update(interactive=True) else: yield history, generating, gr.update(visible=False), gr.update(interactive=True) def stop_generation(): """Stop the generation process""" return False, gr.update(visible=False), gr.update(interactive=True) # Binding events untuk submit message submit_event = msg.submit( user_message, inputs=[msg, chatbot, is_generating], outputs=[msg, chatbot, is_generating, stop_btn, send_btn] ).then( bot_message_stream, inputs=[chatbot, is_generating], outputs=[chatbot, is_generating, stop_btn, send_btn] ) # Binding events untuk send button send_event = send_btn.click( user_message, inputs=[msg, chatbot, is_generating], outputs=[msg, chatbot, is_generating, stop_btn, send_btn] ).then( bot_message_stream, inputs=[chatbot, is_generating], outputs=[chatbot, is_generating, stop_btn, send_btn] ) # Clear chat event clear_btn.click( clear_chat, outputs=[chatbot, msg] ).then( lambda: (False, gr.update(visible=False), gr.update(interactive=True)), outputs=[is_generating, stop_btn, send_btn] ) # Stop generation event stop_btn.click( stop_generation, outputs=[is_generating, stop_btn, send_btn], cancels=[submit_event, send_event] ) # Document upload event upload_btn.click( add_document_to_vectorstore, inputs=[file_upload], outputs=[upload_status] ) # Info panel with gr.Accordion("ℹ️ Info Penggunaan", open=False): gr.Markdown(""" ### Cara Menggunakan: 1. **Chat**: Ketik pertanyaan tentang dokumen yang sudah dimuat 2. **Upload**: Tambahkan dokumen baru ke knowledge base 3. **Stream**: Response akan muncul secara streaming 4. **Stop**: Gunakan tombol stop untuk menghentikan generasi ### Dokumen yang Dimuat: - file2.pdf (dari folder documents) - Dokumen tambahan yang Anda upload ### Teknologi yang Digunakan: - **LLM**: Qwen dengan streaming - **Embedding**: text-embedding-3-small - **Vectorstore**: ChromaDB - **Search**: Hybrid search (dense + sparse) """) # Launch the interface print("Launching Gradio interface...") demo.launch( share=False, server_name="0.0.0.0", server_port=7861, show_error=True, show_api=False )