RAG 数据摄取流水线

import hashlib

def generate_hash(content):
    return hashlib.sha256(content.encode('utf-8')).hexdigest()

dead_letter_queue = []

try:
    process_document(document)
except Exception as e:
    dead_letter_queue.append({"document": document, "error": str(e)})

def save_checkpoint(last_document_id):
    with open("checkpoint.txt", "w") as f:
        f.write(last_document_id)

def load_checkpoint():
    try:
        with open("checkpoint.txt", "r") as f:
            return f.read().strip()
    except FileNotFoundError:
        return None

class FeishuReader:
    def __init__(self, client, app_id, app_secret, sleep=time.sleep, clock=time.monotonic):
        self.client = client
        self.app_id = app_id
        self.app_secret = app_secret
        self.sleep = sleep
        self.clock = clock
        self.token = None
        self.expires_at = 0

    def access_token(self):
        if self.token and self.clock() < self.expires_at:
            return self.token
        r = self.client.post('/open-apis/auth/v3/tenant_access_token/internal',
                             json={'app_id': self.app_id, 'app_secret': self.app_secret})
        r.raise_for_status()
        result = r.json()
        if result.get('code', 0) != 0:
            raise SourceError(f"tenant_token code={result['code']}")
        self.token = result['tenant_access_token']
        self.expires_at = self.clock() + max(0, result['expire'] - 60)
        return self.token

    def request(self, method, path, params=None, body=None):
        refreshed = False
        for attempt in range(5):
            r = self.client.request(method, path, params=params, json=body,
                                    headers={'Authorization': f'Bearer {self.access_token()}'})
            result = r.json() if 'json' in r.headers.get('content-type', '') else {}
            code = result.get('code', 0)
            if (r.status_code == 401 or code in {99991663, 99991668}) and not refreshed:
                self.token = None
                refreshed = True
                continue
            if r.status_code == 429 or code == 99991400 or r.status_code >= 500:
                if attempt == 4:
                    raise SourceError(f'retry_exhausted status={r.status_code} code={code}')
                retry_after = r.headers.get('Retry-After')
                delay = float(retry_after) if retry_after and retry_after.isdecimal() else 2 ** attempt + random.uniform(0, .2)
                self.sleep(min(delay, 30))
                continue
            r.raise_for_status()
            if code != 0:
                raise SourceError(f'feishu status={r.status_code} code={code}')
            return result['data']
        raise SourceError('token_refresh_exhausted')

    def pages(self, method, path, page_size, params=None, body=None):
        cursor = None
        seen = set()
        while True:
            query = dict(params or {}, page_size=page_size)
            if cursor:
                query['page_token'] = cursor
            data = self.request(method, path, query, body)
            yield from data.get('items', [])
            if not data.get('has_more', False):
                return
            cursor = data.get('page_token')
            if not cursor or cursor in seen:
                raise SourceError('invalid_page_token')
            seen.add(cursor)

    def wiki_nodes(self, space_id, parent=None):
        params = {'parent_node_token': parent} if parent else {}
        for node in self.pages('GET', f'/open-apis/wiki/v2/spaces/{token_path(space_id)}/nodes', 50, params):
            yield node
            if node.get('has_child'):
                yield from self.wiki_nodes(space_id, node['node_token'])

    def read_node(self, node, table_id=None):
        kind, token = node['obj_type'], token_path(node['obj_token'])
        if kind == 'docx':
            blocks = list(self.pages('GET', f'/open-apis/docx/v1/documents/{token}/blocks', 500))
            text, omissions = blocks_to_text(blocks)
            return text, omissions
        if kind in {'doc', 'document'}:
            data = self.request('GET', f'/open-apis/doc/v2/{token}/raw_content')
            return data['content'], ['旧版文档纯文本缺少表格结构与附件内容']
        if kind == 'bitable':
            if not table_id:
                raise SourceError('bitable_table_id_required')
            records = list(self.pages('POST', f'/open-apis/bitable/v1/apps/{token}/tables/{token_path(table_id)}/records/search',
                                      500, body={'automatic_fields': False}))
            return '\n'.join(json.dumps({'record_id': r['record_id'], 'fields': r['fields']}, ensure_ascii=False) for r in records), []
        raise SourceError(f'unsupported_obj_type={kind}')

import time
from urllib.parse import quote


class DingTalkReader:
    def __init__(self, client, app_key, app_secret, operator_union_id):
        self.client, self.app_key, self.app_secret = client, app_key, app_secret
        self.operator = operator_union_id
        self.token, self.expires_at = None, 0

    def request(self, method, path, params=None, body=None):
        if not self.token or time.monotonic() >= self.expires_at:
            response = self.client.post('/v1.0/oauth2/accessToken',
                                        json={'appKey': self.app_key, 'appSecret': self.app_secret})
            response.raise_for_status()
            data = response.json()
            self.token, self.expires_at = data['accessToken'], time.monotonic() + max(0, data['expireIn'] - 60)
        response = self.client.request(method, path, params=params, json=body,
            headers={'x-acs-dingtalk-access-token': self.token})
        response.raise_for_status()
        return response.json()

    def records(self, base_id, sheet_id):
        next_token, seen = None, set()
        while True:
            body = {'operatorId': self.operator, 'maxResults': 100}
            if next_token:
                body['nextToken'] = next_token
            data = self.request('POST', f'/v1.0/notable/bases/{quote(base_id, safe="")}/sheets/{quote(sheet_id, safe="")}/records/list', body=body)
            yield from data.get('records', [])
            next_token = data.get('nextToken')
            if not next_token:
                return
            if next_token in seen:
                raise ValueError('repeated_next_token')
            seen.add(next_token)

    def document_blocks(self, doc_key, start_index, end_index):
        if start_index < 0 or end_index <= start_index:
            raise ValueError('invalid_block_range')
        return self.request('GET', f'/v1.0/doc/suites/documents/{quote(doc_key, safe="")}/blocks',
            params={'operatorId': self.operator, 'startIndex': start_index, 'endIndex': end_index})

import pymupdf

with pymupdf.open("demo-policy.pdf") as document:
    pages = [{"page": i + 1, "text": page.get_text(sort=True)}
             for i, page in enumerate(document)]
print(pages)

import pymupdf
from docx import Document
from docx.table import Table
from docx.text.paragraph import Paragraph


def parse_pdf(path):
    pages = []
    with pymupdf.open(path) as document:
        if document.is_encrypted:
            raise PermissionError('encrypted_pdf_requires_authorized_password')
        for index, page in enumerate(document):
            text = page.get_text(sort=True)
            if not text.strip():
                raise ValueError(f'pdf_page_requires_ocr={index + 1}')
            tables = [table.extract() for table in page.find_tables().tables]
            pages.append({'page': index + 1, 'text': text, 'tables': tables})
    return pages


def parse_docx(path):
    document = Document(path)
    if document.inline_shapes:
        raise ValueError('docx_embedded_images_require_review')
    blocks = []
    for index, item in enumerate(document.iter_inner_content()):
        if isinstance(item, Paragraph):
            blocks.append({'block': index, 'type': 'paragraph', 'text': item.text})
        elif isinstance(item, Table):
            if any(cell.tables for row in item.rows for cell in row.cells):
                raise ValueError('docx_nested_tables_require_review')
            blocks.append({'block': index, 'type': 'table', 'rows': [[cell.text for cell in row.cells] for row in item.rows]})
    return blocks