mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 00:16:23 +02:00
Merge branch 'release/v2.3'
This commit is contained in:
commit
222537c26b
18 changed files with 1020 additions and 247 deletions
|
|
@ -18,6 +18,12 @@ import logging
|
|||
# Module logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from google.genai.types import HarmCategory, HarmBlockThreshold
|
||||
from google.genai.errors import ClientError
|
||||
from google.api_core.exceptions import ResourceExhausted
|
||||
|
||||
from .... exceptions import TooManyRequests
|
||||
from .... base import LlmService, LlmResult, LlmChunk
|
||||
|
||||
|
|
@ -71,20 +77,20 @@ class Processor(LlmService):
|
|||
block_level = self.HarmBlockThreshold.BLOCK_ONLY_HIGH
|
||||
|
||||
self.safety_settings = [
|
||||
self.types.SafetySetting(
|
||||
category = self.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
|
||||
types.SafetySetting(
|
||||
category = HarmCategory.HARM_CATEGORY_HATE_SPEECH,
|
||||
threshold = block_level,
|
||||
),
|
||||
self.types.SafetySetting(
|
||||
category = self.HarmCategory.HARM_CATEGORY_HARASSMENT,
|
||||
types.SafetySetting(
|
||||
category = HarmCategory.HARM_CATEGORY_HARASSMENT,
|
||||
threshold = block_level,
|
||||
),
|
||||
self.types.SafetySetting(
|
||||
category = self.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
|
||||
types.SafetySetting(
|
||||
category = HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
|
||||
threshold = block_level,
|
||||
),
|
||||
self.types.SafetySetting(
|
||||
category = self.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
|
||||
types.SafetySetting(
|
||||
category = HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
|
||||
threshold = block_level,
|
||||
),
|
||||
# There is a documentation conflict on whether or not
|
||||
|
|
@ -104,7 +110,7 @@ class Processor(LlmService):
|
|||
|
||||
if cache_key not in self.generation_configs:
|
||||
logger.info(f"Creating generation config for '{model_name}' with temperature {effective_temperature}")
|
||||
self.generation_configs[cache_key] = self.types.GenerateContentConfig(
|
||||
self.generation_configs[cache_key] = types.GenerateContentConfig(
|
||||
temperature = effective_temperature,
|
||||
top_p = 1,
|
||||
top_k = 40,
|
||||
|
|
@ -160,7 +166,7 @@ class Processor(LlmService):
|
|||
# Leave rate limit retries to the default handler
|
||||
raise TooManyRequests()
|
||||
|
||||
except self.ClientError as e:
|
||||
except ClientError as e:
|
||||
# google-genai SDK throws ClientError for 4xx errors
|
||||
if e.code == 429:
|
||||
logger.warning(f"Rate limit exceeded (ClientError 429): {e}")
|
||||
|
|
@ -229,11 +235,11 @@ class Processor(LlmService):
|
|||
|
||||
logger.debug("Streaming complete")
|
||||
|
||||
except self.ResourceExhausted:
|
||||
except ResourceExhausted:
|
||||
logger.warning("Rate limit exceeded during streaming")
|
||||
raise TooManyRequests()
|
||||
|
||||
except self.ClientError as e:
|
||||
except ClientError as e:
|
||||
# google-genai SDK throws ClientError for 4xx errors
|
||||
if e.code == 429:
|
||||
logger.warning(f"Rate limit exceeded during streaming (ClientError 429): {e}")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue