133 lines
5.2 KiB
Python
133 lines
5.2 KiB
Python
import asyncio
|
|
from typing import Any, Dict, List, Optional
|
|
from urllib.parse import quote
|
|
|
|
import httpx
|
|
|
|
from config import RAG_CONFIG, WIKI_CONFIG
|
|
|
|
|
|
class HTKnowClient:
|
|
"""Small async client for the 11000 knowledge service."""
|
|
|
|
def __init__(self, base_url: Optional[str] = None):
|
|
self.base_url = (base_url or WIKI_CONFIG["htknow_base_url"]).rstrip("/")
|
|
|
|
def _headers(self) -> Dict[str, str]:
|
|
return {
|
|
"accept": "application/json",
|
|
"x-user-id": _safe_header(RAG_CONFIG.get("x-user-id") or WIKI_CONFIG["default_user_id"]),
|
|
"x-user-name": _safe_header(RAG_CONFIG.get("x-user-name") or WIKI_CONFIG["default_user_name"]),
|
|
"x-role": _safe_header(RAG_CONFIG.get("x-role") or WIKI_CONFIG["default_role"]),
|
|
}
|
|
|
|
async def list_files(
|
|
self,
|
|
kb_id: Optional[str] = None,
|
|
limit: int = 20,
|
|
offset: int = 0,
|
|
) -> List[Dict[str, Any]]:
|
|
url = f"{self.base_url}/api/v1/knowledge/files/"
|
|
# The 11000 service is more stable with its default-ish page sizes.
|
|
# Fetch at least 10 rows and trim locally instead of requesting tiny pages.
|
|
size = min(max(limit, 10), 100)
|
|
page = max(offset // size, 0) + 1
|
|
collected: List[Dict[str, Any]] = []
|
|
async with httpx.AsyncClient(timeout=60.0, verify=False, trust_env=False) as client:
|
|
while len(collected) < limit:
|
|
params: Dict[str, Any] = {"page": page, "size": size}
|
|
if kb_id:
|
|
params["kb_id"] = kb_id
|
|
payload = await _get_json_with_retry(client, url, params=params, headers=self._headers())
|
|
items = _extract_list(payload, candidates=("results", "data", "items", "files"))
|
|
if not items:
|
|
break
|
|
collected.extend(items)
|
|
if len(items) < params["size"]:
|
|
break
|
|
page += 1
|
|
return collected[:limit]
|
|
|
|
async def get_file_slices(self, file_id: str, limit: Optional[int] = None) -> List[Dict[str, Any]]:
|
|
url = f"{self.base_url}/api/v1/knowledge/files/{file_id}/slices"
|
|
params = {}
|
|
if limit:
|
|
params["limit"] = limit
|
|
async with httpx.AsyncClient(timeout=120.0, verify=False, trust_env=False) as client:
|
|
payload = await _get_json_with_retry(client, url, params=params, headers=self._headers())
|
|
return _extract_list(payload, candidates=("slices", "results", "data", "items"))
|
|
|
|
async def get_file(self, file_id: str) -> Dict[str, Any]:
|
|
url = f"{self.base_url}/api/v1/knowledge/files/{file_id}"
|
|
async with httpx.AsyncClient(timeout=60.0, verify=False, trust_env=False) as client:
|
|
payload = await _get_json_with_retry(client, url, params={}, headers=self._headers())
|
|
if isinstance(payload, dict):
|
|
for key in ("data", "file", "item"):
|
|
value = payload.get(key)
|
|
if isinstance(value, dict):
|
|
return value
|
|
return payload
|
|
return {}
|
|
|
|
async def get_kb_tree(self) -> List[Dict[str, Any]]:
|
|
url = f"{self.base_url}/api/v1/knowledge/knowledge_base/tree"
|
|
async with httpx.AsyncClient(timeout=60.0, verify=False, trust_env=False) as client:
|
|
payload = await _get_json_with_retry(client, url, params={}, headers=self._headers())
|
|
if isinstance(payload, list):
|
|
return [x for x in payload if isinstance(x, dict)]
|
|
if isinstance(payload, dict):
|
|
data = payload.get("data") or payload.get("list") or payload
|
|
if isinstance(data, dict) and isinstance(data.get("children"), list):
|
|
return [x for x in data["children"] if isinstance(x, dict)]
|
|
if isinstance(data, list):
|
|
return [x for x in data if isinstance(x, dict)]
|
|
return []
|
|
|
|
|
|
def _extract_list(payload: Any, candidates: tuple[str, ...]) -> List[Dict[str, Any]]:
|
|
if isinstance(payload, list):
|
|
return [x for x in payload if isinstance(x, dict)]
|
|
if not isinstance(payload, dict):
|
|
return []
|
|
for key in candidates:
|
|
value = payload.get(key)
|
|
if isinstance(value, list):
|
|
return [x for x in value if isinstance(x, dict)]
|
|
if isinstance(value, dict):
|
|
nested = _extract_list(value, candidates)
|
|
if nested:
|
|
return nested
|
|
return []
|
|
|
|
|
|
def _safe_header(value: Any) -> str:
|
|
text = str(value or "")
|
|
try:
|
|
text.encode("latin-1")
|
|
return text
|
|
except UnicodeEncodeError:
|
|
return quote(text, safe="")
|
|
|
|
|
|
async def _get_json_with_retry(
|
|
client: httpx.AsyncClient,
|
|
url: str,
|
|
params: Dict[str, Any],
|
|
headers: Dict[str, str],
|
|
retries: int = 2,
|
|
) -> Any:
|
|
last_exc: Optional[Exception] = None
|
|
for attempt in range(retries + 1):
|
|
try:
|
|
response = await client.get(url, params=params, headers=headers)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except (httpx.HTTPStatusError, httpx.RequestError) as exc:
|
|
last_exc = exc
|
|
if attempt >= retries:
|
|
raise
|
|
await asyncio.sleep(0.5 * (attempt + 1))
|
|
if last_exc:
|
|
raise last_exc
|
|
return {}
|