49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
from langchain import hub
|
|
from langchain_community.vectorstores import Chroma
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
from langchain_core.runnables import RunnablePassthrough
|
|
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
|
|
from dotenv import load_dotenv
|
|
|
|
"""
|
|
User facing RAG application. Mostly adapted from https://github.com/wu4f/cs410g-src/blob/main/03_RAG/08_rag_query.py
|
|
Small changes made regarding OpenAI Embedding, and loading env from dotenv.
|
|
|
|
I use the same rag-prompt since it's a good choice
|
|
|
|
"""
|
|
|
|
load_dotenv()
|
|
|
|
def format_docs(docs):
|
|
return "\n\n".join(doc.page_content for doc in docs)
|
|
|
|
vectorstore = Chroma(
|
|
embedding_function=OpenAIEmbeddings(),
|
|
persist_directory="./rag_data/.chromadb"
|
|
)
|
|
prompt = hub.pull("rlm/rag-prompt")
|
|
retriever = vectorstore.as_retriever()
|
|
llm = ChatOpenAI(model="gpt-4")
|
|
|
|
rag_chain = (
|
|
{"context": retriever | format_docs, "question": RunnablePassthrough()}
|
|
| prompt
|
|
| llm
|
|
| StrOutputParser()
|
|
)
|
|
|
|
print("Welcome to the Kerbal Space Program RAG application. I will try to assist you with any questions ")
|
|
document_data_sources = set()
|
|
for doc_metadata in retriever.vectorstore.get()['metadatas']:
|
|
document_data_sources.add(doc_metadata['sourceURL'])
|
|
|
|
while True:
|
|
line = input("llm>> ")
|
|
if line:
|
|
result = rag_chain.invoke(line)
|
|
print(result)
|
|
else:
|
|
break
|
|
|