feat: add relation between chat_history and documents

app.py

@@ -10,7 +10,7 @@ from langchain.chains import ConversationalRetrievalChain
 from langchain.chains.conversation.memory import ConversationBufferMemory
 from langchain.embeddings import GPT4AllEmbeddings
 
-from chat_history import insert_chat_history
+from chat_history import insert_chat_history, insert_chat_history_articles
 from css import load_css
 from custom_pgvector import CustomPGVector
 from message import Message
@@ -80,6 +80,7 @@ def on_click_callback():
         )
         st.session_state.token_count += cb.total_tokens
         history_id = insert_chat_history(conn, human_prompt, llm_response["answer"])
+        insert_chat_history_articles(conn, history_id, llm_response["source_documents"])
 
 
 load_css()

chat_history.py

@@ -1,3 +1,6 @@
+import json
+from typing import List
+
 import sqlalchemy
 from sqlalchemy import text
 from sqlalchemy.orm import Session
@@ -22,3 +25,24 @@ def insert_chat_history(conn: sqlalchemy.engine.Connection, query: str, answer:
         last_row_id = result.fetchone()[0]
         conn.commit()
     return last_row_id
+
+
+def insert_chat_history_articles(
+    conn: sqlalchemy.engine.Connection, chat_history_id: int, articles: List[str]
+):
+    source_documents = [json.loads(article.page_content) for article in articles]
+
+    with Session(conn) as conn:
+        conn.execute(
+            text(
+                "INSERT INTO chat_history_articles (chat_history_id, article_id) VALUES (:chat_history_id, :article_id) ON CONFLICT DO NOTHING;"
+            ),
+            [
+                {
+                    "chat_history_id": chat_history_id,
+                    "article_id": article["id"],
+                }
+                for article in source_documents
+            ],
+        )
+        conn.commit()