49 lines
1.8 KiB
Python
49 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
|
|
class FastGPTClient:
|
|
def __init__(self, *, base_url: str, dataset_api_key: str, timeout: float = 60.0) -> None:
|
|
self.base_url = base_url.rstrip("/")
|
|
self.dataset_api_key = dataset_api_key.strip()
|
|
self.timeout = timeout
|
|
|
|
@property
|
|
def enabled(self) -> bool:
|
|
return bool(self.base_url and self.dataset_api_key)
|
|
|
|
async def ensure_dataset(self, name: str, intro: str = "") -> dict[str, Any] | None:
|
|
if not self.enabled:
|
|
return None
|
|
payload = {"name": name, "intro": intro}
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
response = await client.post(
|
|
f"{self.base_url}/api/core/dataset/create",
|
|
headers={"Authorization": f"Bearer {self.dataset_api_key}"},
|
|
json=payload,
|
|
)
|
|
response.raise_for_status()
|
|
return response.json().get("data") or response.json()
|
|
|
|
async def add_text_document(self, dataset_id: str, name: str, text: str) -> dict[str, Any] | None:
|
|
if not self.enabled or not dataset_id.strip():
|
|
return None
|
|
payload = {
|
|
"datasetId": dataset_id,
|
|
"type": "text",
|
|
"name": name,
|
|
"trainingType": "chunk",
|
|
"text": text,
|
|
}
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
response = await client.post(
|
|
f"{self.base_url}/api/core/dataset/collection/create/text",
|
|
headers={"Authorization": f"Bearer {self.dataset_api_key}"},
|
|
json=payload,
|
|
)
|
|
response.raise_for_status()
|
|
return response.json().get("data") or response.json()
|