Track live provider instances and update count on success or failure in various providers

This commit is contained in:
hlohaus
2025-09-06 11:28:09 +02:00
parent 6a6a25fba8
commit 1cbf9dfb6a
5 changed files with 55 additions and 37 deletions

View File

@@ -49,6 +49,8 @@ class OpenaiTemplate(AsyncGeneratorProvider, ProviderModelMixin, RaiseErrorMixin
raise_for_status(response)
data = response.json()
data = data.get("data") if isinstance(data, dict) else data
if data:
cls.live += 1
cls.image_models = [model.get("name") if cls.use_model_names else model.get("id", model.get("name")) for model in data if model.get("image") or model.get("type") == "image" or model.get("supports_images")]
cls.vision_models = cls.vision_models.copy()
cls.vision_models += [model.get("name") if cls.use_model_names else model.get("id", model.get("name")) for model in data if model.get("vision")]

View File

@@ -108,6 +108,7 @@ class Api:
"active_by_default": False if provider.active_by_default is None else provider.active_by_default,
"auth": provider.needs_auth,
"login_url": getattr(provider, "login_url", None),
"live": provider.live
} for provider in Provider.__providers__ if provider.working and safe_get_models(provider)]
def get_all_models(self) -> dict[str, list]:

View File

@@ -99,9 +99,11 @@ class RotatedProvider(BaseRetryProvider):
if is_content(chunk):
started = True
if started:
provider.live += 1
# Success, so we return and do not rotate
return
except Exception as e:
provider.live -= 1
exceptions[provider.__name__] = e
debug.error(f"{provider.__name__} failed: {e}")
@@ -161,8 +163,10 @@ class RotatedProvider(BaseRetryProvider):
if is_content(chunk):
started = True
if started:
provider.live += 1
return # Success
except Exception as e:
provider.live -= 1
exceptions[provider.__name__] = e
debug.error(f"{provider.__name__} failed: {e}")

View File

@@ -28,6 +28,7 @@ class BaseProvider(ABC):
params: str
create_function: callable
async_create_function: callable
live: int = 0
@classmethod
def get_dict(cls) -> Dict[str, str]:

View File

@@ -253,18 +253,23 @@ async def async_iter_run_tools(
# Generate response
response = to_async_iterator(provider.async_create_function(model=model, messages=messages, **kwargs))
model_info = model
async for chunk in response:
if isinstance(chunk, ProviderInfo):
model_info = getattr(chunk, 'model', model_info)
elif isinstance(chunk, Usage):
usage = {"user": kwargs.get("user"), "model": model_info, "provider": provider.get_parent(), **chunk.get_dict()}
usage_dir = Path(get_cookies_dir()) / ".usage"
usage_file = usage_dir / f"{datetime.date.today()}.jsonl"
usage_dir.mkdir(parents=True, exist_ok=True)
with usage_file.open("a" if usage_file.exists() else "w") as f:
f.write(f"{json.dumps(usage)}\n")
yield chunk
try:
model_info = model
async for chunk in response:
if isinstance(chunk, ProviderInfo):
model_info = getattr(chunk, 'model', model_info)
elif isinstance(chunk, Usage):
usage = {"user": kwargs.get("user"), "model": model_info, "provider": provider.get_parent(), **chunk.get_dict()}
usage_dir = Path(get_cookies_dir()) / ".usage"
usage_file = usage_dir / f"{datetime.date.today()}.jsonl"
usage_dir.mkdir(parents=True, exist_ok=True)
with usage_file.open("a" if usage_file.exists() else "w") as f:
f.write(f"{json.dumps(usage)}\n")
yield chunk
provider.live += 1
except:
provider.live -= 1
raise
# Yield sources if available
if sources:
@@ -338,33 +343,38 @@ def iter_run_tools(
thinking_start_time = 0
processor = ThinkingProcessor()
model_info = model
for chunk in provider.create_function(model=model, messages=messages, provider=provider, **kwargs):
if isinstance(chunk, FinishReason):
if sources is not None:
yield sources
try:
for chunk in provider.create_function(model=model, messages=messages, provider=provider, **kwargs):
if isinstance(chunk, FinishReason):
if sources is not None:
yield sources
sources = None
yield chunk
continue
elif isinstance(chunk, Sources):
sources = None
yield chunk
continue
elif isinstance(chunk, Sources):
sources = None
elif isinstance(chunk, ProviderInfo):
model_info = getattr(chunk, 'model', model_info)
elif isinstance(chunk, Usage):
usage = {"user": kwargs.get("user"), "model": model_info, "provider": provider.get_parent(), **chunk.get_dict()}
usage_dir = Path(get_cookies_dir()) / ".usage"
usage_file = usage_dir / f"{datetime.date.today()}.jsonl"
usage_dir.mkdir(parents=True, exist_ok=True)
with usage_file.open("a" if usage_file.exists() else "w") as f:
f.write(f"{json.dumps(usage)}\n")
if not isinstance(chunk, str):
yield chunk
continue
elif isinstance(chunk, ProviderInfo):
model_info = getattr(chunk, 'model', model_info)
elif isinstance(chunk, Usage):
usage = {"user": kwargs.get("user"), "model": model_info, "provider": provider.get_parent(), **chunk.get_dict()}
usage_dir = Path(get_cookies_dir()) / ".usage"
usage_file = usage_dir / f"{datetime.date.today()}.jsonl"
usage_dir.mkdir(parents=True, exist_ok=True)
with usage_file.open("a" if usage_file.exists() else "w") as f:
f.write(f"{json.dumps(usage)}\n")
if not isinstance(chunk, str):
yield chunk
continue
thinking_start_time, results = processor.process_thinking_chunk(chunk, thinking_start_time)
thinking_start_time, results = processor.process_thinking_chunk(chunk, thinking_start_time)
for result in results:
yield result
for result in results:
yield result
provider.live += 1
except:
provider.live -= 1
raise
if sources is not None:
yield sources