import streamlit as st import os from llama_index import ( ServiceContext, SimpleDirectoryReader, VectorStoreIndex, ) from llama_index.llms import OpenAI from openai import OpenAI client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) # Define Streamlit layout and interaction st.title("Grounded Generations") # Upload PDF uploaded_file = st.file_uploader("Choose a PDF file", type="pdf") @st.cache_resource(show_spinner=False) def load_data(uploaded_file): with st.spinner('Indexing document...'): # Save the uploaded file temporarily with open("temp.pdf", "wb") as f: f.write(uploaded_file.read()) # Read and index documents using SimpleDirectoryReader reader = SimpleDirectoryReader(input_dir="./", recursive=False) docs = reader.load_data() service_context = ServiceContext.from_defaults( llm=OpenAI( model="gpt-3.5-turbo-16k", temperature=0.1, ), system_prompt="You are an AI assistant that uses context from a PDF to assist the user in generating text." ) index = VectorStoreIndex.from_documents(docs, service_context=service_context) return index # Placeholder for document indexing if uploaded_file is not None: index = load_data(uploaded_file) # Take user query input user_query = st.text_input("Search for the products/info you want to use to ground your generated text content:") # Initialize session_state for retrieved_text if not already present if 'retrieved_text' not in st.session_state: st.session_state['retrieved_text'] = '' # Search and display retrieved text if st.button("Retrieve"): with st.spinner('Retrieving text...'): # Use VectorStoreIndex to search query_engine = index.as_query_engine(similarity_top_k=3) st.session_state['retrieved_text'] = query_engine.query(user_query) st.write(f"Retrieved Text: {st.session_state['retrieved_text']}") # store the retrieved text as a streamlit state variable # Select content type content_type = st.selectbox("Select content type:", ["Blog", "Tweet"]) # make some default nonsense # Generate text based on retrieved text and selected content type if st.button("Generate") and content_type: with st.spinner('Generating text...'): # Generate text using OpenAI API try: if content_type == "Blog": prompt = f"Write a blog about 500 words in length using the {st.session_state['retrieved_text']}" #uses content from the stored state elif content_type == "Tweet": prompt = f"Compose a tweet using the {st.session_state['retrieved_text']}" response = client.chat.completions.create(model="gpt-3.5-turbo-16k", messages=[ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": prompt} ]) generated_text = response.choices[0].message.content st.write(f"Generated Text: {generated_text}") except Exception as e: st.write(f"An error occurred: {e}")