From 9ae3a4af7f1bbe7ab9bf14257a4150d3eaa4353f Mon Sep 17 00:00:00 2001 From: Yuhong Sun Date: Sun, 18 Feb 2024 23:46:59 -0800 Subject: [PATCH] Basic Chat API (#38) --- backend/ee/danswer/main.py | 9 +- .../server/query_and_chat/chat_backend.py | 116 ++++++++++++++++++ .../danswer/server/query_and_chat/models.py | 40 ++++++ 3 files changed, 162 insertions(+), 3 deletions(-) create mode 100644 backend/ee/danswer/server/query_and_chat/chat_backend.py create mode 100644 backend/ee/danswer/server/query_and_chat/models.py diff --git a/backend/ee/danswer/main.py b/backend/ee/danswer/main.py index 173f3c261a45..dfeeb96a77e1 100644 --- a/backend/ee/danswer/main.py +++ b/backend/ee/danswer/main.py @@ -19,8 +19,11 @@ from danswer.utils.variable_functionality import global_version from ee.danswer.configs.app_configs import OPENID_CONFIG_URL from ee.danswer.server.analytics.api import router as analytics_router from ee.danswer.server.api_key.api import router as api_key_router +from ee.danswer.server.query_and_chat.chat_backend import ( + router as chat_router, +) from ee.danswer.server.query_and_chat.query_backend import ( - basic_router as chat_query_router, + basic_router as query_router, ) from ee.danswer.server.query_history.api import router as query_history_router from ee.danswer.server.saml import router as saml_router @@ -69,8 +72,8 @@ def get_ee_application() -> FastAPI: # Api key management include_router_with_global_prefix_prepended(application, api_key_router) # EE only backend APIs - include_router_with_global_prefix_prepended(application, chat_query_router) - + include_router_with_global_prefix_prepended(application, query_router) + include_router_with_global_prefix_prepended(application, chat_router) return application diff --git a/backend/ee/danswer/server/query_and_chat/chat_backend.py b/backend/ee/danswer/server/query_and_chat/chat_backend.py new file mode 100644 index 000000000000..9bf78b121049 --- /dev/null +++ b/backend/ee/danswer/server/query_and_chat/chat_backend.py @@ -0,0 +1,116 @@ +import re + +from fastapi import APIRouter +from fastapi import Depends +from fastapi import HTTPException +from sqlalchemy.orm import Session + +from danswer.auth.users import current_user +from danswer.chat.chat_utils import create_chat_chain +from danswer.chat.models import DanswerAnswerPiece +from danswer.chat.models import QADocsResponse +from danswer.chat.models import StreamingError +from danswer.chat.process_message import stream_chat_message_objects +from danswer.db.chat import get_or_create_root_message +from danswer.db.engine import get_session +from danswer.db.models import User +from danswer.search.models import OptionalSearchSetting +from danswer.search.models import RetrievalDetails +from danswer.server.query_and_chat.models import CreateChatMessageRequest +from danswer.utils.logger import setup_logger +from ee.danswer.server.query_and_chat.models import BasicCreateChatMessageRequest +from ee.danswer.server.query_and_chat.models import ChatBasicResponse +from ee.danswer.server.query_and_chat.models import SimpleDoc + +logger = setup_logger() + +router = APIRouter(prefix="/chat") + + +def translate_doc_response_to_simple_doc( + doc_response: QADocsResponse, +) -> list[SimpleDoc]: + return [ + SimpleDoc( + semantic_identifier=doc.semantic_identifier, + link=doc.link, + blurb=doc.blurb, + match_highlights=[ + highlight for highlight in doc.match_highlights if highlight + ], + source_type=doc.source_type, + ) + for doc in doc_response.top_documents + ] + + +def remove_answer_citations(answer: str) -> str: + pattern = r"\s*\[\[\d+\]\]\(http[s]?://[^\s]+\)" + + return re.sub(pattern, "", answer) + + +@router.post("/send-message-simple-api") +def handle_simplified_chat_message( + chat_message_req: BasicCreateChatMessageRequest, + user: User | None = Depends(current_user), + db_session: Session = Depends(get_session), +) -> ChatBasicResponse: + """This is a Non-Streaming version that only gives back a minimal set of information""" + logger.info(f"Received new chat message: {chat_message_req.message}") + + if not chat_message_req.message and chat_message_req.prompt_id is not None: + raise HTTPException(status_code=400, detail="Empty chat message is invalid") + + try: + parent_message, _ = create_chat_chain( + chat_session_id=chat_message_req.chat_session_id, db_session=db_session + ) + except RuntimeError: + parent_message = get_or_create_root_message( + chat_session_id=chat_message_req.chat_session_id, db_session=db_session + ) + + if ( + chat_message_req.retrieval_options is None + and chat_message_req.search_doc_ids is None + ): + retrieval_options: RetrievalDetails | None = RetrievalDetails( + run_search=OptionalSearchSetting.ALWAYS, + real_time=False, + ) + else: + retrieval_options = chat_message_req.retrieval_options + + full_chat_msg_info = CreateChatMessageRequest( + chat_session_id=chat_message_req.chat_session_id, + parent_message_id=parent_message.id, + message=chat_message_req.message, + prompt_id=chat_message_req.prompt_id, + search_doc_ids=chat_message_req.search_doc_ids, + retrieval_options=retrieval_options, + query_override=chat_message_req.query_override, + ) + + packets = stream_chat_message_objects( + new_msg_req=full_chat_msg_info, + user=user, + db_session=db_session, + ) + + response = ChatBasicResponse() + + answer = "" + for packet in packets: + if isinstance(packet, DanswerAnswerPiece) and packet.answer_piece: + answer += packet.answer_piece + elif isinstance(packet, QADocsResponse): + response.simple_search_docs = translate_doc_response_to_simple_doc(packet) + elif isinstance(packet, StreamingError): + response.error_msg = packet.error + + response.answer = answer + if answer: + response.answer_citationless = remove_answer_citations(answer) + + return response diff --git a/backend/ee/danswer/server/query_and_chat/models.py b/backend/ee/danswer/server/query_and_chat/models.py new file mode 100644 index 000000000000..71715b818239 --- /dev/null +++ b/backend/ee/danswer/server/query_and_chat/models.py @@ -0,0 +1,40 @@ +from pydantic import BaseModel + +from danswer.configs.constants import DocumentSource +from danswer.search.models import RetrievalDetails + + +class BasicCreateChatMessageRequest(BaseModel): + """Before creating messages, be sure to create a chat_session and get an id + Note, for simplicity this option only allows for a single linear chain of messages + """ + + chat_session_id: int + # New message contents + message: str + # Defaults to using retrieval with no additional filters + retrieval_options: RetrievalDetails | None = None + # Allows the caller to specify the exact search query they want to use + # will disable Query Rewording if specified + query_override: str | None = None + # If no prompt provided, provide canned retrieval answer, no actually LLM flow + # Use prompt_id 0 to use the system default prompt which is Answer-Question + prompt_id: int | None = 0 + # If search_doc_ids provided, then retrieval options are unused + search_doc_ids: list[int] | None = None + + +class SimpleDoc(BaseModel): + semantic_identifier: str + link: str | None + blurb: str + match_highlights: list[str] + source_type: DocumentSource + + +class ChatBasicResponse(BaseModel): + # This is built piece by piece, any of these can be None as the flow could break + answer: str | None = None + answer_citationless: str | None = None + simple_search_docs: list[SimpleDoc] | None = None + error_msg: str | None = None