mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-16 01:55:15 +03:00
92 lines
3.8 KiB
Python
92 lines
3.8 KiB
Python
from langchain.embeddings.openai import OpenAIEmbeddings
|
|
from langchain.schema import Document
|
|
from llm.qa import BrainPicking, BrainSettings
|
|
from llm.summarization import llm_evaluate_summaries, llm_summerize
|
|
from logger import get_logger
|
|
from models.chats import ChatMessage
|
|
from models.settings import BrainSettings, CommonsDep
|
|
from pydantic import BaseModel
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
class Neurons(BaseModel):
|
|
|
|
commons: CommonsDep
|
|
settings = BrainSettings()
|
|
|
|
def create_vector(self, user_id, doc, user_openai_api_key=None):
|
|
logger.info(f"Creating vector for document")
|
|
logger.info(f"Document: {doc}")
|
|
if user_openai_api_key:
|
|
self.commons['documents_vector_store']._embedding = OpenAIEmbeddings(openai_api_key=user_openai_api_key)
|
|
try:
|
|
sids = self.commons['documents_vector_store'].add_documents([doc])
|
|
if sids and len(sids) > 0:
|
|
self.commons['supabase'].table("vectors").update({"user_id": user_id}).match({"id": sids[0]}).execute()
|
|
except Exception as e:
|
|
logger.error(f"Error creating vector for document {e}")
|
|
|
|
def create_embedding(self, content):
|
|
return self.commons['embeddings'].embed_query(content)
|
|
|
|
def similarity_search(self, query, table='match_summaries', top_k=5, threshold=0.5):
|
|
query_embedding = self.create_embedding(query)
|
|
summaries = self.commons['supabase'].rpc(
|
|
table, {'query_embedding': query_embedding,
|
|
'match_count': top_k, 'match_threshold': threshold}
|
|
).execute()
|
|
return summaries.data
|
|
|
|
|
|
def create_summary(commons: CommonsDep, document_id, content, metadata):
|
|
logger.info(f"Summarizing document {content[:100]}")
|
|
summary = llm_summerize(content)
|
|
logger.info(f"Summary: {summary}")
|
|
metadata['document_id'] = document_id
|
|
summary_doc_with_metadata = Document(
|
|
page_content=summary, metadata=metadata)
|
|
sids = commons['summaries_vector_store'].add_documents(
|
|
[summary_doc_with_metadata])
|
|
if sids and len(sids) > 0:
|
|
commons['supabase'].table("summaries").update(
|
|
{"document_id": document_id}).match({"id": sids[0]}).execute()
|
|
|
|
|
|
def get_answer(commons: CommonsDep, chat_message: ChatMessage, email: str, user_openai_api_key: str):
|
|
Brain = BrainPicking().init(chat_message.model, email)
|
|
qa = Brain.get_qa(chat_message, user_openai_api_key)
|
|
|
|
neurons = Neurons(commons=commons)
|
|
|
|
if chat_message.use_summarization:
|
|
summaries = neurons.similarity_search(chat_message.question, table='match_summaries')
|
|
evaluations = llm_evaluate_summaries(
|
|
chat_message.question, summaries, chat_message.model)
|
|
if evaluations:
|
|
response = commons['supabase'].from_('vectors').select(
|
|
'*').in_('id', values=[e['document_id'] for e in evaluations]).execute()
|
|
additional_context = '---\nAdditional Context={}'.format(
|
|
'---\n'.join(data['content'] for data in response.data)
|
|
) + '\n'
|
|
model_response = qa(
|
|
{"question": additional_context + chat_message.question})
|
|
else:
|
|
transformed_history = []
|
|
|
|
for i in range(0, len(chat_message.history) - 1, 2):
|
|
user_message = chat_message.history[i][1]
|
|
assistant_message = chat_message.history[i + 1][1]
|
|
transformed_history.append((user_message, assistant_message))
|
|
model_response = qa({"question": chat_message.question, "chat_history": transformed_history})
|
|
|
|
answer = model_response['answer']
|
|
|
|
if "source_documents" in answer:
|
|
sources = [
|
|
doc.metadata["file_name"] for doc in answer["source_documents"]
|
|
if "file_name" in doc.metadata]
|
|
if sources:
|
|
files = dict.fromkeys(sources)
|
|
answer = answer + "\n\nRef: " + "; ".join(files)
|
|
|
|
return answer |