-
Notifications
You must be signed in to change notification settings - Fork 31
/
load_qa_chain.py
108 lines (88 loc) · 4.07 KB
/
load_qa_chain.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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
# Do the below in supabase
# create table public.document_embeds_openai(id bigint generated by default as identity not null,
# content text null,
# sources text null,
# embedding vector(1536) null,
# metadata jsonb null,
# constraint document_embeds_open_pkey primary key(id) ) tablespace pg_default;
# Also create the below function in supabase
# create or replace function match_documents_chain_embeds( query_embedding vector(768),
# match_thresholdfloat,match_count int)
# returns table(id bigint, sources text,content text, similarity float, metadata jsonb)
# language
# sql stable
# as $$
# select
# document_embeds_openai.id,
# document_embeds_openai.sources,
# document_embeds_openai.content,
# 1 - (document_embeds_openai.embedding <= > query_embedding) as similarity,
# metadata from document_embeds_openai where
# 1 - (document_embeds_openai.embedding <= > query_embedding) > match_threshold
# order by similarity desc limit match_count;
# $$;
import os
from typing import List
from dotenv import load_dotenv
from langchain import OpenAI
from langchain.schema import Document
from supabase import Client, create_client
from embeddings.openaiembeddings import OpenAIEmbeds
from llm_prompt_template.load_qa_prompt_template import LoadQATemplate
load_dotenv()
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
url = os.environ.get('SUPABASE_URL')
key = os.environ.get('SUPABASE_SECRET_KEY')
def ingest_documents(emb: OpenAIEmbeds, documents: List[Document], supabase: Client):
document_list = []
for document in documents:
embedded_document = emb.generate_embeddings([document.page_content])
row = {'content': document.page_content,
'embedding': embedded_document,
'sources': document.metadata['source'],
'metadata': document.metadata
}
document_list.append(row)
row = {}
# data = supabase.table('document_embeds').insert(document_list).execute()
data = supabase.table('document_embeds_openai').insert(document_list).execute()
def retrieve_top_k_docs(emb,supabase: Client, query) -> List[Document]:
query_embedding = emb.generate_embeddings(query)
supabase: Client = create_client(url, key)
resp = supabase.rpc("match_documents_chain_embeds", params={'query_embedding': query_embedding,
'match_threshold': 0.7,
'match_count': 1}).execute()
# print(resp.data)
resp_len = len(resp.data)
# for i in range(resp_len):
# print(resp.data[i]['id'])
# print(resp.data[i]['sources'])
# print(resp.data[i]['content'])
# print(resp.data[i]['similarity'])
documents = [Document(
page_content="source::" + resp.data[i]['sources'] + "content::" + resp.data[i]['content'],
metadata=resp.data[i]['metadata']
) for i in range(len(resp.data))]
return documents
if __name__=="__main__":
supabase: Client = create_client(url, key)
emb = OpenAIEmbeds()
# Run the below commented lines once first time to load the vector table in supabase
# Sample docuemnt to load
# docs = [Document(page_content="Langchain is a LLM fraemwork. It abstracts the LLM model operations. "
# "Thus it helps to adopt a polyglot architecture", metadata={"source": "my_website"})]
#
# ingest_documents(emb=emb,documents=docs,supabase=supabase)
llm = OpenAI(model_name="text-davinci-003",
openai_api_key=OPENAI_API_KEY,
temperature=0,
max_tokens=1000)
from langchain.chains.question_answering import load_qa_chain
query = "What is langchain?"
docs = retrieve_top_k_docs(emb=emb,supabase=supabase,query=query)
prompt_template = LoadQATemplate()
prompt = prompt_template.get_prompt()
print(prompt)
chain = load_qa_chain(llm, chain_type="stuff",prompt=prompt)
resp = chain.run(input_documents=docs, question=query, tone="Sad")
print(resp)