feat: add Ollama client and bounded request queue
This commit is contained in:
0
backend/app/llm/__init__.py
Normal file
0
backend/app/llm/__init__.py
Normal file
18
backend/app/llm/client.py
Normal file
18
backend/app/llm/client.py
Normal file
@@ -0,0 +1,18 @@
|
||||
import httpx
|
||||
|
||||
from app.config import settings
|
||||
|
||||
|
||||
class OllamaClient:
|
||||
def __init__(self, base_url: str | None = None):
|
||||
self._base_url = base_url or settings.ollama_base_url
|
||||
|
||||
async def generate(self, model: str, prompt: str, system: str | None = None) -> str:
|
||||
payload = {"model": model, "prompt": prompt, "stream": False}
|
||||
if system is not None:
|
||||
payload["system"] = system
|
||||
|
||||
async with httpx.AsyncClient(base_url=self._base_url, timeout=120.0) as http_client:
|
||||
response = await http_client.post("/api/generate", json=payload)
|
||||
response.raise_for_status()
|
||||
return response.json()["response"]
|
||||
46
backend/app/llm/queue.py
Normal file
46
backend/app/llm/queue.py
Normal file
@@ -0,0 +1,46 @@
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class QueueFullError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class LLMQueue:
|
||||
"""Bounds concurrent Ollama calls and rejects work once too much is queued."""
|
||||
|
||||
def __init__(self, max_concurrency: int, max_queue_depth: int):
|
||||
self._semaphore = asyncio.Semaphore(max_concurrency)
|
||||
self._max_queue_depth = max_queue_depth
|
||||
self._waiting = 0
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def submit(self, coro_factory: Callable[[], Awaitable[T]]) -> T:
|
||||
async with self._lock:
|
||||
if self._waiting >= self._max_queue_depth:
|
||||
raise QueueFullError("too many seekers right now")
|
||||
self._waiting += 1
|
||||
|
||||
acquired = False
|
||||
try:
|
||||
# _waiting counts requests queued behind the concurrency limit,
|
||||
# not requests currently running. It must drop the instant a slot
|
||||
# is won (before coro_factory() runs), otherwise a long-running
|
||||
# call keeps counting against the queue depth for its whole
|
||||
# duration and wrongly rejects callers that should have been
|
||||
# queued. The decrement below has no `await` before it, so it
|
||||
# can't be split by cancellation from the acquire() that preceded
|
||||
# it - either we own a slot and will decrement, or we don't and
|
||||
# the `finally` below decrements instead.
|
||||
await self._semaphore.acquire()
|
||||
acquired = True
|
||||
self._waiting -= 1
|
||||
return await coro_factory()
|
||||
finally:
|
||||
if acquired:
|
||||
self._semaphore.release()
|
||||
else:
|
||||
self._waiting -= 1
|
||||
Reference in New Issue
Block a user