mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-08-17 02:10:31 +03:00
feat: self-reflect brain (#2610)
This pull request adds the SelfBrain integration to the list of available brain integrations.
This commit is contained in:
Stan Girard
GitHub
parent
287cbb3347
commit
c287fceb5b
487
backend/modules/brain/integrations/Self/Brain.py
Normal file
487
backend/modules/brain/integrations/Self/Brain.py
Normal file
@ -0,0 +1,487 @@
|
||||
import json
|
||||
from typing import AsyncIterable, List
|
||||
from uuid import UUID
|
||||
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
from langchain_core.prompts import (
|
||||
ChatPromptTemplate,
|
||||
MessagesPlaceholder,
|
||||
PromptTemplate,
|
||||
)
|
||||
from langchain_core.pydantic_v1 import BaseModel as BaseModelV1
|
||||
from langchain_core.pydantic_v1 import Field as FieldV1
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langgraph.graph import END, StateGraph
|
||||
from logger import get_logger
|
||||
from modules.brain.knowledge_brain_qa import KnowledgeBrainQA
|
||||
from modules.chat.dto.chats import ChatQuestion
|
||||
from modules.chat.dto.outputs import GetChatHistoryOutput
|
||||
from modules.chat.service.chat_service import ChatService
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
|
||||
# Post-processing
|
||||
def format_docs(docs):
|
||||
return "\n\n".join(doc.page_content for doc in docs)
|
||||
|
||||
|
||||
class GraphState(TypedDict):
|
||||
"""
|
||||
Represents the state of our graph.
|
||||
|
||||
Attributes:
|
||||
question: question
|
||||
generation: LLM generation
|
||||
documents: list of documents
|
||||
"""
|
||||
|
||||
question: str
|
||||
generation: str
|
||||
documents: List[str]
|
||||
|
||||
|
||||
# Data model
|
||||
class GradeDocuments(BaseModelV1):
|
||||
"""Binary score for relevance check on retrieved documents."""
|
||||
|
||||
binary_score: str = FieldV1(
|
||||
description="Documents are relevant to the question, 'yes' or 'no'"
|
||||
)
|
||||
|
||||
|
||||
class GradeHallucinations(BaseModelV1):
|
||||
"""Binary score for hallucination present in generation answer."""
|
||||
|
||||
binary_score: str = FieldV1(
|
||||
description="Answer is grounded in the facts, 'yes' or 'no'"
|
||||
)
|
||||
|
||||
|
||||
# Data model
|
||||
class GradeAnswer(BaseModelV1):
|
||||
"""Binary score to assess answer addresses question."""
|
||||
|
||||
binary_score: str = FieldV1(
|
||||
description="Answer addresses the question, 'yes' or 'no'"
|
||||
)
|
||||
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
chat_service = ChatService()
|
||||
|
||||
|
||||
class SelfBrain(KnowledgeBrainQA):
|
||||
"""
|
||||
GPT4Brain integrates with GPT-4 to provide real-time answers and supports various tools to enhance its capabilities.
|
||||
|
||||
Available Tools:
|
||||
- WebSearchTool: Performs web searches to find relevant information.
|
||||
- ImageGeneratorTool: Generates images based on textual descriptions.
|
||||
- URLReaderTool: Reads and summarizes content from URLs.
|
||||
- EmailSenderTool: Sends emails with specified content.
|
||||
|
||||
Use Cases:
|
||||
- WebSearchTool can be used to find the latest news articles on a specific topic or to gather information from various websites.
|
||||
- ImageGeneratorTool is useful for creating visual content based on textual prompts, such as generating a company logo based on a description.
|
||||
- URLReaderTool can be used to summarize articles or web pages, making it easier to quickly understand the content without reading the entire text.
|
||||
- EmailSenderTool enables automated email sending, such as sending a summary of a meeting's minutes to all participants.
|
||||
"""
|
||||
|
||||
max_input: int = 10000
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def calculate_pricing(self):
|
||||
return 3
|
||||
|
||||
def retrieval_grade(self):
|
||||
llm = ChatOpenAI(model="gpt-4o", temperature=0)
|
||||
structured_llm_grader = llm.with_structured_output(GradeDocuments)
|
||||
|
||||
# Prompt
|
||||
system = """You are a grader assessing relevance of a retrieved document to a user question. \n
|
||||
It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n
|
||||
If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
|
||||
Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""
|
||||
grade_prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", system),
|
||||
(
|
||||
"human",
|
||||
"Retrieved document: \n\n {document} \n\n User question: {question}",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
retrieval_grader = grade_prompt | structured_llm_grader
|
||||
|
||||
return retrieval_grader
|
||||
|
||||
def generation_rag(self):
|
||||
# Prompt
|
||||
human_prompt = """You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
|
||||
|
||||
Question: {question}
|
||||
|
||||
Context: {context}
|
||||
|
||||
Answer:
|
||||
"""
|
||||
prompt_human = PromptTemplate.from_template(human_prompt)
|
||||
# LLM
|
||||
llm = ChatOpenAI(model="gpt-4o", temperature=0)
|
||||
|
||||
# Chain
|
||||
rag_chain = prompt_human | llm | StrOutputParser()
|
||||
|
||||
return rag_chain
|
||||
|
||||
def hallucination_grader(self):
|
||||
# LLM with function call
|
||||
llm = ChatOpenAI(model="gpt-4o", temperature=0)
|
||||
structured_llm_grader = llm.with_structured_output(GradeHallucinations)
|
||||
|
||||
# Prompt
|
||||
system = """You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \n
|
||||
Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts."""
|
||||
hallucination_prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", system),
|
||||
(
|
||||
"human",
|
||||
"Set of facts: \n\n {documents} \n\n LLM generation: {generation}",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
hallucination_grader = hallucination_prompt | structured_llm_grader
|
||||
|
||||
return hallucination_grader
|
||||
|
||||
def answer_grader(self):
|
||||
# LLM with function call
|
||||
llm = ChatOpenAI(model="gpt-4o", temperature=0)
|
||||
structured_llm_grader = llm.with_structured_output(GradeAnswer)
|
||||
|
||||
# Prompt
|
||||
system = """You are a grader assessing whether an answer addresses / resolves a question \n
|
||||
Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question."""
|
||||
answer_prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", system),
|
||||
(
|
||||
"human",
|
||||
"User question: \n\n {question} \n\n LLM generation: {generation}",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
answer_grader = answer_prompt | structured_llm_grader
|
||||
|
||||
return answer_grader
|
||||
|
||||
def question_rewriter(self):
|
||||
# LLM
|
||||
llm = ChatOpenAI(model="gpt-4o", temperature=0)
|
||||
|
||||
# Prompt
|
||||
system = """You a question re-writer that converts an input question to a better version that is optimized \n
|
||||
for vectorstore retrieval. Look at the input and try to reason about the underlying semantic intent / meaning."""
|
||||
re_write_prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", system),
|
||||
(
|
||||
"human",
|
||||
"Here is the initial question: \n\n {question} \n Formulate an improved question.",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
question_rewriter = re_write_prompt | llm | StrOutputParser()
|
||||
|
||||
return question_rewriter
|
||||
|
||||
def get_chain(self):
|
||||
|
||||
graph = self.create_graph()
|
||||
|
||||
return graph
|
||||
|
||||
def create_graph(self):
|
||||
|
||||
workflow = StateGraph(GraphState)
|
||||
|
||||
# Define the nodes
|
||||
workflow.add_node("retrieve", self.retrieve) # retrieve
|
||||
workflow.add_node("grade_documents", self.grade_documents) # grade documents
|
||||
workflow.add_node("generate", self.generate) # generatae
|
||||
workflow.add_node("transform_query", self.transform_query) # transform_query
|
||||
|
||||
# Build graph
|
||||
workflow.set_entry_point("retrieve")
|
||||
workflow.add_edge("retrieve", "grade_documents")
|
||||
workflow.add_conditional_edges(
|
||||
"grade_documents",
|
||||
self.decide_to_generate,
|
||||
{
|
||||
"transform_query": "transform_query",
|
||||
"generate": "generate",
|
||||
},
|
||||
)
|
||||
workflow.add_edge("transform_query", "retrieve")
|
||||
workflow.add_conditional_edges(
|
||||
"generate",
|
||||
self.grade_generation_v_documents_and_question,
|
||||
{
|
||||
"not supported": "generate",
|
||||
"useful": END,
|
||||
"not useful": "transform_query",
|
||||
},
|
||||
)
|
||||
|
||||
# Compile
|
||||
app = workflow.compile()
|
||||
return app
|
||||
|
||||
def retrieve(self, state):
|
||||
"""
|
||||
Retrieve documents
|
||||
|
||||
Args:
|
||||
state (dict): The current graph state
|
||||
|
||||
Returns:
|
||||
state (dict): New key added to state, documents, that contains retrieved documents
|
||||
"""
|
||||
print("---RETRIEVE---")
|
||||
logger.info("Retrieving documents")
|
||||
question = state["question"]
|
||||
logger.info(f"Question: {question}")
|
||||
|
||||
# Retrieval
|
||||
retriever = self.knowledge_qa.get_retriever()
|
||||
documents = retriever.get_relevant_documents(question)
|
||||
return {"documents": documents, "question": question}
|
||||
|
||||
def generate(self, state):
|
||||
"""
|
||||
Generate answer
|
||||
|
||||
Args:
|
||||
state (dict): The current graph state
|
||||
|
||||
Returns:
|
||||
state (dict): New key added to state, generation, that contains LLM generation
|
||||
"""
|
||||
print("---GENERATE---")
|
||||
question = state["question"]
|
||||
documents = state["documents"]
|
||||
|
||||
formatted_docs = format_docs(documents)
|
||||
# RAG generation
|
||||
generation = self.generation_rag().invoke(
|
||||
{"context": formatted_docs, "question": question}
|
||||
)
|
||||
return {"documents": documents, "question": question, "generation": generation}
|
||||
|
||||
def grade_documents(self, state):
|
||||
"""
|
||||
Determines whether the retrieved documents are relevant to the question.
|
||||
|
||||
Args:
|
||||
state (dict): The current graph state
|
||||
|
||||
Returns:
|
||||
state (dict): Updates documents key with only filtered relevant documents
|
||||
"""
|
||||
|
||||
print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
|
||||
question = state["question"]
|
||||
documents = state["documents"]
|
||||
|
||||
# Score each doc
|
||||
filtered_docs = []
|
||||
for d in documents:
|
||||
score = self.retrieval_grade().invoke(
|
||||
{"question": question, "document": d.page_content}
|
||||
)
|
||||
grade = score.binary_score
|
||||
if grade == "yes":
|
||||
print("---GRADE: DOCUMENT RELEVANT---")
|
||||
filtered_docs.append(d)
|
||||
else:
|
||||
print("---GRADE: DOCUMENT NOT RELEVANT---")
|
||||
continue
|
||||
return {"documents": filtered_docs, "question": question}
|
||||
|
||||
def transform_query(self, state):
|
||||
"""
|
||||
Transform the query to produce a better question.
|
||||
|
||||
Args:
|
||||
state (dict): The current graph state
|
||||
|
||||
Returns:
|
||||
state (dict): Updates question key with a re-phrased question
|
||||
"""
|
||||
|
||||
print("---TRANSFORM QUERY---")
|
||||
question = state["question"]
|
||||
documents = state["documents"]
|
||||
|
||||
# Re-write question
|
||||
better_question = self.question_rewriter().invoke({"question": question})
|
||||
return {"documents": documents, "question": better_question}
|
||||
|
||||
def decide_to_generate(self, state):
|
||||
"""
|
||||
Determines whether to generate an answer, or re-generate a question.
|
||||
|
||||
Args:
|
||||
state (dict): The current graph state
|
||||
|
||||
Returns:
|
||||
str: Binary decision for next node to call
|
||||
"""
|
||||
|
||||
print("---ASSESS GRADED DOCUMENTS---")
|
||||
question = state["question"]
|
||||
filtered_documents = state["documents"]
|
||||
|
||||
if not filtered_documents:
|
||||
# All documents have been filtered check_relevance
|
||||
# We will re-generate a new query
|
||||
print(
|
||||
"---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---"
|
||||
)
|
||||
return "transform_query"
|
||||
else:
|
||||
# We have relevant documents, so generate answer
|
||||
print("---DECISION: GENERATE---")
|
||||
return "generate"
|
||||
|
||||
def grade_generation_v_documents_and_question(self, state):
|
||||
"""
|
||||
Determines whether the generation is grounded in the document and answers question.
|
||||
|
||||
Args:
|
||||
state (dict): The current graph state
|
||||
|
||||
Returns:
|
||||
str: Decision for next node to call
|
||||
"""
|
||||
|
||||
print("---CHECK HALLUCINATIONS---")
|
||||
question = state["question"]
|
||||
documents = state["documents"]
|
||||
generation = state["generation"]
|
||||
|
||||
score = self.hallucination_grader().invoke(
|
||||
{"documents": documents, "generation": generation}
|
||||
)
|
||||
grade = score.binary_score
|
||||
|
||||
# Check hallucination
|
||||
if grade == "yes":
|
||||
print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
|
||||
# Check question-answering
|
||||
print("---GRADE GENERATION vs QUESTION---")
|
||||
score = self.answer_grader().invoke(
|
||||
{"question": question, "generation": generation}
|
||||
)
|
||||
grade = score.binary_score
|
||||
if grade == "yes":
|
||||
print("---DECISION: GENERATION ADDRESSES QUESTION---")
|
||||
return "useful"
|
||||
else:
|
||||
print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
|
||||
return "not useful"
|
||||
else:
|
||||
print("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")
|
||||
return "not supported"
|
||||
|
||||
async def generate_stream(
|
||||
self, chat_id: UUID, question: ChatQuestion, save_answer: bool = True
|
||||
) -> AsyncIterable:
|
||||
conversational_qa_chain = self.get_chain()
|
||||
transformed_history, streamed_chat_history = (
|
||||
self.initialize_streamed_chat_history(chat_id, question)
|
||||
)
|
||||
filtered_history = self.filter_history(transformed_history, 40, 2000)
|
||||
response_tokens = []
|
||||
config = {"metadata": {"conversation_id": str(chat_id)}}
|
||||
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
(
|
||||
"system",
|
||||
"You are GPT-4 powered by Quivr. You are an assistant. {custom_personality}",
|
||||
),
|
||||
MessagesPlaceholder(variable_name="chat_history"),
|
||||
("human", "{question}"),
|
||||
]
|
||||
)
|
||||
prompt_formated = prompt.format_messages(
|
||||
chat_history=filtered_history,
|
||||
question=question.question,
|
||||
custom_personality=(
|
||||
self.prompt_to_use.content if self.prompt_to_use else None
|
||||
),
|
||||
)
|
||||
|
||||
async for event in conversational_qa_chain.astream(
|
||||
{"question": question.question}, config=config
|
||||
):
|
||||
for key, value in event.items():
|
||||
if "generation" in value and value["generation"] != "":
|
||||
response_tokens.append(value["generation"])
|
||||
streamed_chat_history.assistant = value["generation"]
|
||||
|
||||
yield f"data: {json.dumps(streamed_chat_history.dict())}"
|
||||
|
||||
self.save_answer(question, response_tokens, streamed_chat_history, save_answer)
|
||||
|
||||
def generate_answer(
|
||||
self, chat_id: UUID, question: ChatQuestion, save_answer: bool = True
|
||||
) -> GetChatHistoryOutput:
|
||||
conversational_qa_chain = self.get_chain()
|
||||
transformed_history, _ = self.initialize_streamed_chat_history(
|
||||
chat_id, question
|
||||
)
|
||||
filtered_history = self.filter_history(transformed_history, 40, 2000)
|
||||
config = {"metadata": {"conversation_id": str(chat_id)}}
|
||||
|
||||
prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
(
|
||||
"system",
|
||||
"You are GPT-4 powered by Quivr. You are an assistant. {custom_personality}",
|
||||
),
|
||||
MessagesPlaceholder(variable_name="chat_history"),
|
||||
("human", "{question}"),
|
||||
]
|
||||
)
|
||||
prompt_formated = prompt.format_messages(
|
||||
chat_history=filtered_history,
|
||||
question=question.question,
|
||||
custom_personality=(
|
||||
self.prompt_to_use.content if self.prompt_to_use else None
|
||||
),
|
||||
)
|
||||
model_response = conversational_qa_chain.invoke(
|
||||
{"messages": prompt_formated},
|
||||
config=config,
|
||||
)
|
||||
|
||||
answer = model_response["messages"][-1].content
|
||||
|
||||
return self.save_non_streaming_answer(
|
||||
chat_id=chat_id, question=question, answer=answer, metadata={}
|
||||
)
|
||||
0
backend/modules/brain/integrations/Self/__init__.py
Normal file
0
backend/modules/brain/integrations/Self/__init__.py
Normal file
@ -200,7 +200,7 @@ class KnowledgeBrainQA(BaseModel, QAInterface):
|
||||
max_input=self.max_input,
|
||||
max_tokens=self.max_tokens,
|
||||
**kwargs,
|
||||
)
|
||||
) # type: ignore
|
||||
|
||||
@property
|
||||
def prompt_to_use(self):
|
||||
|
||||
@ -5,6 +5,7 @@ from modules.brain.integrations.Big.Brain import BigBrain
|
||||
from modules.brain.integrations.GPT4.Brain import GPT4Brain
|
||||
from modules.brain.integrations.Notion.Brain import NotionBrain
|
||||
from modules.brain.integrations.Proxy.Brain import ProxyBrain
|
||||
from modules.brain.integrations.Self.Brain import SelfBrain
|
||||
from modules.brain.integrations.SQL.Brain import SQLBrain
|
||||
from modules.brain.knowledge_brain_qa import KnowledgeBrainQA
|
||||
from modules.brain.service.api_brain_definition_service import ApiBrainDefinitionService
|
||||
@ -45,6 +46,7 @@ integration_list = {
|
||||
"big": BigBrain,
|
||||
"doc": KnowledgeBrainQA,
|
||||
"proxy": ProxyBrain,
|
||||
"self": SelfBrain,
|
||||
}
|
||||
|
||||
brain_service = BrainService()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user