services(llm): add on_completion_timeout event

This commit is contained in:
Aleix Conchillo Flaqué
2025-02-24 14:23:03 -08:00
parent fc544fa61c
commit 2110b79507
5 changed files with 25 additions and 7 deletions

View File

@@ -142,6 +142,8 @@ class LLMService(AIService):
self._callbacks = {}
self._start_callbacks = {}
self._register_event_handler("on_completion_timeout")
# TODO-CB: callback function type
def register_function(self, function_name: Optional[str], callback, start_callback=None):
# Registering a function with the function_name set to None will run that callback

View File

@@ -4,15 +4,16 @@
# SPDX-License-Identifier: BSD 2-Clause License
#
import asyncio
import base64
import copy
import io
import json
import re
from asyncio import CancelledError
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union
import httpx
from loguru import logger
from PIL import Image
from pydantic import BaseModel, Field
@@ -251,12 +252,14 @@ class AnthropicLLMService(LLMService):
if total_input_tokens >= 1024:
context.turns_above_cache_threshold += 1
except CancelledError:
except asyncio.CancelledError:
# If we're interrupted, we won't get a complete usage report. So set our flag to use the
# token estimate. The reraise the exception so all the processors running in this task
# also get cancelled.
use_completion_tokens_estimate = True
raise
except httpx.TimeoutException:
await self._call_event_handler("on_completion_timeout")
except Exception as e:
logger.exception(f"{self} exception: {e}")
finally:

View File

@@ -11,6 +11,8 @@ import json
import os
import time
from google.api_core.exceptions import DeadlineExceeded
# Suppress gRPC fork warnings
os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false"
@@ -1126,6 +1128,8 @@ class GoogleLLMService(LLMService):
else:
logger.exception(f"{self} error: {e}")
except DeadlineExceeded:
await self._call_event_handler("on_completion_timeout")
except Exception as e:
logger.exception(f"{self} exception: {e}")
finally:

View File

@@ -310,11 +310,15 @@ class BaseOpenAILLMService(LLMService):
await self.push_frame(frame, direction)
if context:
await self.push_frame(LLMFullResponseStartFrame())
await self.start_processing_metrics()
await self._process_context(context)
await self.stop_processing_metrics()
await self.push_frame(LLMFullResponseEndFrame())
try:
await self.push_frame(LLMFullResponseStartFrame())
await self.start_processing_metrics()
await self._process_context(context)
except httpx.TimeoutException:
await self._call_event_handler("on_completion_timeout")
finally:
await self.stop_processing_metrics()
await self.push_frame(LLMFullResponseEndFrame())
@dataclass