Add lark_sheet_record_for_human_intervention

pqai_agent/toolkit/lark_sheet_record_for_human_intervention.py

@@ -0,0 +1,152 @@
+from datetime import datetime
+from typing import List
+
+import requests
+from pqai_agent.toolkit.base import BaseToolkit
+from pqai_agent.toolkit.function_tool import FunctionTool
+from pqai_agent.logging_service import logger
+
+class LarkSheetRecordForHumanIntervention(BaseToolkit):
+    r"""A toolkit for recording human intervention events into a Feishu spreadsheet."""
+    APP_ID = 'cli_a75d795d877d901c'
+    APP_SECRET = 'nDTqlBmAdKEyPThegK50ZbS4lKsCcYlN'
+    SPREADSHEET_TOKEN = 'PPqNsKJuwhIHLstNZanc6d4Xnrh'
+    SHEET_ID = '3SiNjr'
+
+    BASE_URL = "https://open.feishu.cn/open-apis"
+    TOKEN_URL = f"{BASE_URL}/auth/v3/tenant_access_token/internal"
+    SHEETS_URL_TEMPLATE = f"{BASE_URL}/sheets/v2/spreadsheets/{{spreadsheet_token}}/values_prepend"
+    CELL_STYLE_URL_TEMPLATE = f"{BASE_URL}/sheets/v2/spreadsheets/{{spreadsheet_token}}/styles_batch_update"
+
+    BASE_DATE = datetime(1899, 12, 30)  # Excel base date
+
+    def __init__(self, app_id: str = APP_ID, app_secret: str = APP_SECRET,
+                 spreadsheet_token: str = SPREADSHEET_TOKEN, sheet_id: str = SHEET_ID):
+        """
+        初始化类，设置飞书 API 凭据和表格信息。
+
+        Args:
+            app_id (str): 飞书应用的 App ID。
+            app_secret (str): 飞书应用的 App Secret。
+            spreadsheet_token (str): 飞书电子表格的 Token。
+            sheet_id (str): 飞书电子表格的 Sheet ID。
+        """
+        self.app_id = app_id
+        self.app_secret = app_secret
+        self.spreadsheet_token = spreadsheet_token
+        self.sheet_id = sheet_id
+        self.access_token = self._get_access_token()
+        super().__init__()
+
+    def _get_access_token(self) -> str:
+        """
+        获取飞书 API 的访问令牌。
+
+        Returns:
+            str: 访问令牌。
+        """
+        response = requests.post(self.TOKEN_URL, json={
+            "app_id": self.app_id,
+            "app_secret": self.app_secret
+        })
+        response_data = response.json()
+        if response.status_code == 200 and response_data.get("code") == 0:
+            return response_data["tenant_access_token"]
+        else:
+            raise Exception(f"Failed to get access token: {response_data}")
+
+    def send_lark_sheet_record_for_human_intervention(
+            self, employee: str, user: str, dialogue: str, reason: str
+    ) -> str:
+        """
+        Records a human intervention event into the Feishu spreadsheet.
+
+        Args:
+            employee (str): Name of the employee handling the intervention.
+            user (str): Name of the user involved.
+            dialogue (str): Content of the intervention dialogue.
+            reason (str): Reason for the intervention.
+
+        Returns:
+            str: A confirmation message.
+        """
+        # Prepare the data to be appended
+        now = datetime.now()
+        date_value = (now - self.BASE_DATE).days
+        time_value = round((now.hour * 3600 + now.minute * 60 + now.second) / 86400, 10)
+        values = [[date_value, time_value, employee, user, dialogue, reason]]
+
+        url = self.SHEETS_URL_TEMPLATE.format(spreadsheet_token=self.spreadsheet_token)
+        headers = {
+            "Authorization": f"Bearer {self.access_token}",
+            "Content-Type": "application/json"
+        }
+        payload = {
+            "valueRange": {
+                "range": f"{self.sheet_id}!A2:F2",
+                "values": values
+            }
+        }
+
+        try:
+            response = requests.post(url, headers=headers, json=payload)
+            response_data = response.json()
+
+            if response.status_code == 200 and response_data.get("code") == 0:
+                self._set_date_time_styles()
+                return "Record added successfully."
+            else:
+                return f"Failed to add record: {response_data.get('msg', 'Unknown error')}"
+        except Exception as e:
+            logger.error("Error occurred while adding record: %s", e)
+            return "Error occurred while adding the record."
+
+    def _set_date_time_styles(self):
+        # Define the style for date and time columns
+        style_payload = {
+            "data": [
+                {
+                    "ranges": [f"{self.sheet_id}!A2:A2"],
+                    "style": {
+                        "formatter": "yyyy/MM/dd"
+                    },
+                },
+                {
+                    "ranges": [f"{self.sheet_id}!B2:B2"],
+                    "style": {
+                        "formatter": "HH:mm:ss"
+                    },
+                }
+            ]
+        }
+
+        style_url = self.CELL_STYLE_URL_TEMPLATE.format(spreadsheet_token=self.spreadsheet_token)
+        headers = {
+            "Authorization": f"Bearer {self.access_token}",
+            "Content-Type": "application/json"
+        }
+        try:
+            response = requests.put(style_url, headers=headers, json=style_payload)
+            if response.status_code != 200:
+                logger.error("Failed to update cell styles: %s", response.json())
+        except Exception as e:
+            logger.error("Error occurred while updating cell styles: %s", e)
+
+    def get_tools(self) -> List[FunctionTool]:
+        return [FunctionTool(self.send_lark_sheet_record_for_human_intervention)]
+
+
+if __name__ == '__main__':
+    toolkit = LarkSheetRecordForHumanIntervention()
+    tools = toolkit.get_tools()
+    for tool in tools:
+        print(f"Tool schema: {tool.get_openai_tool_schema()}")
+
+    # Test the function
+    resp = toolkit.send_lark_sheet_record_for_human_intervention(
+        employee="张三",
+        user="李四",
+        dialogue="错误",
+        reason="系统故障"
+    )
+    print(resp)