sorbobotapp/chat_history.py

import json
from typing import List

import sqlalchemy
from sqlalchemy import text
from sqlalchemy.orm import Session


def insert_chat_history(conn: sqlalchemy.engine.Connection, query: str, answer: str):
    with Session(conn) as conn:
        conn.execute(
            text("INSERT INTO chat_history (query, answer) VALUES (:query, :answer);"),
            [
                {
                    "query": query,
                    "answer": answer,
                }
            ],
        )
        conn.commit()

        result = conn.execute(
            text("SELECT id FROM chat_history ORDER BY id DESC LIMIT 1;")
        )
        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]
):
    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.metadata["id"],
                }
                for article in articles
            ],
        )
        conn.commit()