1 line
11 KiB
JSON
1 line
11 KiB
JSON
{"repo": "ankitacharya01/rag-agent", "files": {"app.py": {"1e9536c2afa1fa544294cf3b2c81a3e8ed52ca03": {"t": 1785674395.3585272, "n": 10474, "h": "7505831178ce79a9", "c": "import streamlit as st\n\nfrom pathlib import Path\nfrom dotenv import load_dotenv \nimport os\nfrom smolagents import Tool, ToolCallingAgent, OpenAIServerModel\nfrom huggingface_hub import InferenceClient\nimport time\nfrom openai import OpenAI\nfrom io import BytesIO\n\nfrom vectorstore import get_vectorstore\nfrom tt_speech import tts, render_audio_bubble\nfrom rag_main import build_agent, run_rag\nload_dotenv() \n\ndef get_tts_client():\n return InferenceClient()\n\n#build cached vectorDB\n@st.cache_resource\ndef build_vectorstore():\n return get_vectorstore()\n\n\n\n#pdfretrievertool taken out\n \n#agent creation\n@st.cache_resource\ndef get_agent(groq_key):\n return build_agent(groq_key)\n\n@st.cache_resource\ndef get_transcription_client():\n return OpenAI(\n api_key=os.getenv(\"GROQ_API_KEY\"),\n base_url=\"https://api.groq.com/openai/v1\"\n )\n\ndef login_verification(token):\n with st.spinner(\"Verifying...\"):\n try:\n login_test = OpenAI(api_key=token, base_url=\"https://api.groq.com/openai/v1\")\n login_test.models.list()\n return True\n except Exception as e: \n return False\n \n\n# login logic flow: check .env for token, if available:try login, \n# else, require login manually by pasting api key from UI\n\nTOKEN_CACHE_FILE = Path(\".cache/groq_token.txt\")\nst.title(\"RAG-based Assistant\")\nif(\"groq_logged_in\" not in st.session_state):\n \n cached_token = TOKEN_CACHE_FILE.read_text().strip() if TOKEN_CACHE_FILE.exists() else None\n\n if cached_token: #first check cache for api key\n st.session_state[\"groq_token\"] = cached_token\n st.session_state[\"groq_logged_in\"] = True\n \n elif os.getenv(\"GROQ_API_KEY\"): #then check env file\n if login_verification(os.getenv(\"GROQ_API_KEY\")):\n st.session_state[\"groq_token\"] = os.getenv(\"GROQ_API_KEY\")\n st.session_state[\"groq_logged_in\"]=True\n print(\"for debug: logged in through env\")\n \n else: #invalid env token, user logs in manually \n st.session_state[\"groq_logged_in\"]=False\n print(\"Invalid token in .env file\")\n \n else:\n st.session_state[\"groq_logged_in\"]=False\n print(\"env login error!\")\n \n \n\nif not st.session_state[\"groq_logged_in\"]:\n with st.sidebar: \n st.link_button(\"Get Free Groq API Key\", \"https://console.groq.com/keys\")\n groq_token = st.text_input(\n \"Groq API Key\",\n type=\"password\",\n help=\"Paste your Groq API Key (gsk_...)\"\n )\n\n #login button\n if st.button(\"login\"):\n if not groq_token:\n st.error(\"Please Enter a Groq API Key!\")\n st.stop()\n if groq_token and not groq_token.startswith(\"gsk_\"):\n st.warning(\"This doesn't look like a valid Groq API Key!\")\n st.stop()\n \n #verification of what user typed\n if not login_verification(groq_token):\n st.error(\"Invalid Groq API Key, Login Failed!\")\n else:\n TOKEN_CACHE_FILE.parent.mkdir(parents=True, exist_ok=True)\n TOKEN_CACHE_FILE.write_text(groq_token)\n\n st.session_state[\"groq_token\"] = groq_token\n st.session_state[\"groq_logged_in\"] = True\n \n \n st.success(\"Successfully logged in!\")\n \n st.rerun()\n \nelse:\n #pdf management sidebar\n PDF_DIR=Path(\"pdfs\")\n PDF_DIR.mkdir(exist_ok=True)\n\n #logout button\n with st.sidebar:\n col1, col2, col3 = st.columns([1, 5, 1])\n with col2:\n if st.button(\"Use a different Groq Key\"):\n if TOKEN_CACHE_FILE.exists():\n TOKEN_CACHE_FILE.unlink()\n st.session_state[\"groq_logged_in\"]=False\n st.session_state.pop(\"groq_token\", None)\n get_agent.clear()\n st.rerun()\n\n #audio playback toggle btn\n col1, col2, col3 = st.columns([1, 3, 2])\n with col2: \n st.markdown(\"### Audio Playback\")\n with col3:\n to_play_audio = st.toggle(\"AudioToggle\", value=True, label_visibility=\"hidden\") \n st.divider()\n st.header(\"PDF Management\")\n\n if \"uploader_key\" not in st.session_state:\n st.session_state[\"uploader_key\"] = 0\n\n\n uploaded_files = st.file_uploader(\n \"Upload PDFs\",\n type=['pdf'],\n accept_multiple_files=True,\n key=f\"uploader_{st.session_state['uploader_key']}\"\n )\n \n if uploaded_files:\n with st.status(\"Uploading and indexing...\", expanded=True) as status:\n for uploaded_file in uploaded_files:\n save_path = PDF_DIR / uploaded_file.name\n\n with open(save_path, \"wb\") as f:\n f.write(uploaded_file.getbuffer())\n \n st.write(\"Rebuilding Index\")\n build_vectorstore.clear()\n get_agent.clear()\n st.session_state[\"agent_ready\"] = False\n get_agent(st.session_state.get(\"groq_token\"))\n st.session_state[\"agent_ready\"] = True\n status.update(label=\"Done!\", state=\"complete\") \n st.session_state[\"uploader_key\"]+=1\n st.success(\"PDFs uploaded.\")\n st.rerun()\n \n st.divider()\n #shows the list of pdfs \n st.subheader(\"Current PDFs\")\n pdfs = sorted(PDF_DIR.glob(\"*.pdf\"))\n if len(pdfs)==0:\n st.info(\"Please upload one or more PDFs to chat with the agent.\")\n st.stop()\n for pdf in pdfs:\n col1, col2 = st.columns([4,1])\n\n with col1:\n st.text(pdf.name)\n with col2:\n if st.button(\"⛔\", key=f\"delete_{pdf.name}\"):\n if pdf.exists():\n pdf.unlink()\n\n build_vectorstore.clear()\n get_agent.clear()\n st.session_state[\"agent_ready\"] = False \n #also clear agent cache if pdf deleted\n\n st.success(\"Deleted pdf\") \n st.rerun()\n \n groq_key = st.session_state.get(\"groq_token\")\n if not groq_key:\n st.error(\"Groq API key missing\")\n st.stop()\n if not st.session_state.get(\"agent_ready\"):\n with st.spinner(\"Loading Knowledge base...\"):\n \n agent = get_agent(groq_key)\n st.session_state[\"agent_ready\"] = True\n else:\n agent = get_agent(groq_key)\n\n\n if \"messages\" not in st.session_state:\n st.session_state.messages=[{\"role\": \"assistant\", \n \"content\": \"How can I help you?\"}]\n\n #display existing messages\n for msg in st.session_state.messages:\n with st.chat_message(msg[\"role\"]):\n st.markdown(msg[\"content\"])\n if msg[\"role\"]==\"assistant\" and msg.get(\"audio\") and to_play_audio:\n render_audio_bubble(msg[\"audio\"], autoplay=False)\n \n \n \n\n #text or audio input\n prompt = st.chat_input(\"Ask a question about your PDFs...\", \n accept_audio=True,\n audio_sample_rate=16000)\n \n if prompt:\n query=None\n \n #voice input\n if prompt.audio:\n with st.spinner(\"Analysing your voice...\"):\n audio_bytes = prompt.audio.getvalue()\n audio_file = BytesIO(audio_bytes)\n audio_file.name = \"audio.wav\"\n\n transcript = get_transcription_client().audio.transcriptions.create(\n model = \"whisper-large-v3-turbo\",\n file = audio_file\n )\n \n query = transcript.text\n\n \n print(\"type of audio:\")\n print(type(prompt.audio))\n\n elif prompt.text:\n query=prompt.text\n \n if query:\n st.session_state.messages.append({\n \"role\": \"user\", \"content\": query\n })\n with st.chat_message(\"user\"):\n st.markdown(query)\n\n try:\n with st.spinner(\"Thinking...\"):\n\n response = run_rag(agent = agent,\n query=query,\n history=st.session_state.messages)\n audio_bytes = None\n with st.chat_message(\"assistant\"):\n \n st.session_state.messages.append({\n \"role\":\"assistant\", \"content\": response, \"audio\": audio_bytes \n })\n st.markdown(response)\n\n if to_play_audio:\n with st.spinner(\"Generating audio...\"):\n try:\n time1 = time.time()\n print(\"time1 checkpoint!\")\n audio_bytes = tts(response) \n print(\"time2 checkpoint!\")\n # autoplay_audio(audio_bytes) \n except Exception as e:\n print(f\"Error occured: {e}\")\n\n time2 = time.time()\n print(\"time3 checkpoint!\")\n if audio_bytes:\n print(\"time4 checkpoint!\")\n render_audio_bubble(audio_bytes, autoplay=True)\n print(\"time5 checkpoint!\")\n time3 = time.time()\n print(f\"TTs took: {time2-time1}\")\n print(f\"Auto play took: {time3-time2}\")\n print(f\"Total: {time3-time1}\") \n else:\n st.info(\"Audio playback is disabled. Enable it in the sidebar to hear the response.\") \n time.sleep(3) \n st.rerun()\n st.session_state.messages[-1][\"audio\"] = audio_bytes\n \n \n except Exception as e:\n st.error(f\"An Error occured: {e}\")\n "}}}} |