Files
DM_rewrite_3.31/booway_kg_api/export_like_msg.py
T
2025-03-31 15:17:47 +08:00

89 lines
2.4 KiB
Python

# coding:utf-8
# @Time : 2024/8/27 下午3:08
# @Author : ouyangyouzhang
# @FileName : export_like_msg.py
# @Describe : 导出点赞的消息
import psycopg2
from psycopg2 import sql
import os
import pandas as pd
class ExportMsg:
def __init__(self):
try:
# 连接数据库
self.connection = psycopg2.connect(
user="postgres",
password="difyai123456",
host="172.20.0.145",
port=5432,
database="dify"
)
except (Exception, psycopg2.Error) as error:
print("Error while connecting to PostgreSQL", error)
def export_like_msg_id(self, user_id: str) -> list:
cursor = self.connection.cursor()
# 构建SQL查询
query = sql.SQL("""
SELECT message_id, from_end_user_id, rating
FROM message_feedbacks
WHERE from_end_user_id = %s AND rating = %s
""")
# 执行查询并获取结果
cursor.execute(query, (user_id, "like"))
records = cursor.fetchall()
if cursor:
cursor.close()
return [item[0] for item in records]
def export_QA_to_excel(self, msg_id_list, save_file) -> bool:
q_list = []
a_list = []
cursor = self.connection.cursor()
for item in msg_id_list:
# 构建SQL查询
query = sql.SQL("""
SELECT id, query, answer
FROM messages
WHERE id = %s
""")
# 执行查询并获取结果
cursor.execute(query, (item,))
records = cursor.fetchall()
q_list.append(records[0][1])
a_list.append(records[0][2])
if cursor:
cursor.close()
data = {"query": q_list, "answer": a_list}
df = pd.DataFrame(data)
df.to_excel(save_file, index=False)
def main(arg1: list) -> dict:
if arg1 is None or len(arg1) == 0:
return {
"result": [],
}
new_list = [{"title": item["title"], "content": item["content"]} for item in arg1]
return {
"result": new_list,
}
if __name__ == '__main__':
print(e)
ex = ExportMsg()
msg_list = ex.export_like_msg_id("de014752-5509-4c9e-a26d-72d7d955f895")
ex.export_QA_to_excel(msg_list, "D:\\Code\\PGSql\\like_msg.xlsx")
print('Python')