from datetime import datetime from db.init import get_db_connection, DB_TIMEZONE def save_conversation(record, question, rewritten_question): timestamp = datetime.now(DB_TIMEZONE) conn = get_db_connection() try: with conn.cursor() as cur: cur.execute( """ INSERT INTO conversations ( question, rewritten_question, answer, model, instructions, prompt, prompt_tokens, completion_tokens, total_tokens, response_time, cost, timestamp ) VALUES ( %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) RETURNING id """, ( question, rewritten_question, record.answer, record.model, record.instructions, record.prompt, record.prompt_tokens, record.completion_tokens, record.total_tokens, record.response_time, record.cost, timestamp, ), ) conversation_id = cur.fetchone()[0] conn.commit() finally: conn.close() return conversation_id