mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 00:16:23 +02:00
283 lines
10 KiB
Python
283 lines
10 KiB
Python
|
|
"""
|
|
Simple LLM service, performs text prompt completion using GoogleAIStudio.
|
|
Input is prompt, output is response.
|
|
"""
|
|
|
|
#
|
|
# Using this SDK:
|
|
# https://googleapis.github.io/python-genai/genai.html#module-genai.client
|
|
#
|
|
# Seems to have simpler dependencies on the 'VertexAI' service, which
|
|
# TrustGraph implements in the trustgraph-vertexai package.
|
|
#
|
|
|
|
import os
|
|
import logging
|
|
|
|
# Module logger
|
|
logger = logging.getLogger(__name__)
|
|
|
|
from .... exceptions import TooManyRequests
|
|
from .... base import LlmService, LlmResult, LlmChunk
|
|
|
|
default_ident = "text-completion"
|
|
|
|
default_model = 'gemini-2.0-flash-001'
|
|
default_temperature = 0.0
|
|
default_max_output = 8192
|
|
default_api_key = os.getenv("GOOGLE_AI_STUDIO_KEY")
|
|
|
|
class Processor(LlmService):
|
|
|
|
def __init__(self, **params):
|
|
|
|
model = params.get("model", default_model)
|
|
api_key = params.get("api_key", default_api_key)
|
|
temperature = params.get("temperature", default_temperature)
|
|
max_output = params.get("max_output", default_max_output)
|
|
|
|
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
|
|
self.genai = genai
|
|
self.types = types
|
|
self.HarmCategory = HarmCategory
|
|
self.HarmBlockThreshold = HarmBlockThreshold
|
|
self.ClientError = ClientError
|
|
self.ResourceExhausted = ResourceExhausted
|
|
|
|
if api_key is None:
|
|
raise RuntimeError("Google AI Studio API key not specified")
|
|
|
|
super(Processor, self).__init__(
|
|
**params | {
|
|
"model": model,
|
|
"temperature": temperature,
|
|
"max_output": max_output,
|
|
}
|
|
)
|
|
|
|
self.client = self.genai.Client(api_key=api_key, vertexai=False)
|
|
self.default_model = model
|
|
self.temperature = temperature
|
|
self.max_output = max_output
|
|
|
|
# Cache for generation configs per model
|
|
self.generation_configs = {}
|
|
|
|
block_level = self.HarmBlockThreshold.BLOCK_ONLY_HIGH
|
|
|
|
self.safety_settings = [
|
|
self.types.SafetySetting(
|
|
category = self.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
|
|
threshold = block_level,
|
|
),
|
|
self.types.SafetySetting(
|
|
category = self.HarmCategory.HARM_CATEGORY_HARASSMENT,
|
|
threshold = block_level,
|
|
),
|
|
self.types.SafetySetting(
|
|
category = self.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
|
|
threshold = block_level,
|
|
),
|
|
self.types.SafetySetting(
|
|
category = self.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
|
|
threshold = block_level,
|
|
),
|
|
# There is a documentation conflict on whether or not
|
|
# CIVIC_INTEGRITY is a valid category
|
|
# HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY: block_level,
|
|
]
|
|
|
|
logger.info("GoogleAIStudio LLM service initialized")
|
|
|
|
def _get_or_create_config(self, model_name, temperature=None):
|
|
"""Get or create generation config with dynamic temperature"""
|
|
# Use provided temperature or fall back to default
|
|
effective_temperature = temperature if temperature is not None else self.temperature
|
|
|
|
# Create cache key that includes temperature to avoid conflicts
|
|
cache_key = f"{model_name}:{effective_temperature}"
|
|
|
|
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(
|
|
temperature = effective_temperature,
|
|
top_p = 1,
|
|
top_k = 40,
|
|
max_output_tokens = self.max_output,
|
|
response_mime_type = "text/plain",
|
|
safety_settings = self.safety_settings,
|
|
)
|
|
|
|
return self.generation_configs[cache_key]
|
|
|
|
async def generate_content(self, system, prompt, model=None, temperature=None):
|
|
|
|
# Use provided model or fall back to default
|
|
model_name = model or self.default_model
|
|
# Use provided temperature or fall back to default
|
|
effective_temperature = temperature if temperature is not None else self.temperature
|
|
|
|
logger.debug(f"Using model: {model_name}")
|
|
logger.debug(f"Using temperature: {effective_temperature}")
|
|
|
|
generation_config = self._get_or_create_config(model_name, effective_temperature)
|
|
# Set system instruction per request (can't be cached)
|
|
generation_config.system_instruction = system
|
|
|
|
try:
|
|
|
|
response = self.client.models.generate_content(
|
|
model=model_name,
|
|
config=generation_config,
|
|
contents=prompt,
|
|
)
|
|
|
|
resp = response.text
|
|
inputtokens = int(response.usage_metadata.prompt_token_count)
|
|
outputtokens = int(response.usage_metadata.candidates_token_count)
|
|
logger.debug(f"LLM response: {resp}")
|
|
logger.info(f"Input Tokens: {inputtokens}")
|
|
logger.info(f"Output Tokens: {outputtokens}")
|
|
|
|
resp = LlmResult(
|
|
text = resp,
|
|
in_token = inputtokens,
|
|
out_token = outputtokens,
|
|
model = model_name
|
|
)
|
|
|
|
return resp
|
|
|
|
except self.ResourceExhausted as e:
|
|
|
|
logger.warning("Rate limit exceeded")
|
|
|
|
# Leave rate limit retries to the default handler
|
|
raise TooManyRequests()
|
|
|
|
except self.ClientError as e:
|
|
# google-genai SDK throws ClientError for 4xx errors
|
|
if e.code == 429:
|
|
logger.warning(f"Rate limit exceeded (ClientError 429): {e}")
|
|
raise TooManyRequests()
|
|
# Other client errors are unrecoverable
|
|
logger.error(f"GoogleAIStudio ClientError: {e}", exc_info=True)
|
|
raise e
|
|
|
|
except Exception as e:
|
|
|
|
# Apart from rate limits, treat all exceptions as unrecoverable
|
|
|
|
logger.error(f"GoogleAIStudio LLM exception ({type(e).__name__}): {e}", exc_info=True)
|
|
raise e
|
|
|
|
def supports_streaming(self):
|
|
"""Google AI Studio supports streaming"""
|
|
return True
|
|
|
|
async def generate_content_stream(self, system, prompt, model=None, temperature=None):
|
|
"""Stream content generation from Google AI Studio"""
|
|
model_name = model or self.default_model
|
|
effective_temperature = temperature if temperature is not None else self.temperature
|
|
|
|
logger.debug(f"Using model (streaming): {model_name}")
|
|
logger.debug(f"Using temperature: {effective_temperature}")
|
|
|
|
generation_config = self._get_or_create_config(model_name, effective_temperature)
|
|
generation_config.system_instruction = system
|
|
|
|
try:
|
|
response = self.client.models.generate_content_stream(
|
|
model=model_name,
|
|
config=generation_config,
|
|
contents=prompt,
|
|
)
|
|
|
|
total_input_tokens = 0
|
|
total_output_tokens = 0
|
|
|
|
for chunk in response:
|
|
if hasattr(chunk, 'text') and chunk.text:
|
|
yield LlmChunk(
|
|
text=chunk.text,
|
|
in_token=None,
|
|
out_token=None,
|
|
model=model_name,
|
|
is_final=False
|
|
)
|
|
|
|
# Accumulate token counts if available
|
|
if hasattr(chunk, 'usage_metadata'):
|
|
if hasattr(chunk.usage_metadata, 'prompt_token_count'):
|
|
total_input_tokens = int(chunk.usage_metadata.prompt_token_count)
|
|
if hasattr(chunk.usage_metadata, 'candidates_token_count'):
|
|
total_output_tokens = int(chunk.usage_metadata.candidates_token_count)
|
|
|
|
# Send final chunk with token counts
|
|
yield LlmChunk(
|
|
text="",
|
|
in_token=total_input_tokens,
|
|
out_token=total_output_tokens,
|
|
model=model_name,
|
|
is_final=True
|
|
)
|
|
|
|
logger.debug("Streaming complete")
|
|
|
|
except self.ResourceExhausted:
|
|
logger.warning("Rate limit exceeded during streaming")
|
|
raise TooManyRequests()
|
|
|
|
except self.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}")
|
|
raise TooManyRequests()
|
|
# Other client errors are unrecoverable
|
|
logger.error(f"GoogleAIStudio streaming ClientError: {e}", exc_info=True)
|
|
raise e
|
|
|
|
except Exception as e:
|
|
logger.error(f"GoogleAIStudio streaming exception ({type(e).__name__}): {e}", exc_info=True)
|
|
raise e
|
|
|
|
@staticmethod
|
|
def add_args(parser):
|
|
|
|
LlmService.add_args(parser)
|
|
|
|
parser.add_argument(
|
|
'-m', '--model',
|
|
default=default_model,
|
|
help=f'LLM model (default: {default_model})'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-k', '--api-key',
|
|
default=default_api_key,
|
|
help=f'GoogleAIStudio API key'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-t', '--temperature',
|
|
type=float,
|
|
default=default_temperature,
|
|
help=f'LLM temperature parameter (default: {default_temperature})'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-x', '--max-output',
|
|
type=int,
|
|
default=default_max_output,
|
|
help=f'LLM max output tokens (default: {default_max_output})'
|
|
)
|
|
|
|
def run():
|
|
|
|
Processor.launch(default_ident, __doc__)
|
|
|