@ -1,6 +1,7 @@
""" Async Ollama HTTP client with retry logic. """
from __future__ import annotations
import asyncio
import json
import time
from typing import Any
@ -15,66 +16,111 @@ logger = get_logger(__name__)
_CHAT_PATH = " /api/chat "
_TAGS_PATH = " /api/tags "
# Per-chunk read timeout for streaming; model can take minutes to start tokens on CPU.
_STREAM_CHUNK_TIMEOUT = 600.0
class OllamaClient :
""" Async client for Ollama REST API. """
""" Async client for Ollama REST API.
Uses synchronous httpx inside asyncio . to_thread ( ) to fully isolate Ollama
HTTP calls from the asyncio event loop , preventing hangs when called from
within an asyncpg / SQLAlchemy session context on Python 3.13 .
"""
def __init__ ( self , base_url : str , model : str , timeout : float = 60.0 ) - > None :
self . _base_url = base_url . rstrip ( " / " )
self . model = model
self . _timeout = timeout
self . _client : httpx . AsyncClient | None = None
async def __aenter__ ( self ) - > OllamaClient :
self . _client = httpx . AsyncClient ( base_url = self . _base_url , timeout = self . _timeout )
return self
async def __aexit__ ( self , * _ : object ) - > None :
if self . _client is not None :
await self . _client . aclose ( )
self . _client = None
pass
@with_retry ( max_attempts = 2 , min_wait = 0.5 , max_wait = 10.0 , multiplier = 2.0 )
async def chat (
def _sync_call (
self ,
messages : list [ dict [ str , str ] ] ,
response_format : str = " json " ,
temperature : float = 0.0 ,
response_format : str ,
temperature : float ,
) - > tuple [ dict [ str , Any ] , dict [ str , int ] , int ] :
""" Call Ollama /api/chat and return (parsed_json, token_usage, elapsed_ms) .
""" Synchronous HTTP call to Ollama /api/chat .
Args :
messages : List of { role , content } message dicts .
response_format : " json " forces JSON output mode .
temperature : Sampling temperature ( 0.0 = deterministic ) .
Runs inside asyncio . to_thread ( ) so it does not block the event loop .
Returns :
Tuple of ( parsed response dict , token usage dict , elapsed_ms ) .
Raises :
LLMTimeoutError : On request timeout .
LLMError : On non - retryable HTTP errors .
RetryableError : On 5 xx server errors .
LLMError : On non - retryable HTTP errors or bad JSON .
RetryableError : On 5 xx server errors or connection failures .
"""
if self . _client is None :
raise LLMError ( " OllamaClient must be used as an async context manager " )
payload : dict [ str , Any ] = {
" model " : self . model ,
" messages " : messages ,
" stream " : False ,
" options " : { " temperature " : temperature } ,
" stream " : True ,
# Disable Qwen3/Qwen3.5 extended thinking mode — thinking tokens add
# 30-60s latency with no benefit for structured JSON extraction tasks.
" think " : False ,
" options " : {
" temperature " : temperature ,
# Limit KV cache to actual prompt size; default 262144 for Qwen3.5
# allocates 20 GB of VRAM which slows cold-start significantly.
" num_ctx " : 8192 ,
} ,
}
if response_format == " json " :
payload [ " format " ] = " json "
stream_timeout = httpx . Timeout (
connect = 10.0 ,
read = _STREAM_CHUNK_TIMEOUT ,
write = 30.0 ,
pool = 10.0 ,
)
t0 = time . monotonic ( )
try :
response = await self . _client . post ( _CHAT_PATH , json = payload )
chunks : list [ str ] = [ ]
prompt_tokens = 0
completion_tokens = 0
with httpx . Client ( base_url = self . _base_url , timeout = stream_timeout ) as http :
with http . stream ( " POST " , _CHAT_PATH , json = payload ) as response :
if response . status_code > = 500 :
body = response . read ( )
raise RetryableError (
f " Ollama server error { response . status_code } " ,
source = " ollama " ,
context = { " status " : response . status_code , " body " : body [ : 200 ] . decode ( ) } ,
)
if response . status_code > = 400 :
body = response . read ( )
raise LLMError (
f " Ollama client error { response . status_code } : { body [ : 200 ] . decode ( ) } " ,
source = " ollama " ,
context = { " status " : response . status_code } ,
)
for line in response . iter_lines ( ) :
if not line . strip ( ) :
continue
try :
chunk = json . loads ( line )
except json . JSONDecodeError :
continue
content = chunk . get ( " message " , { } ) . get ( " content " , " " )
if content :
chunks . append ( content )
if chunk . get ( " done " ) :
prompt_tokens = chunk . get ( " prompt_eval_count " , 0 )
completion_tokens = chunk . get ( " eval_count " , 0 )
break
except ( LLMTimeoutError , LLMError , RetryableError ) :
raise
except httpx . TimeoutException as exc :
raise LLMTimeoutError (
f " Ollama request timed out after { self . _timeout } s " ,
f " Ollama request timed out after { _STREAM_CHUNK_TIMEOUT } s " ,
source = " ollama " ,
context = { " model " : self . model } ,
) from exc
@ -86,25 +132,10 @@ class OllamaClient:
) from exc
elapsed_ms = int ( ( time . monotonic ( ) - t0 ) * 1000 )
if response . status_code > = 500 :
raise RetryableError (
f " Ollama server error { response . status_code } " ,
source = " ollama " ,
context = { " status " : response . status_code , " body " : response . text [ : 200 ] } ,
)
if response . status_code > = 400 :
raise LLMError (
f " Ollama client error { response . status_code } : { response . text [ : 200 ] } " ,
source = " ollama " ,
context = { " status " : response . status_code } ,
)
data = response . json ( )
raw_content : str = data . get ( " message " , { } ) . get ( " content " , " " )
raw_content = " " . join ( chunks )
token_usage = {
" prompt_tokens " : data. get ( " prompt_eval_count " , 0 ) ,
" completion_tokens " : data. get ( " eval_count " , 0 ) ,
" prompt_tokens " : prompt_tokens ,
" completion_tokens " : completion_tokens ,
}
try :
@ -124,12 +155,48 @@ class OllamaClient:
)
return parsed , token_usage , elapsed_ms
@with_retry ( max_attempts = 2 , min_wait = 0.5 , max_wait = 10.0 , multiplier = 2.0 )
async def chat (
self ,
messages : list [ dict [ str , str ] ] ,
response_format : str = " json " ,
temperature : float = 0.0 ,
) - > tuple [ dict [ str , Any ] , dict [ str , int ] , int ] :
""" Call Ollama /api/chat and return (parsed_json, token_usage, elapsed_ms).
Args :
messages : List of { role , content } message dicts .
response_format : " json " forces JSON output mode .
temperature : Sampling temperature ( 0.0 = deterministic ) .
Returns :
Tuple of ( parsed response dict , token usage dict , elapsed_ms ) .
Raises :
LLMTimeoutError : On request timeout .
LLMError : On non - retryable HTTP errors .
RetryableError : On 5 xx server errors .
"""
try :
return await asyncio . to_thread (
self . _sync_call , messages , response_format , temperature
)
except ( LLMTimeoutError , LLMError , RetryableError ) :
raise
except Exception as exc :
raise LLMError (
f " Unexpected error calling Ollama: { exc } " ,
source = " ollama " ,
context = { " model " : self . model } ,
) from exc
async def health_check ( self ) - > bool :
""" Return True if Ollama is reachable and the model is available. """
if self . _client is None :
raise LLMError ( " OllamaClient must be used as an async context manager " )
def _sync_health ( ) - > bool :
try :
response = await self . _client . get ( _TAGS_PATH )
with httpx . Client ( base_url = self . _base_url , timeout = 10.0 ) as http :
response = http . get ( _TAGS_PATH )
if response . status_code != 200 :
return False
data = response . json ( )
@ -138,6 +205,8 @@ class OllamaClient:
except Exception :
return False
return await asyncio . to_thread ( _sync_health )
def make_ollama_client ( ) - > OllamaClient :
""" Factory that reads config from settings. """