knowledge/utils/fei_shu.py

import os
import uuid
from typing import Any, Dict, Optional

import lark_oapi as lark
from lark_oapi.core.enum import LogLevel
import config


class FeiShu:

    def __init__(self, file_token:str = None):
        app_id = os.getenv('FEISHU_APP_ID', config.FEISHU_APP_ID)
        app_secret = os.getenv('FEISHU_APP_SECRET', config.FEISHU_APP_SECRET)
        self.file_token = file_token if file_token is not None else os.getenv('FEISHU_FILE_TOKEN')
        self.client = (lark.client.Client().builder()
                       .app_id(app_id)
                       .app_secret(app_secret)
                       .log_level(LogLevel.INFO)
                       .build())

    def create_table(self, name: str) -> str:
        request = (lark.bitable.v1.CreateAppTableRequest.builder()
                   .app_token(self.file_token)
                   .request_body(lark.bitable.v1.CreateAppTableRequestBody.builder()
                                 .table(lark.bitable.v1.ReqTable.builder()
                                        .name(name)
                                        .fields([
                                            lark.bitable.v1.AppTableCreateHeader.builder().field_name('原文链接').type(15).build(),
                                            lark.bitable.v1.AppTableCreateHeader.builder().field_name('抓取结果').type(1).build(),
                                            lark.bitable.v1.AppTableCreateHeader.builder().field_name('标题').type(1).build(),
                                            lark.bitable.v1.AppTableCreateHeader.builder().field_name('用户链接').type(15).build(),
                                            lark.bitable.v1.AppTableCreateHeader.builder().field_name('识别结果').type(1).build(),
                                            lark.bitable.v1.AppTableCreateHeader.builder().field_name('初步理解').type(1).build(),
                                            lark.bitable.v1.AppTableCreateHeader.builder().field_name('物理聚合').type(1).build(),
                                        ])
                                        .build())
                                 .build())
                   .build())
        response = self.client.bitable.v1.app_table.create(request)
        if not response.success():
            raise RuntimeError(f'多维表格创建新数据表失败: {lark.JSON.marshal(response.error)}')
        return response.data.table_id

    def get_all_records(self, table_id: str, page_token: Optional[str] = None):
        request = (lark.bitable.v1.SearchAppTableRecordRequest.builder()
                   .app_token(self.file_token)
                   .table_id(table_id)
                   .user_id_type('open_id')
                   .page_size(500)
                   .request_body(lark.bitable.v1.SearchAppTableRecordRequestBody.builder()
                                 .automatic_fields(True)
                                 .build()))
        if page_token:
            request = request.page_token(page_token)
        request = request.build()
        response = self.client.bitable.v1.app_table_record.search(request)
        if not response.success():
            raise RuntimeError(f'获取多维表格记录失败: {lark.JSON.marshal(response.error)}')
        return response.data

    def get_record(self, table_id: str, *record_ids: str):
        request = (lark.bitable.v1.BatchGetAppTableRecordRequest.builder()
                   .app_token(self.file_token)
                   .table_id(table_id)
                   .request_body(lark.bitable.v1.BatchGetAppTableRecordRequestBody.builder()
                                 .record_ids(list(record_ids))
                                 .user_id_type('open_id')
                                 .with_shared_url(True)
                                 .automatic_fields(True)
                                 .build())
                   .build())
        response = self.client.bitable.v1.app_table_record.batch_get(request)
        if not response.success():
            raise RuntimeError(f'获取多维表格指定记录失败: {lark.JSON.marshal(response.error)}')
        return response.data

    def create_record(self, table_id: str, *fields: Dict[str, Any]):
        request = (lark.bitable.v1.BatchCreateAppTableRecordRequest.builder()
                   .app_token(self.file_token)
                   .table_id(table_id)
                   .user_id_type('open_id')
                   .client_token(str(uuid.uuid4()))
                   .ignore_consistency_check(False)
                   .request_body(lark.bitable.v1.BatchCreateAppTableRecordRequestBody.builder()
                                 .records([lark.bitable.v1.AppTableRecord.builder()
                                          .fields(item)
                                          .build() for item in fields])
                                 .build())
                   .build())
        response = self.client.bitable.v1.app_table_record.batch_create(request)
        if not response.success():
            raise RuntimeError(f'向多维表格添加记录失败: {lark.JSON.marshal(response.error)}')
        return response.data

    def update_record(self, table_id: str, *records: lark.bitable.v1.AppTableRecord):
        request = (lark.bitable.v1.BatchUpdateAppTableRecordRequest.builder()
                   .app_token(self.file_token)
                   .table_id(table_id)
                   .user_id_type('open_id')
                   .ignore_consistency_check(False)
                   .request_body(lark.bitable.v1.BatchUpdateAppTableRecordRequestBody.builder()
                                 .records(list(records))
                                 .build())
                   .build())
        response = self.client.bitable.v1.app_table_record.batch_update(request)
        if not response.success():
            raise RuntimeError(f'更新多维表格指定记录失败: {lark.JSON.marshal(response.error)}')
        return response.data

    def delete_record(self, table_id: str, *record_ids: str):
        request = (lark.bitable.v1.BatchDeleteAppTableRecordRequest.builder()
                   .app_token(self.file_token)
                   .table_id(table_id)
                   .request_body(lark.bitable.v1.BatchDeleteAppTableRecordRequestBody.builder()
                                 .records(list(record_ids))
                                 .build())
                   .build())
        response = self.client.bitable.v1.app_table_record.batch_delete(request)
        if not response.success():
            raise RuntimeError(f'删除多维表格指定记录失败: {lark.JSON.marshal(response.error)}')
        return response.data


if __name__ == '__main__':
    from dotenv import load_dotenv

    try:
        load_dotenv()
    except ImportError:
        raise RuntimeError('导入环境变量失败')

    feishu = FeiShu()
    # 创建数据表
    new_table_id = feishu.create_table('测试数据表')
    # 新增记录
    new_fields = [
        {
            '原文链接': {
                'link': 'https://www.baidu.com',
                'text': 'https://www.baidu.com',
            },
            '抓取结果': '这是抓取结果1',
        },
        {
            '原文链接': {
                'link': 'https://www.qq.com',
                'text': 'https://www.qq.com',
            },
            '抓取结果': '这是抓取结果2',
        }
    ]
    feishu.create_record(new_table_id, *new_fields)
    # 获取全部记录
    get_result = feishu.get_all_records(new_table_id)
    has_more = get_result.has_more  # 是否有下一页
    next_page_token = get_result.page_token  # 下一页token
    new_record_ids = []
    for record in get_result.items:
        new_record_ids.append(record.record_id)
        print(record.fields)
    # 更新记录
    new_record = (lark.bitable.v1.AppTableRecord.builder()
                  .record_id(new_record_ids[0])
                  .fields({'识别结果': '这是识别结果'})
                  .build())
    feishu.update_record(new_table_id, new_record)
    # 获取指定ID记录
    get_result = feishu.get_record(new_table_id, *new_record_ids)
    for record in get_result.records:
        new_record_ids.append(record.record_id)
        print(record.fields)
    # 删除指定ID记录
    feishu.delete_record(new_table_id, new_record_ids[1])