410 lines
14 KiB
Python
410 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from typing import Iterator
|
|
|
|
import httpx
|
|
|
|
from .errors import SharedAIPlatformError
|
|
from .types import PollOptions, RequestContext, SSEEvent, SyncRecord
|
|
|
|
|
|
class SharedAIPlatformClient:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
base_url: str = "http://127.0.0.1:8011",
|
|
api_key: str,
|
|
app_id: str,
|
|
timeout_seconds: float = 30.0,
|
|
client: httpx.Client | None = None,
|
|
api_prefix: str = "/api/v0.2",
|
|
):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.api_prefix = self._normalize_api_prefix(api_prefix)
|
|
self.api_key = api_key
|
|
self.app_id = app_id
|
|
self._owns_client = client is None
|
|
self._client = client or httpx.Client(base_url=self.base_url, timeout=timeout_seconds)
|
|
|
|
def close(self):
|
|
if self._owns_client:
|
|
self._client.close()
|
|
|
|
def __enter__(self) -> "SharedAIPlatformClient":
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
self.close()
|
|
|
|
def list_tasks(self, context: RequestContext) -> list[dict]:
|
|
return self._request_json("GET", self._path("/tasks"), context=context)["data"]
|
|
|
|
def run_task(
|
|
self,
|
|
context: RequestContext,
|
|
task_key: str,
|
|
*,
|
|
inputs: dict,
|
|
provider: str | None = None,
|
|
metadata: dict | None = None,
|
|
) -> dict:
|
|
return self._request_json(
|
|
"POST",
|
|
self._path(f"/tasks/{task_key}/run"),
|
|
context=context,
|
|
json_body={"inputs": inputs, "provider": provider, "metadata": metadata or {}},
|
|
)["data"]
|
|
|
|
def enqueue_task(
|
|
self,
|
|
context: RequestContext,
|
|
task_key: str,
|
|
*,
|
|
inputs: dict,
|
|
provider: str | None = None,
|
|
metadata: dict | None = None,
|
|
) -> dict:
|
|
return self._request_json(
|
|
"POST",
|
|
self._path(f"/tasks/{task_key}/enqueue"),
|
|
context=context,
|
|
json_body={"inputs": inputs, "provider": provider, "metadata": metadata or {}},
|
|
)["data"]
|
|
|
|
def get_task_run(self, context: RequestContext, run_id: str) -> dict:
|
|
return self._request_json("GET", self._path(f"/task-runs/{run_id}"), context=context)["data"]
|
|
|
|
def list_jobs(self, context: RequestContext, *, status_value: str | None = None, limit: int = 50) -> list[dict]:
|
|
params = {"limit": limit}
|
|
if status_value:
|
|
params["status_value"] = status_value
|
|
return self._request_json("GET", self._path("/jobs"), context=context, params=params)["data"]
|
|
|
|
def get_job(self, context: RequestContext, job_id: str) -> dict:
|
|
return self._request_json("GET", self._path(f"/jobs/{job_id}"), context=context)["data"]
|
|
|
|
def run_next_job(self, context: RequestContext) -> dict | None:
|
|
return self._request_json("POST", self._path("/jobs/run-next"), context=context)["data"]["job"]
|
|
|
|
def poll_job_until_terminal(self, context: RequestContext, job_id: str, options: PollOptions | None = None) -> dict:
|
|
opts = options or PollOptions()
|
|
deadline = time.time() + opts.timeout_seconds
|
|
while True:
|
|
job = self.get_job(context, job_id)
|
|
if job["status"] in opts.terminal_statuses:
|
|
return job
|
|
if time.time() >= deadline:
|
|
raise TimeoutError(f"job {job_id} did not reach a terminal state within {opts.timeout_seconds} seconds")
|
|
time.sleep(opts.interval_seconds)
|
|
|
|
def list_conversations(self, context: RequestContext, *, page: int = 1, page_size: int = 20) -> dict:
|
|
return self._request_json(
|
|
"GET",
|
|
self._path("/conversations"),
|
|
context=context,
|
|
params={"page": page, "page_size": page_size},
|
|
)["data"]
|
|
|
|
def get_conversation_detail(self, context: RequestContext, conversation_id: str) -> dict:
|
|
return self._request_json("GET", self._path(f"/conversations/{conversation_id}"), context=context)["data"]
|
|
|
|
def list_knowledge_bases(self, context: RequestContext, *, page: int = 1, page_size: int = 20) -> dict:
|
|
return self._request_json(
|
|
"GET",
|
|
self._path("/knowledge-bases"),
|
|
context=context,
|
|
params={"page": page, "page_size": page_size},
|
|
)["data"]
|
|
|
|
def list_documents(self, context: RequestContext, knowledge_base_id: str, *, page: int = 1, page_size: int = 20) -> dict:
|
|
return self._request_json(
|
|
"GET",
|
|
self._path(f"/knowledge-bases/{knowledge_base_id}/documents"),
|
|
context=context,
|
|
params={"page": page, "page_size": page_size},
|
|
)["data"]
|
|
|
|
def create_knowledge_base(
|
|
self,
|
|
context: RequestContext,
|
|
*,
|
|
name: str,
|
|
description: str | None = None,
|
|
retrieval_mode: str = "hybrid",
|
|
) -> dict:
|
|
return self._request_json(
|
|
"POST",
|
|
self._path("/knowledge-bases"),
|
|
context=context,
|
|
json_body={"name": name, "description": description, "retrieval_mode": retrieval_mode},
|
|
)["data"]
|
|
|
|
def list_knowledge_base_versions(self, context: RequestContext, knowledge_base_id: str) -> list[dict]:
|
|
return self._request_json(
|
|
"GET",
|
|
self._path(f"/knowledge-bases/{knowledge_base_id}/versions"),
|
|
context=context,
|
|
)["data"]
|
|
|
|
def publish_knowledge_base(self, context: RequestContext, knowledge_base_id: str) -> dict:
|
|
return self._request_json(
|
|
"POST",
|
|
self._path(f"/knowledge-bases/{knowledge_base_id}/publish"),
|
|
context=context,
|
|
)["data"]
|
|
|
|
def unpublish_knowledge_base(self, context: RequestContext, knowledge_base_id: str) -> dict:
|
|
return self._request_json(
|
|
"POST",
|
|
self._path(f"/knowledge-bases/{knowledge_base_id}/unpublish"),
|
|
context=context,
|
|
)["data"]
|
|
|
|
def rollback_knowledge_base(
|
|
self,
|
|
context: RequestContext,
|
|
knowledge_base_id: str,
|
|
*,
|
|
target_version_label: str,
|
|
change_summary: str | None = None,
|
|
) -> dict:
|
|
return self._request_json(
|
|
"POST",
|
|
self._path(f"/knowledge-bases/{knowledge_base_id}/rollback"),
|
|
context=context,
|
|
json_body={"target_version_label": target_version_label, "change_summary": change_summary},
|
|
)["data"]
|
|
|
|
def upload_document(
|
|
self,
|
|
context: RequestContext,
|
|
knowledge_base_id: str,
|
|
*,
|
|
file_name: str,
|
|
mime_type: str | None,
|
|
content: bytes,
|
|
) -> dict:
|
|
response = self._client.post(
|
|
self._path(f"/knowledge-bases/{knowledge_base_id}/documents"),
|
|
headers=self._headers(context),
|
|
files={"file": (file_name, content, mime_type or "application/octet-stream")},
|
|
)
|
|
self._raise_for_status(response)
|
|
payload = response.json()
|
|
if payload.get("code") not in (0, None):
|
|
raise SharedAIPlatformError(
|
|
status_code=response.status_code,
|
|
message=str(payload.get("message") or "request failed"),
|
|
code=payload.get("code"),
|
|
data=payload.get("data"),
|
|
)
|
|
return payload["data"]
|
|
|
|
def retry_document(
|
|
self,
|
|
context: RequestContext,
|
|
knowledge_base_id: str,
|
|
document_id: str,
|
|
*,
|
|
file_name: str | None = None,
|
|
mime_type: str | None = None,
|
|
content: bytes | None = None,
|
|
) -> dict:
|
|
request_kwargs = {"headers": self._headers(context)}
|
|
if content is not None:
|
|
request_kwargs["files"] = {"file": (file_name or "document.txt", content, mime_type or "application/octet-stream")}
|
|
response = self._client.post(
|
|
self._path(f"/knowledge-bases/{knowledge_base_id}/documents/{document_id}/retry"),
|
|
**request_kwargs,
|
|
)
|
|
self._raise_for_status(response)
|
|
payload = response.json()
|
|
if payload.get("code") not in (0, None):
|
|
raise SharedAIPlatformError(
|
|
status_code=response.status_code,
|
|
message=str(payload.get("message") or "request failed"),
|
|
code=payload.get("code"),
|
|
data=payload.get("data"),
|
|
)
|
|
return payload["data"]
|
|
|
|
def sync_content(
|
|
self,
|
|
context: RequestContext,
|
|
knowledge_base_id: str,
|
|
*,
|
|
records: list[SyncRecord | dict],
|
|
schema_version: str = "content-sync.v1",
|
|
sync_mode: str = "incremental",
|
|
batch_id: str | None = None,
|
|
cursor_token: str | None = None,
|
|
) -> dict:
|
|
normalized_records = [record.to_dict() if isinstance(record, SyncRecord) else record for record in records]
|
|
return self._request_json(
|
|
"POST",
|
|
self._path(f"/knowledge-bases/{knowledge_base_id}/sync-content"),
|
|
context=context,
|
|
json_body={
|
|
"schema_version": schema_version,
|
|
"sync_mode": sync_mode,
|
|
"batch_id": batch_id,
|
|
"cursor_token": cursor_token,
|
|
"records": normalized_records,
|
|
},
|
|
)["data"]
|
|
|
|
def list_sync_jobs(self, context: RequestContext, knowledge_base_id: str) -> dict:
|
|
return self._request_json("GET", self._path(f"/knowledge-bases/{knowledge_base_id}/sync-jobs"), context=context)["data"]
|
|
|
|
def run_legal_review(
|
|
self,
|
|
context: RequestContext,
|
|
*,
|
|
title: str,
|
|
text: str,
|
|
contract_type: str | None = None,
|
|
review_as: str = "party_a",
|
|
provider: str | None = None,
|
|
) -> dict:
|
|
return self._request_json(
|
|
"POST",
|
|
self._path("/domain-packs/legal-review/reviews/run"),
|
|
context=context,
|
|
json_body={
|
|
"title": title,
|
|
"contract_type": contract_type,
|
|
"review_as": review_as,
|
|
"text": text,
|
|
"provider": provider,
|
|
},
|
|
)["data"]
|
|
|
|
def stream_conversation(
|
|
self,
|
|
context: RequestContext,
|
|
*,
|
|
message: str,
|
|
chat_context: dict | None = None,
|
|
provider: str | None = None,
|
|
template_key: str | None = None,
|
|
conversation_id: str | None = None,
|
|
) -> Iterator[SSEEvent]:
|
|
headers = self._headers(context)
|
|
with self._client.stream(
|
|
"POST",
|
|
self._path("/conversations/stream"),
|
|
headers=headers,
|
|
json={
|
|
"conversation_id": conversation_id,
|
|
"message": message,
|
|
"chat_context": chat_context or {},
|
|
"provider": provider,
|
|
"template_key": template_key,
|
|
},
|
|
) as response:
|
|
self._raise_for_status(response)
|
|
yield from self._iter_sse_events(response)
|
|
|
|
def _iter_sse_events(self, response: httpx.Response) -> Iterator[SSEEvent]:
|
|
data_lines: list[str] = []
|
|
|
|
for line in response.iter_lines():
|
|
if isinstance(line, bytes):
|
|
line = line.decode("utf-8")
|
|
|
|
if line == "":
|
|
payload = self._parse_sse_payload(data_lines)
|
|
data_lines = []
|
|
if payload is None:
|
|
continue
|
|
yield SSEEvent(event_type=str(payload.get("type") or "unknown"), payload=payload)
|
|
continue
|
|
|
|
if line.startswith("data:"):
|
|
data_part = line.removeprefix("data:").lstrip()
|
|
if data_part:
|
|
data_lines.append(data_part)
|
|
continue
|
|
|
|
if data_lines or line.startswith("{") or line.startswith("["):
|
|
data_lines.append(line)
|
|
|
|
payload = self._parse_sse_payload(data_lines)
|
|
if payload is not None:
|
|
yield SSEEvent(event_type=str(payload.get("type") or "unknown"), payload=payload)
|
|
|
|
@staticmethod
|
|
def _parse_sse_payload(data_lines: list[str]) -> dict | None:
|
|
if not data_lines:
|
|
return None
|
|
payload_text = "\n".join(part for part in data_lines if part).strip()
|
|
if not payload_text:
|
|
return None
|
|
return json.loads(payload_text)
|
|
|
|
def _request_json(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
*,
|
|
context: RequestContext | None = None,
|
|
json_body: dict | None = None,
|
|
params: dict | None = None,
|
|
) -> dict:
|
|
response = self._client.request(method, path, headers=self._headers(context), json=json_body, params=params)
|
|
self._raise_for_status(response)
|
|
payload = response.json()
|
|
if payload.get("code") not in (0, None):
|
|
raise SharedAIPlatformError(
|
|
status_code=response.status_code,
|
|
message=str(payload.get("message") or "request failed"),
|
|
code=payload.get("code"),
|
|
data=payload.get("data"),
|
|
)
|
|
return payload
|
|
|
|
def _headers(self, context: RequestContext | None = None) -> dict[str, str]:
|
|
headers: dict[str, str] = {}
|
|
if self.api_key.strip():
|
|
headers["Authorization"] = f"Bearer {self.api_key}"
|
|
if context is not None:
|
|
headers.update(
|
|
{
|
|
"X-Shared-AI-App-Id": self.app_id,
|
|
"X-Shared-AI-Tenant-Id": context.tenant_id,
|
|
"X-Shared-AI-User-Id": context.user_id,
|
|
"X-Shared-AI-User-Role": context.user_role,
|
|
}
|
|
)
|
|
return headers
|
|
|
|
def _raise_for_status(self, response: httpx.Response):
|
|
try:
|
|
response.raise_for_status()
|
|
except httpx.HTTPStatusError as exc:
|
|
if not response.is_closed:
|
|
response.read()
|
|
message = response.text
|
|
try:
|
|
payload = response.json()
|
|
message = str(payload.get("message") or message)
|
|
code = payload.get("code")
|
|
data = payload.get("data")
|
|
except Exception:
|
|
code = None
|
|
data = None
|
|
raise SharedAIPlatformError(response.status_code, message, code, data) from exc
|
|
|
|
@staticmethod
|
|
def _normalize_api_prefix(api_prefix: str) -> str:
|
|
normalized = (api_prefix or "/api/v0.2").strip().rstrip("/")
|
|
if not normalized.startswith("/"):
|
|
normalized = f"/{normalized}"
|
|
return normalized
|
|
|
|
def _path(self, suffix: str) -> str:
|
|
normalized_suffix = suffix if suffix.startswith("/") else f"/{suffix}"
|
|
return f"{self.api_prefix}{normalized_suffix}"
|