feat: add optional source documents to message

app.py

@@ -54,7 +54,9 @@ def on_click_callback():
         human_prompt = st.session_state.human_prompt
         llm_response = st.session_state.conversation(human_prompt)
         st.session_state.history.append(Message("human", human_prompt))
-        st.session_state.history.append(Message("ai", llm_response["answer"]))
+        st.session_state.history.append(
+            Message("ai", llm_response["answer"], llm_response["source_documents"])
+        )
         st.session_state.token_count += cb.total_tokens
 
 

message.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Literal
+from typing import Literal, Optional
 
 
 @dataclass
@@ -8,3 +8,4 @@ class Message:
 
     origin: Literal["human", "ai"]
     message: str
+    metadata: Optional[dict] = None