|
import streamlit as st |
|
from datetime import time as t |
|
import time |
|
|
|
from operator import itemgetter |
|
import os |
|
import json |
|
import getpass |
|
import openai |
|
|
|
from langchain.vectorstores import Pinecone |
|
from langchain.embeddings import OpenAIEmbeddings |
|
import pinecone |
|
|
|
|
|
from results import results_agent |
|
from filter import filter_agent |
|
from reranker import reranker |
|
from utils import build_filter, clean_pinecone |
|
from router import routing_agent |
|
|
|
OPENAI_API = st.secrets["OPENAI_API"] |
|
PINECONE_API = st.secrets["PINECONE_API"] |
|
openai.api_key = OPENAI_API |
|
|
|
|
|
pinecone.init( |
|
api_key= PINECONE_API, |
|
environment="gcp-starter" |
|
) |
|
index_name = "use-class-db" |
|
|
|
embeddings = OpenAIEmbeddings(openai_api_key = OPENAI_API) |
|
|
|
index = pinecone.Index(index_name) |
|
|
|
k = 5 |
|
|
|
st.title("USC GPT - Find the perfect class") |
|
|
|
class_time = st.slider( |
|
"Filter Class Times:", |
|
value=((11, 30), (12, 45)), |
|
min_value=(0, 0), |
|
max_value=(23, 59), |
|
step=(1, 15), |
|
) |
|
|
|
units = st.slider( |
|
"Number of units", |
|
1, 4, 4 |
|
) |
|
|
|
days = st.multiselect("What days are you free?", |
|
options = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat"], |
|
default = None, |
|
placeholder = "Any day" |
|
) |
|
|
|
assistant = st.chat_message("assistant") |
|
initial_message = "How can I help you today?" |
|
|
|
|
|
|
|
def get_rag_results(prompt): |
|
''' |
|
1. Remove filters from the prompt to optimize success of the RAG-based step. |
|
2. Query the Pinecone DB and return the top 25 results based on cosine similarity |
|
3. Rerank the results from vector DB using a BERT-based cross encoder |
|
''' |
|
query = filter_agent(prompt, OPENAI_API) |
|
|
|
|
|
days_filter = list() |
|
start = float(class_time[0][0]) + float(class_time[0][1]) / 100.0 |
|
end = float(class_time[1][0]) + float(class_time[1][1]) / 100.0 |
|
query_filter = { |
|
"start": {"$gte": start}, |
|
"end": {"$lte": end} |
|
} |
|
|
|
if units != "any": |
|
query_filter["units"] = str(int(units)) + ".0 units" |
|
|
|
if len(days) > 0: |
|
for i in range(len(days)): |
|
days_filter.append(days[i]) |
|
for j in range(i+1, len(days)): |
|
two_day = days[i] + ", " + days[j] |
|
days_filter.append(two_day) |
|
query_filter["days"] = {"$in": days_filter} |
|
|
|
|
|
response = index.query( |
|
vector = embeddings.embed_query(query), |
|
top_k = 45, |
|
filter = query_filter, |
|
include_metadata = True |
|
) |
|
|
|
response, additional_metadata = clean_pinecone(response) |
|
response = reranker(query, response) |
|
|
|
return response, additional_metadata |
|
|
|
|
|
|
|
if "messages" not in st.session_state: |
|
st.session_state.messages = [] |
|
with st.chat_message("assistant"): |
|
st.markdown(initial_message) |
|
st.session_state.messages.append({"role": "assistant", "content": initial_message}) |
|
|
|
|
|
if prompt := st.chat_input("What kind of class are you looking for?"): |
|
st.session_state.messages.append({"role": "user", "content": prompt}) |
|
with st.chat_message("user"): |
|
st.markdown(prompt) |
|
|
|
with st.chat_message("assistant"): |
|
message_placeholder = st.empty() |
|
full_response = "" |
|
|
|
messages = [{"role": m["role"], "content": m["content"]} |
|
for m in st.session_state.messages] |
|
message_history = " ".join([message["content"] for message in messages]) |
|
|
|
route = routing_agent(prompt, OPENAI_API, message_history) |
|
|
|
if route == "1": |
|
|
|
rag_response, additional_metadata = get_rag_results(prompt) |
|
result_query = 'Original Query:' + prompt + 'Query Results:' + str(rag_response) + '\n Additional Class Times:' + str(additional_metadata) |
|
assistant_response = results_agent(result_query, OPENAI_API) |
|
else: |
|
|
|
assistant_response = openai.ChatCompletion.create( |
|
model = "gpt-4", |
|
messages = [ |
|
{"role": m["role"], "content": m["content"]} |
|
for m in st.session_state.messages |
|
] |
|
)["choices"][0]["message"]["content"] |
|
|
|
|
|
for chunk in assistant_response.split(): |
|
full_response += chunk + " " |
|
time.sleep(0.05) |
|
message_placeholder.markdown(full_response + "▌") |
|
message_placeholder.markdown(full_response) |
|
st.session_state.messages.append({"role": "assistant", "content": full_response}) |
|
|
|
|