| """Thin async client for the internal llama-server instance. |
| |
| llama-server is started with --parallel N --cont-batching (see entrypoint.sh), |
| which makes it batch multiple in-flight requests' forward passes together |
| instead of running them one at a time. This is where the actual CPU |
| efficiency gain for many short strings comes from: we just need to keep |
| enough requests in flight concurrently, and llama.cpp does the batching. |
| """ |
|
|
| import asyncio |
|
|
| import httpx |
|
|
| from . import config |
|
|
| _client: httpx.AsyncClient | None = None |
| |
| |
| |
| _semaphore = asyncio.Semaphore(config.PARALLEL_SLOTS) |
|
|
|
|
| def get_client() -> httpx.AsyncClient: |
| global _client |
| if _client is None: |
| _client = httpx.AsyncClient( |
| base_url=config.LLAMA_SERVER_URL, |
| timeout=httpx.Timeout(config.REQUEST_TIMEOUT_SECONDS), |
| ) |
| return _client |
|
|
|
|
| async def aclose() -> None: |
| global _client |
| if _client is not None: |
| await _client.aclose() |
| _client = None |
|
|
|
|
| class LlamaServerError(RuntimeError): |
| pass |
|
|
|
|
| async def chat_complete( |
| messages: list[dict[str, str]], max_tokens: int | None = None |
| ) -> str: |
| client = get_client() |
| payload = { |
| "messages": messages, |
| "temperature": config.TEMPERATURE, |
| "top_p": config.TOP_P, |
| "top_k": config.TOP_K, |
| "repeat_penalty": config.REPEAT_PENALTY, |
| "max_tokens": max_tokens or config.MAX_TOKENS, |
| "stream": False, |
| } |
| async with _semaphore: |
| try: |
| resp = await client.post("/v1/chat/completions", json=payload) |
| resp.raise_for_status() |
| except httpx.HTTPError as exc: |
| raise LlamaServerError(str(exc)) from exc |
|
|
| data = resp.json() |
| try: |
| return data["choices"][0]["message"]["content"].strip() |
| except (KeyError, IndexError) as exc: |
| raise LlamaServerError(f"Unexpected llama-server response: {data}") from exc |
|
|
|
|
| async def health() -> bool: |
| client = get_client() |
| try: |
| resp = await client.get("/health", timeout=5.0) |
| return resp.status_code == 200 |
| except httpx.HTTPError: |
| return False |
|
|
