Spaces:
Runtime error
Runtime error
File size: 2,856 Bytes
1f8a981 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
from __future__ import annotations
import os
from typing import Iterator, List, Dict, Any
import httpx
class APIClient:
"""Simple client for the LLM backend API."""
def __init__(self, server: str = "http://localhost:8000", api_key: str | None = None) -> None:
self._server = server.rstrip("/")
self._headers = {"X-API-Key": api_key} if api_key else {}
# ------------------------------------------------------------------
# Helper methods
# ------------------------------------------------------------------
def _url(self, path: str) -> str:
return f"{self._server}{path}"
# ------------------------------------------------------------------
# Public API methods
# ------------------------------------------------------------------
def list_sessions(self, user: str) -> List[str]:
resp = httpx.get(self._url(f"/sessions/{user}"), headers=self._headers)
resp.raise_for_status()
data = resp.json()
return data.get("sessions", [])
def stream_chat(self, user: str, session: str, prompt: str) -> Iterator[str]:
with httpx.stream(
"POST",
self._url("/chat/stream"),
json={"user": user, "session": session, "prompt": prompt},
headers=self._headers,
timeout=None,
) as resp:
resp.raise_for_status()
for line in resp.iter_lines():
if line:
yield line.decode()
def upload_document(self, user: str, session: str, path: str) -> str:
name = os.path.basename(path)
with open(path, "rb") as f:
files = {"file": (name, f)}
data = {"user": user, "session": session}
resp = httpx.post(self._url("/upload"), data=data, files=files, headers=self._headers)
resp.raise_for_status()
return resp.json()["path"]
def list_vm_dir(self, user: str, path: str = "/data") -> List[Dict[str, Any]]:
resp = httpx.get(self._url(f"/vm/{user}/list"), params={"path": path}, headers=self._headers)
resp.raise_for_status()
return resp.json().get("entries", [])
def read_vm_file(self, user: str, path: str) -> str:
resp = httpx.get(self._url(f"/vm/{user}/file"), params={"path": path}, headers=self._headers)
resp.raise_for_status()
return resp.json().get("content", "")
def write_vm_file(self, user: str, path: str, content: str) -> None:
payload = {"path": path, "content": content}
resp = httpx.post(self._url(f"/vm/{user}/file"), json=payload, headers=self._headers)
resp.raise_for_status()
def delete_vm_file(self, user: str, path: str) -> None:
resp = httpx.delete(self._url(f"/vm/{user}/file"), params={"path": path}, headers=self._headers)
resp.raise_for_status()
|