leo-bourrel commited on
Commit
b8c8744
·
1 Parent(s): 7b594ac

feat: Add source documents on right of display

Browse files
Files changed (2) hide show
  1. app.py +10 -2
  2. message.py +4 -2
app.py CHANGED
@@ -19,6 +19,7 @@ st.title("Sorbobot - Le futur de la recherche scientifique interactive")
19
 
20
  chat_column, doc_column = st.columns([2, 1])
21
 
 
22
  def initialize_session_state():
23
  if "history" not in st.session_state:
24
  st.session_state.history = []
@@ -53,16 +54,20 @@ def initialize_session_state():
53
  return_source_documents=True,
54
  )
55
 
 
56
  def on_click_callback():
57
  with get_openai_callback() as cb:
58
  human_prompt = st.session_state.human_prompt
59
  llm_response = st.session_state.conversation(human_prompt)
60
  st.session_state.history.append(Message("human", human_prompt))
61
  st.session_state.history.append(
62
- Message("ai", llm_response["answer"], llm_response["source_documents"])
 
 
63
  )
64
  st.session_state.token_count += cb.total_tokens
65
 
 
66
  load_css()
67
  initialize_session_state()
68
 
@@ -142,4 +147,7 @@ with chat_column:
142
  with doc_column:
143
  if len(st.session_state.history) > 0:
144
  st.markdown("**Source document**")
145
- st.markdown(st.session_state.history[-1].metadata)
 
 
 
 
19
 
20
  chat_column, doc_column = st.columns([2, 1])
21
 
22
+
23
  def initialize_session_state():
24
  if "history" not in st.session_state:
25
  st.session_state.history = []
 
54
  return_source_documents=True,
55
  )
56
 
57
+
58
  def on_click_callback():
59
  with get_openai_callback() as cb:
60
  human_prompt = st.session_state.human_prompt
61
  llm_response = st.session_state.conversation(human_prompt)
62
  st.session_state.history.append(Message("human", human_prompt))
63
  st.session_state.history.append(
64
+ Message(
65
+ "ai", llm_response["answer"], documents=llm_response["source_documents"]
66
+ )
67
  )
68
  st.session_state.token_count += cb.total_tokens
69
 
70
+
71
  load_css()
72
  initialize_session_state()
73
 
 
147
  with doc_column:
148
  if len(st.session_state.history) > 0:
149
  st.markdown("**Source document**")
150
+ for doc in st.session_state.history[-1].documents:
151
+ expander = st.expander(doc.metadata["title"])
152
+ expander.markdown("**" + doc.metadata["doi"] + "**")
153
+ expander.markdown(doc.page_content)
message.py CHANGED
@@ -1,5 +1,7 @@
1
  from dataclasses import dataclass
2
- from typing import Literal, Optional
 
 
3
 
4
 
5
  @dataclass
@@ -8,4 +10,4 @@ class Message:
8
 
9
  origin: Literal["human", "ai"]
10
  message: str
11
- metadata: Optional[dict] = None
 
1
  from dataclasses import dataclass
2
+ from typing import Literal, Optional, List
3
+
4
+ from langchain.schema import Document
5
 
6
 
7
  @dataclass
 
10
 
11
  origin: Literal["human", "ai"]
12
  message: str
13
+ documents: Optional[List[Document]] = None