数据源

import os
import pathlib
import sys
from urllib.parse import quote
import requests
import pdftotext

base_url = os.getenv("COUNSO_BASE_URL", "https://app.counso.ai")
workspace_id = os.environ["COUNSO_WORKSPACE_ID"]
space_id = os.environ["COUNSO_SPACE_ID"]
data_source_id = os.environ["COUNSO_DATA_SOURCE_ID"]
api_key = os.environ["COUNSO_API_KEY"]

def upload(text, file):
    document_id = file.relative_to(directory).as_posix()
    url = (f"{base_url}/api/v1/w/{workspace_id}/spaces/{space_id}"
           f"/data_sources/{data_source_id}/documents/{quote(document_id, safe='')}")
    return requests.post(
        url,
        headers={"Authorization": f"Bearer {api_key}"},
        json={"text": text},
    )

directory = pathlib.Path(sys.argv[1])
for file in directory.rglob("*"):
    if not file.is_file():
        continue
    if file.suffix.lower() == ".pdf":
        with file.open("rb") as source:
            text = "\n\n".join(pdftotext.PDF(source))
    elif file.suffix.lower() in {".txt", ".md"}:
        text = file.read_text()
    else:
        continue
    response = upload(text, file)
    if response.status_code != 200:
        print("导入失败：", file, response.status_code, response.text)
    else:
        print("已导入：", file)