3.31 上传 dm rewrite

booway_kg_api/export_like_msg.py

@@ -0,0 +1,88 @@
+# 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')