Zdao032
89 lines
2.4 KiB
Python
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')
|