1 line
8.2 KiB
JSON
1 line
8.2 KiB
JSON
{"repo": "alayadb-ai/alayalite", "files": {"examples/rag/ui.py": {"87a79a2411e172e5c2c56feef64d14fff25fe389": {"t": 1785768171.08095, "n": 7781, "h": "6fd2dbc2f7aa3753", "c": "# SPDX-FileCopyrightText: 2025 AlayaDB.AI\n#\n# SPDX-License-Identifier: AGPL-3.0-only\n\n\"\"\"\nThis module provides the Streamlit user interface for a Retrieval-Augmented Generation (RAG) Q&A application.\nIt handles file uploads, knowledge base management, and the chat interface.\n\"\"\"\n\nimport json\nimport os\nfrom datetime import datetime\nfrom typing import Callable, Generator, Tuple\n\nimport streamlit as st\nimport torch\nfrom db import insert_text, query_text, reset_db\nfrom docx import Document\nfrom llm import ask_llm\nfrom pypdf import PdfReader\n\ntorch.classes.__path__ = [os.path.join(torch.__path__[0], torch.classes.__file__)]\nUSE_STREAM = True\nst.set_page_config(page_title=\"My RAG\", page_icon=\"🤖\", layout=\"wide\", initial_sidebar_state=\"expanded\")\n\n# Init session\nif \"user\" not in st.session_state:\n st.session_state.update({\"user\": None, \"chat_history\": [], \"current_db\": \"default\"})\n\n\ndef read_file(uploaded_file):\n content = \"\"\n\n file_type = uploaded_file.type\n\n if uploaded_file.name.endswith(\".md\"):\n file_type = \"text/markdown\"\n\n try:\n if file_type in [\"text/plain\", \"text/x-markdown\", \"text/markdown\"]:\n content = uploaded_file.read().decode(\"utf-8\")\n elif file_type == \"application/pdf\":\n pdf_reader = PdfReader(uploaded_file)\n content = \"\\n\".join([page.extract_text() for page in pdf_reader.pages])\n elif file_type == \"application/vnd.openxmlformats-officedocument.wordprocessingml.document\":\n doc = Document(uploaded_file)\n content = \"\\n\".join([para.text for para in doc.paragraphs])\n else:\n raise ValueError(f\"Unsupported file types: {uploaded_file.type} for file {uploaded_file.name}\")\n except Exception as e:\n raise RuntimeError(f\"Fail to read files: {uploaded_file.name}. Type: {uploaded_file.type}.\") from e\n return content\n\n\ndef enhanced_file_processor(file):\n content = read_file(file)\n content = content.replace(\"\\n\", \" \")\n content = \" \".join(content.split())\n return content\n\n\ndef context_aware_query(\n collection_name: str,\n query: str,\n history: list,\n *,\n llm_url: str,\n llm_api_key: str,\n llm_model: str,\n embed_model_path: str,\n is_stream: bool = True,\n) -> Tuple[str | Callable[[], Generator[str, None, None]], str]:\n context = \"\\n\".join([f\"Q: {q}\\nA: {a}\" for q, a, d in history[-3:]])\n enhanced_query = f\"{context}\\n\\nNew question：{query}\"\n\n retrieved_docs = query_text(\n collection_name=collection_name,\n query=enhanced_query,\n embed_model_path=embed_model_path,\n )\n return ask_llm(\n llm_url, llm_api_key, llm_model, query=query, retrieved_docs=retrieved_docs, is_stream=is_stream\n ), retrieved_docs\n\n\ndef main_interface():\n with st.sidebar:\n with st.expander(\"🛠️ Service Setting\", expanded=True):\n llm_url = st.text_input(\"LLM Base URL\", value=\"https://api.lkeap.cloud.tencent.com/v1/chat\")\n llm_api_key = st.text_input(\"LLM API Key\", value=\"Your API Key\")\n llm_model = st.text_input(\"LLM Model\", value=\"deepseek-v3\")\n embed_model_path = st.text_input(\"Embedding Model\", value=\"BAAI/bge-small-zh-v1.5\")\n\n with st.expander(\"📚 Knowledge Base Management\", expanded=True):\n with st.form(\"upload_form\"):\n collection_name = \"rag_collection\"\n uploaded_file = st.file_uploader(\n \"Upload documents\", type=[\"txt\", \"pdf\", \"docx\", \"md\"], accept_multiple_files=True\n )\n if st.form_submit_button(\"🚀 Start processing\"):\n reset_db()\n if uploaded_file:\n success = True\n for file in uploaded_file:\n with st.spinner(f\"Process {file.name}...\"):\n content = enhanced_file_processor(file)\n success = insert_text(\n collection_name=collection_name,\n docs=content,\n embed_model_path=embed_model_path,\n chunksize=256,\n overlap=25,\n )\n if not success:\n break\n if success:\n st.success(\"Document processing completed\")\n else:\n st.error(\"Document processing failed!\")\n else:\n st.error(\"No document uploaded yet!\")\n\n with st.expander(\"🗃️ Dialogue management\", expanded=True):\n if st.button(\"🔄 Clear all records\", use_container_width=True):\n st.session_state.chat_history = []\n st.rerun()\n\n timestamp = datetime.now().strftime(\"%Y%m%d\")\n if st.download_button(\n \"💾 Export all records\",\n data=json.dumps(st.session_state.chat_history),\n file_name=f\"chat_{timestamp}.json\",\n mime=\"application/json\",\n use_container_width=True,\n ):\n st.toast(\"Exported!\")\n\n with st.container():\n st.header(\"💬 RAG QA\")\n\n # Real-time dialog display container\n chat_container = st.container(height=600, border=False)\n # Render the existing history first\n with chat_container:\n for _, (q, a, d) in enumerate(st.session_state.chat_history):\n with st.chat_message(\"user\", avatar=\"👤\"):\n st.markdown(f\"{q}\")\n with st.chat_message(\"assistant\", avatar=\"🤖\"):\n st.markdown(a)\n if d:\n with st.popover(\"References\"):\n st.markdown(d)\n\n # Input processing\n if prompt := st.chat_input(\"Please enter your question...\"):\n # Display user input\n with chat_container:\n with st.chat_message(\"user\", avatar=\"👤\"):\n st.markdown(f\"{prompt}\")\n\n # Add a placeholder response\n with st.chat_message(\"assistant\", avatar=\"🤖\"):\n answer_placeholder = st.empty()\n answer_placeholder.markdown(\"▌\")\n\n # Execute the query (keep the original logic)\n try:\n resp_or_stream, retrieved_docs = context_aware_query(\n collection_name=collection_name,\n query=prompt,\n history=st.session_state.chat_history,\n llm_url=llm_url,\n llm_api_key=llm_api_key,\n llm_model=llm_model,\n embed_model_path=embed_model_path,\n is_stream=USE_STREAM,\n )\n\n # Update the answer\n full_resp = \"\"\n if USE_STREAM:\n for chunk in resp_or_stream():\n full_resp += chunk\n answer_placeholder.markdown(full_resp)\n else:\n full_resp = resp_or_stream\n answer_placeholder.markdown(full_resp)\n\n # Update session history\n st.session_state.chat_history.append((prompt, full_resp, retrieved_docs))\n # pylint: disable=broad-exception-caught\n except Exception as e:\n answer_placeholder.error(f\"Handling failures: {str(e)}\")\n st.session_state.chat_history.append((prompt, f\"⚠️ Error: {str(e)}\", \"\"))\n\n st.rerun()\n\n\ndef main():\n main_interface()\n\n\nif __name__ == \"__main__\":\n main()\n"}}}} |