Spaces:
Sleeping
Sleeping
| import ollama | |
| import streamlit as st | |
| import httpx | |
| import io | |
| import pandas as pd | |
| from PyPDF2 import PdfReader | |
| from PIL import Image | |
| st.set_page_config(layout="wide") | |
| st.title("InsightBot LLM Chatbot Plus Rag") | |
| st.text("Analyzing data for making business-critical decisions and effectively handling complex analysis") | |
| # Initialize history | |
| if "messages" not in st.session_state: | |
| st.session_state["messages"] = [] | |
| # Initialize models and system prompt | |
| if "model" not in st.session_state: | |
| st.session_state["model"] = "" | |
| if "system_prompt" not in st.session_state: | |
| st.session_state["system_prompt"] = "analyze file and summarize in bullet points" | |
| if "new_message" not in st.session_state: | |
| st.session_state["new_message"] = False | |
| if "user_query" not in st.session_state: | |
| st.session_state["user_query"] = "" | |
| if "uploaded_file_content" not in st.session_state: | |
| st.session_state["uploaded_file_content"] = "" | |
| if "uploaded_files" not in st.session_state: | |
| st.session_state["uploaded_files"] = [] | |
| st.sidebar.write("Query Assist AI") | |
| # Sidebar menu | |
| with st.sidebar: | |
| try: | |
| models = [model["name"] for model in ollama.list()["models"]] | |
| st.session_state["model"] = st.selectbox("Choose your model", models) | |
| except httpx.ConnectError: | |
| st.error("Unable to connect") | |
| st.session_state["system_prompt"] = st.text_area("System Prompt", value="analyze file and summarize in bullet points") | |
| if st.button("Reset"): | |
| st.session_state["messages"] = [] | |
| st.session_state["new_message"] = False | |
| st.session_state["user_query"] = "" | |
| st.session_state["uploaded_file_content"] = "" | |
| st.session_state["uploaded_files"] = [] | |
| st.rerun() | |
| uploaded_files = st.file_uploader("Upload images (PNG, JPG) or text files (PDF, CSV)", type=["png", "jpg", "pdf", "csv"], accept_multiple_files=True) | |
| if uploaded_files: | |
| st.session_state["uploaded_files"] = uploaded_files | |
| # Add a radio button to decide if the uploaded file should be part of the query | |
| include_files_in_query = st.radio("Include uploaded files in query?", ("Yes", "No")) | |
| def process_uploaded_files(): | |
| file_contents = [] | |
| image_files = [] | |
| for uploaded_file in st.session_state["uploaded_files"]: | |
| file_type = uploaded_file.type | |
| if "image" in file_type: | |
| image_files.append(uploaded_file) | |
| file_contents.append(f"Image file: {uploaded_file.name}") | |
| elif "pdf" in file_type: | |
| pdf_reader = PdfReader(io.BytesIO(uploaded_file.read())) | |
| pdf_text = "" | |
| for page in pdf_reader.pages: | |
| pdf_text += page.extract_text() | |
| file_contents.append(f"PDF file: {uploaded_file.name}\nContent:\n{pdf_text}") | |
| elif "csv" in file_type: | |
| csv_data = pd.read_csv(uploaded_file) | |
| csv_text = csv_data.to_string() | |
| file_contents.append(f"CSV file: {uploaded_file.name}\nContent:\n{csv_text}") | |
| return "\n".join(file_contents), image_files | |
| def model_res_generator(messages): | |
| try: | |
| stream = ollama.chat( | |
| model=st.session_state["model"], | |
| messages=messages, | |
| stream=True, | |
| ) | |
| for chunk in stream: | |
| yield chunk["message"]["content"] | |
| except (httpx.ConnectError, ollama.ResponseError) as e: | |
| st.error(f"An error occurred: {e}") | |
| return | |
| # Display chat messages from history on app rerun | |
| for message in st.session_state["messages"]: | |
| with st.chat_message(message["role"]): | |
| st.markdown(message["content"]) | |
| # Add some space before the query box | |
| st.write("") | |
| st.write("") | |
| if st.session_state["new_message"]: | |
| st.session_state["user_query"] = "" | |
| st.session_state["new_message"] = False | |
| st.rerun() | |
| if prompt := st.text_input("What is your query?", key="user_query"): | |
| # Process and include uploaded file content in the query if the user chose to include it | |
| if include_files_in_query == "Yes": | |
| st.session_state["uploaded_file_content"], image_files = process_uploaded_files() | |
| augmented_prompt = prompt + "\n\n" + st.session_state["uploaded_file_content"] | |
| else: | |
| augmented_prompt = prompt | |
| # Prepare messages with possible images | |
| messages = [{"role": "user", "content": augmented_prompt}] | |
| if st.session_state["system_prompt"]: | |
| messages.insert(0, {"role": "system", "content": st.session_state["system_prompt"]}) | |
| if include_files_in_query == "Yes" and image_files: | |
| with io.BytesIO() as file_obj: | |
| file_obj.write(image_files[0].read()) | |
| file_obj.seek(0) | |
| messages[0]["images"] = [file_obj.read()] | |
| # Add latest message to history in format {role, content} | |
| st.session_state["messages"].append({"role": "user", "content": prompt}) | |
| with st.chat_message("user"): | |
| st.markdown(prompt) | |
| with st.chat_message("assistant"): | |
| with st.spinner("Thinking..."): | |
| # Generate response based on the augmented prompt | |
| try: | |
| message = "".join(model_res_generator(messages)) | |
| st.session_state["messages"].append({"role": "assistant", "content": message}) | |
| except Exception as e: | |
| st.error(f"Failed to generate response: {e}") | |
| # Set flag for new message | |
| st.session_state["new_message"] = True | |
| st.rerun() | |
| st.sidebar.info("built by dw") |