How do you return source_documents from a ConversationalRetrievalChain?

[auxon]

Here's my code:

from enum import StrEnum
import os
from fastapi import FastAPI
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from dotenv import load_dotenv
from lanarky import LangchainRouter, StreamingResponse
from langchain.chains import LLMChain, ConversationalRetrievalChain
from langchain.embeddings import GPT4AllEmbeddings
from langchain.vectorstores import Chroma
from langchain.document_loaders import DirectoryLoader
from langchain.document_loaders import UnstructuredHTMLLoader
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.callbacks.streaming_aiter import AsyncIteratorCallbackHandler
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.memory import ConversationSummaryBufferMemory
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts.chat import (
    ChatPromptTemplate,
    PromptTemplate)
from langchain.memory.chat_message_histories.in_memory import (
    ChatMessageHistory)
from pydantic import BaseModel, constr
from AGPT4All import AGPT4All

load_dotenv()  # load environment variables from .env file
model_path = os.environ.get("MODEL_PATH")
titleDownloads = os.environ.get("TITLE_DOWNLOADS")
print(f"TITLE_DOWNLOADS={titleDownloads}")
print(f"MODEL_PATH={model_path}")

app = FastAPI(title="AI", version="0.0.1",
              description="AI Service")

app.mount("/static", StaticFiles(directory="static"), name="static")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)

callbacks = [AsyncIteratorCallbackHandler(), StreamingStdOutCallbackHandler()]

llm = AGPT4All(model=model_path, callbacks=callbacks, verbose=True,
               max_tokens=4096, n_predict=4096, streaming=True)

embeddings = GPT4AllEmbeddings()

if os.path.exists("./chroma_db") and os.path.isdir("./chroma_db"):
    vectorstore = Chroma(persist_directory="./chroma_db",
                         embedding_function=embeddings)
else:
    loader = DirectoryLoader(
        f'{titleDownloads}',
        loader_cls=UnstructuredHTMLLoader,
        recursive=True, glob="**/*.html", show_progress=True)
    documents = loader.load()
    print(f"Total docs: {len(documents)}")
    chunks = []
    splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=32)
    for chunk in splitter.split_documents(documents):
        chunks.append(chunk)
    texts = [doc.page_content for doc in chunks]
    metadatas = [doc.metadata for doc in chunks]
    vectorstore = Chroma.from_texts(texts=texts,
                                    embedding=embeddings,
                                    metadatas=metadatas,
                                    persist_directory="./chroma_db")

retriever = vectorstore.as_retriever()

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"


class Role(StrEnum):
    SYSTEM = "system"
    ASSISTANT = "assistant"
    USER = "user"


class Message(BaseModel):
    role: constr(
        regex=f"^({Role.ASSISTANT}|{Role.USER}|{Role.SYSTEM})$")  # NOQA
    content: str


class ChatRequest(BaseModel):
    model: str
    messages: list[Message]
    max_tokens: int
    temperature: float


def create_chain(messages: list[Message]):

    condense_template_str = (
        "You are an expert at summarizing chat histories and questions. "
        "Given the following Chat History and a Follow Up Question, "
        "rephrase the follow up question to be a new Standalone Question. \n "
        "Chat History: \n"
        "{chat_history} \n"
        "Follow Up Question: {question} \n"
        "Standalone question:")
    condense_template = PromptTemplate.from_template(
        f"{B_INST} {B_SYS}{condense_template_str.strip()}{E_SYS} {E_INST}")

    prompt_template_str = (
        "You are a helpful AI. "
        "Use the following Context and Chat History to answer the "
        "question at the end with a helpful and detailed answer. "
        "If you don't know the answer, just say "
        "'I don't know'; don't try to make up an answer. \n"
        "Context: {context} \n"
        "Chat History: {chat_history} \n"
        "Question: {question} \n"
        "Helpful Answer:")
    qa_prompt_template = ChatPromptTemplate.from_template(
        f"{B_INST} {B_SYS}{prompt_template_str.strip()}{E_SYS} {E_INST}")

    chat_memory = ChatMessageHistory()
    for message in messages:
        if message.role == Role.USER:
            chat_memory.add_user_message(message.content)
        elif message.role == Role.ASSISTANT:
            chat_memory.add_ai_message(message.content)

    memory = ConversationSummaryBufferMemory(
        llm=llm,
        chat_memory=chat_memory,
        memory_key="chat_history",
        input_key="question",
        return_messages=True)

    question_generator = LLMChain(llm=llm, prompt=condense_template,
                                  memory=memory, verbose=True)

    doc_chain = load_qa_chain(llm=llm, chain_type="stuff",
                              prompt=qa_prompt_template,  verbose=True)

    return ConversationalRetrievalChain(
            combine_docs_chain=doc_chain,
            memory=memory,
            retriever=retriever,
            question_generator=question_generator,
            return_generated_question=True,
            return_source_documents=True,
            output_key="answer",
            verbose=True)


router = LangchainRouter(
    streaming_mode=1,
    # llm_cache_mode=3,  # GPTCache
)


@router.post(
    "/chat",
    summary="AI Chat",
    description="Chat with AI Service",
)
def chat(request: ChatRequest):
    chain = create_chain(
        # model=request.model,
        messages=request.messages[:-1],
        # max_tokens=request.max_tokens,
        # temperature=request.temperature
        )
    return StreamingResponse.from_chain(
        chain, request.messages[-1].content, as_json=False)


app.include_router(router, tags=["chat"])

Everything is working except I want the source documents and metadata returned. If I set as_json=True in StreamingResponse.from_chain it returns a ton of "token"s but it also includes the sources and metadata, but it's an unexpected response format.

[auxon]

The JSON response looks like this (maybe it's correct except the error at the end:

{"token":" "}{"token":" I"}{"token":"'"}{"token":"m"}{"token":" happy"}{"token":" to"}{"token":" help"}{"token":" you"}{"token":" with"}{"token":" your"}{"token":" question"}{"token":"!"}{"token":" However"}{"token":","}{"token":" I"}{"token":" must"}{"token":" inform"}{"token":" you"}{"token":" that"}{"token":" M"}{"token":"oby"}{"token":" Dick"}{"token":" does"}{"token":" not"}{"token":" actually"}{"token":" kill"}{"token":" anyone"}{"token":" in"}{"token":" the"}{"token":" novel"}{"token":"."}{"token":" While"}{"token":" A"}{"token":"hab"}{"token":" is"}{"token":" obs"}{"token":"essed"}{"token":" with"}{"token":" hunting"}{"token":" and"}{"token":" killing"}{"token":" the"}{"token":" white"}{"token":" wh"}{"token":"ale"}{"token":","}{"token":" no"}{"token":" one"}{"token":" else"}{"token":" on"}{"token":" the"}{"token":" crew"}{"token":" is"}{"token":" har"}{"token":"med"}{"token":" or"}{"token":" killed"}{"token":" by"}{"token":" M"}{"token":"oby"}{"token":" Dick"}{"token":" directly"}{"token":"."}{"token":" In"}{"token":" fact"}{"token":","}{"token":" many"}{"token":" of"}{"token":" the"}{"token":" characters"}{"token":" on"}{"token":" the"}{"token":" ship"}{"token":" surv"}{"token":"ive"}{"token":" their"}{"token":" encounter"}{"token":" with"}{"token":" the"}{"token":" wh"}{"token":"ale"}{"token":" and"}{"token":" even"}{"token":" return"}{"token":" home"}{"token":" safely"}{"token":" after"}{"token":" their"}{"token":" journey"}{"token":"."}{"token":"\n"}{"token":"So"}{"token":" to"}{"token":" answer"}{"token":" your"}{"token":" question"}{"token":" directly"}{"token":","}{"token":" there"}{"token":" is"}{"token":" no"}{"token":" chapter"}{"token":" in"}{"token":" which"}{"token":" M"}{"token":"oby"}{"token":" Dick"}{"token":" k"}{"token":"ills"}{"token":" everyone"}{"token":"."}{"token":" The"}{"token":" novel"}{"token":"'"}{"token":"s"}{"token":" focus"}{"token":" is"}{"token":" more"}{"token":" on"}{"token":" A"}{"token":"hab"}{"token":"'"}{"token":"s"}{"token":" ob"}{"token":"session"}{"token":" and"}{"token":" his"}{"token":" un"}{"token":"wa"}{"token":"ver"}{"token":"ing"}{"token":" determ"}{"token":"ination"}{"token":" to"}{"token":" h"}{"token":"unt"}{"token":" down"}{"token":" and"}{"token":" kill"}{"token":" the"}{"token":" white"}{"token":" wh"}{"token":"ale"}{"token":","}{"token":" rather"}{"token":" than"}{"token":" any"}{"token":" actual"}{"token":" violence"}{"token":" or"}{"token":" harm"}{"token":" inf"}{"token":"lic"}{"token":"ted"}{"token":" by"}{"token":" the"}{"token":" wh"}{"token":"ale"}{"token":" itself"}{"token":"."}{"token":"\n"}{"token":"I"}{"token":" hope"}{"token":" this"}{"token":" helps"}{"token":" clarify"}{"token":" things"}{"token":" for"}{"token":" you"}{"token":"!"}{"token":" Let"}{"token":" me"}{"token":" know"}{"token":" if"}{"token":" you"}{"token":" have"}{"token":" any"}{"token":" other"}{"token":" questions"}{"token":"."}{"token":"","source_documents":[{"page_content":"Chapter 41. Moby Dick.\n\nI, Ishmael, was one of that crew; my shouts had gone up with the rest; my oath had been welded with theirs; and stronger I shouted, and more did I hammer and clinch my oath, because of the dread in my soul. A wild, mystical, sympathetical feeling was in me; Ahab’s quenchless feud seemed mine. With greedy ears I learned the history of that murderous monster against whom I and all the others had taken our oaths of violence and revenge.","metadata":{"source":"/Users/rah/titleDownloads/2c169b4c-8e8d-4dd4-8b57-82d5a21e4a1b/fae7a981-6f18-423d-b8e1-a7b0c73101fc.html"}},{"page_content":"Chapter 41. Moby Dick.\n\nI, Ishmael, was one of that crew; my shouts had gone up with the rest; my oath had been welded with theirs; and stronger I shouted, and more did I hammer and clinch my oath, because of the dread in my soul. A wild, mystical, sympathetical feeling was in me; Ahab’s quenchless feud seemed mine. With greedy ears I learned the history of that murderous monster against whom I and all the others had taken our oaths of violence and revenge.","metadata":{"source":"/Users/rah/titleDownloads/2c169b4c-8e8d-4dd4-8b57-82d5a21e4a1b/47.html"}},{"page_content":"it. All that most maddens and torments; all that stirs up the lees of things; all truth with malice in it; all that cracks the sinews and cakes the brain; all the subtle demonisms of life and thought; all evil, to crazy Ahab, were visibly personified, and made practically assailable in Moby Dick. He piled upon the whale’s white hump the sum of all the general rage and hate felt by his whole race from Adam down; and then, as if his chest had been a mortar, he burst his hot heart’s shell upon it.","metadata":{"source":"/Users/rah/titleDownloads/2c169b4c-8e8d-4dd4-8b57-82d5a21e4a1b/fae7a981-6f18-423d-b8e1-a7b0c73101fc.html"}},{"page_content":"it. All that most maddens and torments; all that stirs up the lees of things; all truth with malice in it; all that cracks the sinews and cakes the brain; all the subtle demonisms of life and thought; all evil, to crazy Ahab, were visibly personified, and made practically assailable in Moby Dick. He piled upon the whale’s white hump the sum of all the general rage and hate felt by his whole race from Adam down; and then, as if his chest had been a mortar, he burst his hot heart’s shell upon it.","metadata":{"source":"/Users/rah/titleDownloads/2c169b4c-8e8d-4dd4-8b57-82d5a21e4a1b/47.html"}}]}One output key expected, got dict_keys(['answer', 'source_documents', 'generated_question'])

So, maybe it's supposed to do this but I get "Unrecognized response type; displaying content as text.", using Swagger to test, and there is the "One output key expected, got dict_keys(['answer', 'source_documents', 'generated_question'])" at the end of the response that concerns me.

[auxon]

I fixed the "One output key expected" error by simply adding the output_key parameter to the ConversationSummaryBufferMemory:

    memory = ConversationSummaryBufferMemory(
        llm=llm,
        chat_memory=chat_memory,
        memory_key="chat_history",
        input_key="question",
        output_key="answer",
        return_messages=True)

Now, how do I get valid JSON back? Or is this the way JSON is streamed in response, but Swaggar just can't handle it?

[auxon]

I guess I forgot to change the StreamingMode to JSON!

router = LangchainRouter(
streaming_mode=2, # 1 = text, 2 = JSON
# llm_cache_mode=3, # GPTCache
)

[auxon]

Well, that didn't change the response at all. Still getting "Unrecognized response type; displaying content as text."

[auxon]

Ok, now I have finally resolved it, by adding the correct media_type to the StreamingResponse:

return StreamingResponse.from_chain(
    chain, request.messages[-1].content, as_json=True,
    media_type="text/event-stream")

[auxon]

I also had to add output_key="answer" to the question_generator LLMChain to get the follow up questions to work:

question_generator = LLMChain(llm=llm, prompt=condense_template,
                                  memory=memory, verbose=True,
                                  output_key="answer")