-
Notifications
You must be signed in to change notification settings - Fork 0
/
app.py
85 lines (66 loc) · 3.41 KB
/
app.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
#
# Copyright (c) 2024 swsychen.
# All rights reserved. Licensed under the MIT license.
# See LICENSE file in the project root for details.
#
import os
import streamlit as st
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import ChatMessage
from langchain_openai import ChatOpenAI
from langchain_pinecone import PineconeVectorStore
from langchain_community import vectorstores
from langchain_openai import OpenAIEmbeddings
from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA,RetrievalQAWithSourcesChain
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") or "OPENAI_API_KEY"
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY") or "PINECONE_API_KEY"
PINECONE_ENVIRONMENT = os.getenv("PINECONE_ENVIRONMENT") or "PINECONE_ENVIRONMENT"
model_name = "text-embedding-ada-002"
text_field = "flow_document" # default value for estuary flow
index_name = "YOUR_INDEX_NAME" # Replace with your index name
namespace = "" # Replace with your namespace name, empty for free tier
def configure_retriever(model_name,openai_api_key,pinecone_api_key,namespace,index_name,text_field,stream_handler):
"""
Here we define the retriever chain with the OpenAI model and Pinecone vector store.
"""
embed = OpenAIEmbeddings(model=model_name, openai_api_key=openai_api_key)
vectorstore= PineconeVectorStore(
pinecone_api_key=pinecone_api_key,namespace=namespace,index_name=index_name,embedding=embed,text_key=text_field
)
llm = ChatOpenAI(openai_api_key=openai_api_key, streaming=True, callbacks=[stream_handler])
retriever = RetrievalQA.from_chain_type(llm=llm, retriever=vectorstore.as_retriever(search_kwargs={'k': 10}))
return retriever
class StreamHandler(BaseCallbackHandler):
"""
Here is the stream handler that will be used to stream the output of the model to the UI word by word.
"""
def __init__(self, container, initial_text=""):
self.container = container
self.text = initial_text
def on_llm_new_token(self, token: str, **kwargs) -> None:
self.text += token
self.container.markdown(self.text)
# with st.sidebar:
# openai_api_key = st.text_input("OpenAI API Key", type="password")
openai_api_key=OPENAI_API_KEY
# The majority of the code below is for the Streamlit UI
# We get the user prompt from the chat and input it to the retriever chain
# The model response from the chain is then streamed to the UI.
if "messages" not in st.session_state:
st.session_state["messages"] = [ChatMessage(role="assistant", content="How can I help you?")]
# The function of this loop is to preserve the chat history and show them always
# each `st.session_state.messages.append` will add a new message to the chat history
for msg in st.session_state.messages:
st.chat_message(msg.role).write(msg.content)
if prompt := st.chat_input():
st.session_state.messages.append(ChatMessage(role="user", content=prompt))
st.chat_message("user").write(prompt)
if not openai_api_key:
st.info("Please add your OpenAI API key to continue.")
st.stop()
with st.chat_message("assistant"):
stream_handler = StreamHandler(st.empty())
qa=configure_retriever(model_name,openai_api_key,PINECONE_API_KEY,namespace,index_name,text_field,stream_handler)
response = qa.invoke(prompt)
st.session_state.messages.append(ChatMessage(role="assistant", content=response['result']))