Basic Chat API (#38)

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
 
 

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

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