Spaces:
Running
Running
File size: 10,602 Bytes
a55ea57 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 | """
File: /client_small.py
Created Date: Sunday December 21st 2025
Author: Christian Nonis <alch.infoemail@gmail.com>
-----
Last Modified: Thursday February 19th 2026 7:45:12 pm
Modified By: Christian Nonis <alch.infoemail@gmail.com>
-----
"""
import json
import os
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError
import threading
os.environ["GRPC_DNS_RESOLVER"] = "native"
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
RetryError,
)
from src.adapters.interfaces.llm import LLM
from src.config import config
_parent_pid = os.getpid()
class LLMClientSmall(LLM):
"""
Large language model client, used for main agents and other high-level tasks.
"""
def __init__(self):
self._client = None
self._model = None
self._langchain_model = None
self._lock = threading.Lock()
self._pid = os.getpid()
self.default_timeout = 120
def _check_fork(self):
if os.getpid() != self._pid:
self._client = None
self._langchain_model = None
self._pid = os.getpid()
@property
def client(self):
self._check_fork()
if self._client is None:
with self._lock:
if self._client is None:
from google import genai
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = (
config.gcp.credentials_path
)
self._client = genai.Client(
vertexai=True,
project=config.gcp.project_id,
location="global",
http_options={"api_version": "v1"},
)
return self._client
@property
def model(self):
if self._model is None:
self._model = config.gcp.small_llm_model
return self._model
@property
def langchain_model(self):
"""
Lazily initialize (thread-safe) and return a LangChain ChatVertexAI wrapper configured for Vertex AI.
This property constructs the ChatVertexAI instance on first access, sets the
GOOGLE_APPLICATION_CREDENTIALS environment variable from config.gcp.credentials_path,
and caches the resulting model for subsequent calls. Initialization is guarded
with a lock to be safe across threads.
Returns:
ChatVertexAI: The initialized LangChain ChatVertexAI model instance.
"""
if self._langchain_model is None:
with self._lock:
if self._langchain_model is None:
from langchain_google_vertexai import ChatVertexAI
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = (
config.gcp.credentials_path
)
kwargs = {
"model_name": self.model,
"project": config.gcp.project_id,
"max_retries": 5,
"location": "global",
"request_timeout": 120,
"streaming": False,
"wait_exponential_kwargs": {
"multiplier": 1,
"min": 2,
"max": 60,
},
}
kwargs["max_output_tokens"] = 65000
kwargs["thinking_budget"] = 0
self._langchain_model = ChatVertexAI(**kwargs)
return self._langchain_model
def generate_text(
self, prompt: str, max_new_tokens: int = 65536, timeout: int = None
) -> str:
"""
Generate plain-text completion for a prompt using the configured Vertex AI model.
If `prompt` is empty or only whitespace, returns the fixed message "Input prompt is empty". The method will retry transient failures up to 3 times and enforce the provided `timeout` per attempt.
Parameters:
prompt (str): The input prompt to complete.
max_new_tokens (int): Maximum number of tokens to generate for the completion. If falsy, no token limit is passed to the API.
timeout (int): Maximum seconds to wait for a single generation attempt; defaults to the client's default timeout.
Returns:
str: The generated plain-text completion, or "Input prompt is empty" when the prompt is empty.
Raises:
TimeoutError: If generation times out or all retry attempts fail.
"""
if not prompt or len(prompt.strip()) == 0:
return "Input prompt is empty"
timeout = timeout or self.default_timeout
from google.genai.types import GenerateContentConfig
def _generate():
"""
Call the LLM client's generate_content API for the current model and request a plain-text response.
If `max_new_tokens` is set in the enclosing scope, the request will include `max_output_tokens` with that value.
Returns:
The raw response from the client's `generate_content` call containing the generated plain-text output.
"""
return self.client.models.generate_content(
model=self.model,
contents=[prompt],
config=GenerateContentConfig(
response_mime_type="text/plain",
**({"max_output_tokens": max_new_tokens} if max_new_tokens else {}),
),
)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=2, max=30),
retry=retry_if_exception_type((TimeoutError, Exception)),
reraise=True,
)
def _generate_with_retry():
executor = ThreadPoolExecutor(max_workers=1)
future = None
try:
future = executor.submit(_generate)
response = future.result(timeout=timeout)
return response.text
except FutureTimeoutError:
if future:
future.cancel()
raise TimeoutError(
f"LLM generate_text timed out after {timeout} seconds. "
"This may indicate a network issue or the LLM service is unresponsive."
)
finally:
executor.shutdown(wait=False, cancel_futures=True)
try:
return _generate_with_retry()
except RetryError as e:
last_attempt = e.last_attempt
raise TimeoutError(
f"LLM generate_text failed after {last_attempt.attempt_number} attempts. "
f"Last error: {last_attempt.exception()}"
) from last_attempt.exception()
except TimeoutError:
raise
def generate_json(
self,
prompt: str,
max_new_tokens: int = 65536,
max_retries: int = 3,
timeout: int = None,
) -> dict:
"""
Generate JSON content from a given prompt using the configured language model.
Attempts generation up to a specified number of retries with exponential backoff and enforces a timeout for each attempt. Parses the generated text as JSON and returns it as a dictionary.
Parameters:
max_new_tokens (int): Maximum number of tokens to generate in the response.
max_retries (int): Number of retry attempts for generation on failure.
timeout (int): Maximum time in seconds to wait for each generation attempt.
Returns:
dict: The generated content parsed as a JSON dictionary.
"""
timeout = timeout or self.default_timeout
from google.genai.types import GenerateContentConfig
try:
from google.genai.types import ThinkingConfig
_thinking_config = ThinkingConfig(thinking_budget=0)
except ImportError:
_thinking_config = None
def _generate():
"""
Generate JSON content from the model based on the given prompt.
Returns:
dict: The parsed JSON response generated from the prompt.
"""
config_kw = {
"response_mime_type": "application/json",
**({"max_output_tokens": max_new_tokens} if max_new_tokens else {}),
}
if _thinking_config is not None:
config_kw["thinking_config"] = _thinking_config
_response = self.client.models.generate_content(
model=self.model,
contents=[prompt],
config=GenerateContentConfig(**config_kw),
)
_response = _response.text.strip("").strip("```")
return json.loads(_response)
@retry(
stop=stop_after_attempt(max_retries),
wait=wait_exponential(multiplier=1, min=2, max=30),
retry=retry_if_exception_type(
(TimeoutError, json.JSONDecodeError, Exception)
),
reraise=True,
)
def _generate_with_retry():
executor = ThreadPoolExecutor(max_workers=1)
future = None
try:
future = executor.submit(_generate)
response = future.result(timeout=timeout)
return response
except FutureTimeoutError:
if future:
future.cancel()
raise TimeoutError(
f"LLM generate_json timed out after {timeout} seconds. "
"This may indicate a network issue or the LLM service is unresponsive."
)
except json.JSONDecodeError as e:
raise json.JSONDecodeError(f"Failed to parse JSON: {e}", e.doc, e.pos)
finally:
executor.shutdown(wait=False, cancel_futures=True)
try:
return _generate_with_retry()
except RetryError as e:
last_attempt = e.last_attempt
raise TimeoutError(
f"LLM generate_json failed after {last_attempt.attempt_number} attempts. "
f"Last error: {last_attempt.exception()}"
) from last_attempt.exception()
except (TimeoutError, json.JSONDecodeError):
raise
_llm_small_client = LLMClientSmall() |