Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from llama_index import ( | |
| VectorStoreIndex, | |
| SummaryIndex, | |
| ServiceContext, | |
| StorageContext, | |
| Document, | |
| load_index_from_storage, ChatPromptTemplate, | |
| ) | |
| from llama_index import set_global_service_context | |
| from llama_index.embeddings.openai import OpenAIEmbedding | |
| from llama_index import SimpleDirectoryReader | |
| from llama_index.llms import OpenAI, ChatMessage, MessageRole | |
| import gradio as gr | |
| from typing import List, cast, Optional | |
| service_context = ServiceContext.from_defaults( | |
| chunk_size=512, | |
| llm=OpenAI(model="gpt-4"), | |
| embed_model=OpenAIEmbedding(), | |
| system_prompt=""" | |
| You are tasked with answering queries related to the content found within a provided PDF document. | |
| When responding to queries, you must ALWAYS utilize the tools available to you to analyze the content of the PDF. | |
| NEVER provide an answer based solely on general knowledge or without employing the tools designed for this task. | |
| Your responses should be informed by the information extracted from the PDF content using these tools, ensuring accuracy and relevance to the user's query. | |
| Remember, the data source for your answers is implicit in the queries you receive, and you should not reference or assume any external data sources when formulating your responses. Your primary objective is to provide precise and contextually correct information based on the PDF content analysis. | |
| """, | |
| ) | |
| # Function to load the dataframe and chat with it | |
| def chat_with_dataframe(question): | |
| storage_context = StorageContext.from_defaults( | |
| persist_dir=str("storage") | |
| ) | |
| vector_index = cast(VectorStoreIndex, load_index_from_storage(storage_context)) | |
| vector_query_engine = vector_index.as_query_engine( | |
| similarity_top_k=2 | |
| ) | |
| response = vector_query_engine.query(question) | |
| return response.response | |
| set_global_service_context(service_context) | |
| # Streamlit interface | |
| def main(): | |
| st.sidebar.title('Virtimo - INUBIT') | |
| st.sidebar.subheader('Ask me anything about the product!') | |
| # File uploader in the sidebar | |
| # st.title('Chat with Excel') | |
| # Store LLM generated responses | |
| if "messages" not in st.session_state.keys(): | |
| st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}] | |
| # Display or clear chat messages | |
| for message in st.session_state.messages: | |
| with st.chat_message(message["role"]): | |
| st.write(message["content"]) | |
| # # Initialize conversation in session state if it doesn't exist | |
| # if 'conversation' not in st.session_state: | |
| # st.session_state['conversation'] = [] | |
| # | |
| # # Initialize user input in session state if it doesn't exist | |
| # if 'user_input' not in st.session_state: | |
| # st.session_state['user_input'] = '' | |
| # Chat interface similar to the provided image | |
| # if st.session_state.uploaded_file: | |
| # User-provided prompt | |
| if prompt := st.chat_input(): | |
| st.session_state.messages.append({"role": "user", "content": prompt}) | |
| with st.chat_message("user"): | |
| st.write(prompt) | |
| # Generate a new response if last message is not from assistant | |
| if st.session_state.messages[-1]["role"] != "assistant": | |
| with st.chat_message("assistant"): | |
| with st.spinner("Thinking..."): | |
| response = chat_with_dataframe(prompt) | |
| placeholder = st.empty() | |
| # full_response = '' | |
| if response is not None: | |
| st.write(response) | |
| message = {"role": "assistant", "content": response} | |
| st.session_state.messages.append(message) | |
| # # Check if the response is iterable | |
| # if isinstance(response, (list, tuple)): | |
| # full_response = ''.join(response) | |
| # elif isinstance(response, str): | |
| # full_response = response | |
| # else: | |
| # full_response = str(response) # Convert non-string, | |
| # # placeholder.markdown(full_response) | |
| # Run the main function | |
| if __name__ == '__main__': | |
| main() | |