mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-19 12:21:46 +03:00
0fd99fe5b2
This pull request updates the API documentation to include new sections on configuring Quivr and contacting the Quivr team. It also removes the "API Brains" section from the documentation.
72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
import json
|
|
from typing import AsyncIterable
|
|
from uuid import UUID
|
|
|
|
from langchain_community.chat_models import ChatLiteLLM
|
|
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
|
|
from modules.brain.knowledge_brain_qa import KnowledgeBrainQA
|
|
from modules.chat.dto.chats import ChatQuestion
|
|
|
|
|
|
class ClaudeBrain(KnowledgeBrainQA):
|
|
"""This is the Notion brain class. it is a KnowledgeBrainQA has the data is stored locally.
|
|
It is going to call the Data Store internally to get the data.
|
|
|
|
Args:
|
|
KnowledgeBrainQA (_type_): A brain that store the knowledge internaly
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
**kwargs,
|
|
):
|
|
super().__init__(
|
|
**kwargs,
|
|
)
|
|
|
|
def calculate_pricing(self):
|
|
return 3
|
|
|
|
def get_chain(self):
|
|
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
(
|
|
"system",
|
|
"You are Claude powered by Quivr. You are an assistant. {custom_personality}",
|
|
),
|
|
MessagesPlaceholder(variable_name="chat_history"),
|
|
("human", "{question}"),
|
|
]
|
|
)
|
|
|
|
chain = prompt | ChatLiteLLM(
|
|
model="claude-3-haiku-20240307", max_tokens=self.max_tokens
|
|
)
|
|
|
|
return chain
|
|
|
|
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)
|
|
)
|
|
response_tokens = []
|
|
|
|
async for chunk in conversational_qa_chain.astream(
|
|
{
|
|
"question": question.question,
|
|
"chat_history": transformed_history,
|
|
"custom_personality": (
|
|
self.prompt_to_use.content if self.prompt_to_use else None
|
|
),
|
|
}
|
|
):
|
|
response_tokens.append(chunk.content)
|
|
streamed_chat_history.assistant = chunk.content
|
|
yield f"data: {json.dumps(streamed_chat_history.dict())}"
|
|
|
|
self.save_answer(question, response_tokens, streamed_chat_history, save_answer)
|