wehub-resource-sync
792 lines
25 KiB
Python
792 lines
25 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
ABOUTME: Shared token estimation utilities for audit scripts
|
|
ABOUTME: XML sanitization helpers for document processing
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
|
|
try:
|
|
from google import genai
|
|
from google.genai import types
|
|
|
|
HAS_GEMINI = True
|
|
except ImportError: # pragma: no cover - optional dependency
|
|
genai = None
|
|
types = None
|
|
HAS_GEMINI = False
|
|
|
|
try:
|
|
import openai
|
|
|
|
HAS_OPENAI = True
|
|
except ImportError: # pragma: no cover - optional dependency
|
|
openai = None
|
|
HAS_OPENAI = False
|
|
|
|
|
|
def estimate_tokens(text: str) -> int:
|
|
"""
|
|
Estimate token count for LLM context management.
|
|
|
|
Uses a weighted formula based on character types:
|
|
- Chinese characters: ~0.75 tokens per character (subword tokenization)
|
|
- JSON structural characters (brackets, quotes, commas): ~1 tokens per character
|
|
- Other characters (English, numbers, symbols): ~0.4 tokens per character (~3 chars/token)
|
|
|
|
Includes 5% buffer and safety offset for special formatting and system prompt overhead.
|
|
|
|
Args:
|
|
text: Input text to estimate tokens for
|
|
|
|
Returns:
|
|
int: Estimated token count
|
|
"""
|
|
if not text:
|
|
return 0
|
|
|
|
chinese_count = len(re.findall(r"[\u4e00-\u9fa5]", text))
|
|
json_chars_count = len(re.findall(r'[\[\]",{}]', text))
|
|
other_count = len(text) - chinese_count - json_chars_count
|
|
|
|
base_estimate = (
|
|
(chinese_count * 0.75) + (json_chars_count * 1) + (other_count * 0.4)
|
|
)
|
|
final_tokens = int(base_estimate * 1.05) + 2
|
|
return final_tokens
|
|
|
|
|
|
def sanitize_xml_string(text: str) -> str:
|
|
"""
|
|
Remove control characters that are illegal in XML 1.0.
|
|
|
|
XML 1.0 allows: #x9 (tab), #xA (LF), #xD (CR), and #x20-#xD7FF, #xE000-#xFFFD, #x10000-#x10FFFF
|
|
This function removes all other control characters (0x00-0x08, 0x0B, 0x0C, 0x0E-0x1F).
|
|
|
|
Args:
|
|
text: Text that may contain control characters
|
|
|
|
Returns:
|
|
Sanitized text safe for XML. Returns input unchanged if not a non-empty string.
|
|
"""
|
|
if not text or not isinstance(text, str):
|
|
return text
|
|
# Build a translation table to remove illegal control characters
|
|
# Keep: \t (0x09), \n (0x0A), \r (0x0D)
|
|
# Remove: 0x00-0x08, 0x0B, 0x0C, 0x0E-0x1F
|
|
illegal_chars = "".join(chr(c) for c in range(0x20) if c not in (0x09, 0x0A, 0x0D))
|
|
return text.translate(str.maketrans("", "", illegal_chars))
|
|
|
|
|
|
def is_vertex_ai_mode() -> bool:
|
|
"""
|
|
Check if Vertex AI mode is enabled via environment variable.
|
|
|
|
Returns:
|
|
True if GOOGLE_GENAI_USE_VERTEXAI is set to 'true', False otherwise
|
|
"""
|
|
return os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true"
|
|
|
|
|
|
def create_gemini_client(use_async: bool = False):
|
|
"""
|
|
Create Gemini client for AI Studio or Vertex AI.
|
|
|
|
Supports two modes:
|
|
- AI Studio (default): Uses GOOGLE_API_KEY for authentication
|
|
- Vertex AI: Uses ADC (GOOGLE_APPLICATION_CREDENTIALS or gcloud auth)
|
|
|
|
Environment variables for Vertex AI mode:
|
|
- GOOGLE_GENAI_USE_VERTEXAI: Set to 'true' to enable Vertex AI mode
|
|
- GOOGLE_CLOUD_PROJECT: Required GCP project ID
|
|
- GOOGLE_CLOUD_LOCATION: Optional region (default: us-central1)
|
|
- GOOGLE_VERTEX_BASE_URL: Optional custom API endpoint (for API gateway proxies)
|
|
- GOOGLE_APPLICATION_CREDENTIALS: Path to service account JSON (or use gcloud auth)
|
|
|
|
Args:
|
|
use_async: If True, return the async client (.aio), otherwise return sync client
|
|
|
|
Returns:
|
|
Gemini client instance (sync or async based on use_async parameter)
|
|
|
|
Raises:
|
|
ValueError: If required environment variables are not set
|
|
"""
|
|
use_vertex = is_vertex_ai_mode()
|
|
|
|
if use_vertex:
|
|
# Vertex AI mode - uses ADC (GOOGLE_APPLICATION_CREDENTIALS or gcloud auth)
|
|
project = os.getenv("GOOGLE_CLOUD_PROJECT")
|
|
location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")
|
|
base_url = os.getenv("GOOGLE_VERTEX_BASE_URL")
|
|
|
|
if not project:
|
|
raise ValueError(
|
|
"GOOGLE_CLOUD_PROJECT is required for Vertex AI mode. "
|
|
"Set GOOGLE_GENAI_USE_VERTEXAI=false to use AI Studio mode instead."
|
|
)
|
|
|
|
# Build http_options only if custom base_url is specified
|
|
http_options = None
|
|
if base_url:
|
|
http_options = {"base_url": base_url}
|
|
|
|
# Note: ADC handles authentication automatically
|
|
# via GOOGLE_APPLICATION_CREDENTIALS env var or gcloud auth
|
|
client = genai.Client(
|
|
vertexai=True, project=project, location=location, http_options=http_options
|
|
)
|
|
else:
|
|
# AI Studio mode - requires API key
|
|
api_key = os.getenv("GOOGLE_API_KEY")
|
|
if not api_key:
|
|
raise ValueError(
|
|
"GOOGLE_API_KEY is required for AI Studio mode. "
|
|
"Set GOOGLE_GENAI_USE_VERTEXAI=true and configure GCP credentials for Vertex AI mode."
|
|
)
|
|
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
# Return async or sync client based on parameter
|
|
return client.aio if use_async else client
|
|
|
|
|
|
def get_gemini_provider_name() -> str:
|
|
"""
|
|
Get the Gemini provider name based on current mode.
|
|
|
|
Returns:
|
|
Provider name string for display purposes
|
|
"""
|
|
if is_vertex_ai_mode():
|
|
project = os.getenv("GOOGLE_CLOUD_PROJECT", "unknown")
|
|
location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")
|
|
return f"Google Gemini (Vertex AI: {project}/{location})"
|
|
return "Google Gemini (AI Studio)"
|
|
|
|
|
|
def create_openai_client(use_async: bool = True):
|
|
"""
|
|
Create OpenAI client with optional custom base URL.
|
|
|
|
Environment variables:
|
|
- OPENAI_API_KEY: Required API key
|
|
- OPENAI_BASE_URL: Optional custom API endpoint (for proxies, Azure, etc.)
|
|
|
|
Args:
|
|
use_async: If True, return AsyncOpenAI, otherwise return OpenAI
|
|
|
|
Returns:
|
|
OpenAI client instance (async or sync based on use_async parameter)
|
|
|
|
Raises:
|
|
ValueError: If OPENAI_API_KEY is not set
|
|
"""
|
|
if not HAS_OPENAI:
|
|
raise ValueError("openai library is not installed.")
|
|
api_key = os.getenv("OPENAI_API_KEY")
|
|
if not api_key:
|
|
raise ValueError("OPENAI_API_KEY is required for OpenAI mode.")
|
|
|
|
base_url = os.getenv("OPENAI_BASE_URL")
|
|
|
|
if use_async:
|
|
return openai.AsyncOpenAI(base_url=base_url)
|
|
return openai.OpenAI(base_url=base_url)
|
|
|
|
|
|
def get_openai_provider_name() -> str:
|
|
"""
|
|
Get the OpenAI provider name, including custom endpoint if configured.
|
|
|
|
Returns:
|
|
Provider name string for display purposes
|
|
"""
|
|
base_url = os.getenv("OPENAI_BASE_URL")
|
|
if base_url:
|
|
return f"OpenAI (Custom: {base_url})"
|
|
return "OpenAI"
|
|
|
|
|
|
def is_openai_reasoning_model(model_name: str) -> bool:
|
|
"""
|
|
Check if the OpenAI model supports reasoning_effort parameter.
|
|
|
|
Models that support reasoning_effort:
|
|
- o-series: o1, o3, o4 and their variants (o1-mini, o1-2024-12-17, etc.)
|
|
- gpt-5 series: gpt-5, gpt-5.2, gpt-5-turbo, etc.
|
|
|
|
Non-reasoning models like gpt-4.1, gpt-4o, etc. will reject this parameter.
|
|
|
|
Handles proxy/router prefixes like "openai/o1-mini" or "openrouter/gpt-5.2".
|
|
|
|
Args:
|
|
model_name: The OpenAI model name (may include path prefix)
|
|
|
|
Returns:
|
|
True if the model supports reasoning_effort, False otherwise
|
|
"""
|
|
model_lower = model_name.lower()
|
|
|
|
# Handle proxy/router prefixes like "openai/o1-mini", "openrouter/gpt-5.2"
|
|
# Extract the base model name after the last "/"
|
|
if "/" in model_lower:
|
|
model_lower = model_lower.rsplit("/", 1)[-1]
|
|
|
|
# Match o-series and gpt-5 series
|
|
return model_lower.startswith(("o1", "o3", "o4", "gpt-5"))
|
|
|
|
|
|
def is_openai_retryable(error: Exception) -> bool:
|
|
"""
|
|
Determine if an OpenAI error should be retried.
|
|
|
|
Non-retryable errors:
|
|
- AuthenticationError (401): Invalid API key
|
|
- PermissionDeniedError (403): No access to resource
|
|
- BadRequestError (400): Invalid request format
|
|
- NotFoundError (404): Model or resource not found
|
|
|
|
Retryable errors:
|
|
- RateLimitError (429): Rate limit exceeded
|
|
- APIConnectionError: Network issues
|
|
- InternalServerError (500): Server errors
|
|
- APIStatusError with 502, 503, 504: Gateway/service errors
|
|
|
|
Args:
|
|
error: The exception from OpenAI API call
|
|
|
|
Returns:
|
|
True if the error should be retried, False otherwise
|
|
"""
|
|
if not HAS_OPENAI:
|
|
return True
|
|
|
|
# Authentication error - invalid API key (401)
|
|
if isinstance(error, openai.AuthenticationError):
|
|
return False
|
|
|
|
# Permission denied - no access to resource (403)
|
|
if isinstance(error, openai.PermissionDeniedError):
|
|
return False
|
|
|
|
# Bad request - invalid request format (400)
|
|
if isinstance(error, openai.BadRequestError):
|
|
return False
|
|
|
|
# Not found - model or resource doesn't exist (404)
|
|
if isinstance(error, openai.NotFoundError):
|
|
return False
|
|
|
|
# Rate limit exceeded - should retry with backoff (429)
|
|
if isinstance(error, openai.RateLimitError):
|
|
return True
|
|
|
|
# API connection error - network issues, should retry
|
|
if isinstance(error, openai.APIConnectionError):
|
|
return True
|
|
|
|
# Internal server error - should retry (500)
|
|
if isinstance(error, openai.InternalServerError):
|
|
return True
|
|
|
|
# For other APIStatusError, check HTTP status code
|
|
if isinstance(error, openai.APIStatusError):
|
|
# Retryable server-side errors
|
|
return error.status_code in (429, 500, 502, 503, 504)
|
|
|
|
# For unknown errors, default to retry (network issues, timeouts, etc.)
|
|
return True
|
|
|
|
|
|
def is_gemini_retryable(error: Exception) -> bool:
|
|
"""
|
|
Determine if a Gemini error should be retried.
|
|
|
|
Uses string matching on error messages since google-genai may not have
|
|
well-defined exception types for all error cases.
|
|
|
|
Non-retryable errors:
|
|
- API key errors
|
|
- Authentication/permission errors
|
|
- Invalid request errors
|
|
- Model not found errors
|
|
- Billing/quota permanently exceeded
|
|
|
|
Retryable errors:
|
|
- Rate limit (429)
|
|
- Server errors (500, 502, 503, 504)
|
|
- Timeout/connection errors
|
|
|
|
Args:
|
|
error: The exception from Gemini API call
|
|
|
|
Returns:
|
|
True if the error should be retried, False otherwise
|
|
"""
|
|
error_str = str(error).lower()
|
|
|
|
# API key / authentication errors - do not retry
|
|
if "api_key" in error_str or "api key" in error_str:
|
|
return False
|
|
if "authentication" in error_str or "authenticate" in error_str:
|
|
return False
|
|
if "invalid_api_key" in error_str or "invalid api key" in error_str:
|
|
return False
|
|
|
|
# Permission / forbidden errors - do not retry
|
|
if "permission" in error_str and "denied" in error_str:
|
|
return False
|
|
if "forbidden" in error_str or "403" in error_str:
|
|
return False
|
|
|
|
# Invalid request errors - do not retry
|
|
if "invalid" in error_str and ("request" in error_str or "argument" in error_str):
|
|
return False
|
|
if "400" in error_str and "bad request" in error_str:
|
|
return False
|
|
|
|
# Model not found - do not retry
|
|
if "model" in error_str and ("not found" in error_str or "not exist" in error_str):
|
|
return False
|
|
if "404" in error_str:
|
|
return False
|
|
|
|
# Billing / permanent quota errors - do not retry
|
|
if "billing" in error_str:
|
|
return False
|
|
if "quota" in error_str and ("exceeded" in error_str or "exhausted" in error_str):
|
|
# Check if it mentions billing which indicates permanent quota issue
|
|
if "billing" in error_str or "payment" in error_str:
|
|
return False
|
|
# Temporary quota (rate limit) - should retry
|
|
return True
|
|
|
|
# Rate limit errors - should retry (429)
|
|
if "rate" in error_str and "limit" in error_str:
|
|
return True
|
|
if "429" in error_str or "resource_exhausted" in error_str:
|
|
return True
|
|
|
|
# Server errors - should retry (500, 502, 503, 504)
|
|
if any(code in error_str for code in ["500", "502", "503", "504"]):
|
|
return True
|
|
if "internal" in error_str and ("error" in error_str or "server" in error_str):
|
|
return True
|
|
if "service" in error_str and "unavailable" in error_str:
|
|
return True
|
|
if "gateway" in error_str:
|
|
return True
|
|
|
|
# Timeout / connection errors - should retry
|
|
if "timeout" in error_str or "timed out" in error_str:
|
|
return True
|
|
if "connection" in error_str:
|
|
return True
|
|
if "network" in error_str:
|
|
return True
|
|
|
|
# Unknown errors - default to retry with limited attempts
|
|
return True
|
|
|
|
|
|
# JSON Schema for LLM structured output
|
|
AUDIT_RESULT_SCHEMA = {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
"is_violation": {
|
|
"type": "boolean",
|
|
"description": "Whether any violations were found",
|
|
},
|
|
"violations": {
|
|
"type": "array",
|
|
"description": "List of violations found",
|
|
"items": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
"rule_id": {
|
|
"type": "string",
|
|
"description": "ID of the violated rule (e.g., R001)",
|
|
},
|
|
"violation_text": {
|
|
"type": "string",
|
|
"description": "The problematic text directly verbatim quote from the source content, and not span multiple cells",
|
|
},
|
|
"violation_reason": {
|
|
"type": "string",
|
|
"description": "Explanation of why this violates the rule",
|
|
},
|
|
"fix_action": {
|
|
"type": "string",
|
|
"enum": ["replace", "manual"],
|
|
"description": "Action type: replace substitutes text (including deletion-via-replace), manual requires human review",
|
|
},
|
|
"revised_text": {
|
|
"type": "string",
|
|
"description": "For replace: complete replacement text (including deletion-via-replace). For manual: additional guidance for human reviewer",
|
|
},
|
|
},
|
|
"required": [
|
|
"rule_id",
|
|
"violation_text",
|
|
"violation_reason",
|
|
"fix_action",
|
|
"revised_text",
|
|
],
|
|
},
|
|
},
|
|
},
|
|
"required": ["is_violation", "violations"],
|
|
}
|
|
|
|
# JSON Schema for global extraction output
|
|
GLOBAL_EXTRACT_SCHEMA = {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
"results": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
"rule_id": {"type": "string"},
|
|
"extracted_results": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
"entity": {"type": "string"},
|
|
"fields": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
"name": {"type": "string"},
|
|
"value": {"type": "string"},
|
|
"evidence": {"type": "string"},
|
|
},
|
|
"required": ["name", "value", "evidence"],
|
|
},
|
|
},
|
|
},
|
|
"required": ["entity", "fields"],
|
|
},
|
|
},
|
|
},
|
|
"required": ["rule_id", "extracted_results"],
|
|
},
|
|
}
|
|
},
|
|
"required": ["results"],
|
|
}
|
|
|
|
# JSON Schema for global verification output
|
|
GLOBAL_VERIFY_SCHEMA = {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
"violations": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
"rule_id": {"type": "string"},
|
|
"uuid": {"type": "string"},
|
|
"uuid_end": {"type": "string"},
|
|
"violation_text": {"type": "string"},
|
|
"violation_reason": {"type": "string"},
|
|
"fix_action": {"type": "string", "enum": ["replace", "manual"]},
|
|
"revised_text": {"type": "string"},
|
|
},
|
|
"required": [
|
|
"rule_id",
|
|
"uuid",
|
|
"uuid_end",
|
|
"violation_text",
|
|
"violation_reason",
|
|
"fix_action",
|
|
"revised_text",
|
|
],
|
|
},
|
|
}
|
|
},
|
|
"required": ["violations"],
|
|
}
|
|
|
|
|
|
async def global_extract_gemini_async(
|
|
user_prompt: str,
|
|
system_prompt: str,
|
|
model_name: str,
|
|
client,
|
|
thinking_level: str = None,
|
|
thinking_budget: int = None,
|
|
) -> dict:
|
|
thinking_config = None
|
|
if thinking_level and thinking_level.upper() in (
|
|
"MINIMAL",
|
|
"LOW",
|
|
"MEDIUM",
|
|
"HIGH",
|
|
):
|
|
level_map = {
|
|
"MINIMAL": types.ThinkingLevel.MINIMAL,
|
|
"LOW": types.ThinkingLevel.LOW,
|
|
"MEDIUM": types.ThinkingLevel.MEDIUM,
|
|
"HIGH": types.ThinkingLevel.HIGH,
|
|
}
|
|
thinking_config = types.ThinkingConfig(
|
|
thinking_level=level_map[thinking_level.upper()]
|
|
)
|
|
elif thinking_budget is not None:
|
|
thinking_config = types.ThinkingConfig(thinking_budget=int(thinking_budget))
|
|
|
|
config_params = {
|
|
"system_instruction": system_prompt,
|
|
"response_mime_type": "application/json",
|
|
"response_schema": GLOBAL_EXTRACT_SCHEMA,
|
|
}
|
|
if thinking_config:
|
|
config_params["thinking_config"] = thinking_config
|
|
|
|
response = await client.models.generate_content(
|
|
model=model_name,
|
|
contents=user_prompt,
|
|
config=types.GenerateContentConfig(**config_params),
|
|
)
|
|
return json.loads(response.text)
|
|
|
|
|
|
async def global_extract_openai_async(
|
|
user_prompt: str,
|
|
system_prompt: str,
|
|
model_name: str,
|
|
client,
|
|
reasoning_effort: str = None,
|
|
) -> dict:
|
|
request_params = {
|
|
"model": model_name,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
],
|
|
"response_format": {
|
|
"type": "json_schema",
|
|
"json_schema": {
|
|
"name": "global_extract",
|
|
"strict": True,
|
|
"schema": GLOBAL_EXTRACT_SCHEMA,
|
|
},
|
|
},
|
|
}
|
|
if (
|
|
reasoning_effort
|
|
and reasoning_effort.lower() in ("low", "medium", "high")
|
|
and is_openai_reasoning_model(model_name)
|
|
):
|
|
request_params["reasoning_effort"] = reasoning_effort.lower()
|
|
|
|
response = await client.chat.completions.create(**request_params)
|
|
return json.loads(response.choices[0].message.content)
|
|
|
|
|
|
async def global_verify_gemini_async(
|
|
user_prompt: str,
|
|
system_prompt: str,
|
|
model_name: str,
|
|
client,
|
|
thinking_level: str = None,
|
|
thinking_budget: int = None,
|
|
) -> dict:
|
|
thinking_config = None
|
|
if thinking_level and thinking_level.upper() in (
|
|
"MINIMAL",
|
|
"LOW",
|
|
"MEDIUM",
|
|
"HIGH",
|
|
):
|
|
level_map = {
|
|
"MINIMAL": types.ThinkingLevel.MINIMAL,
|
|
"LOW": types.ThinkingLevel.LOW,
|
|
"MEDIUM": types.ThinkingLevel.MEDIUM,
|
|
"HIGH": types.ThinkingLevel.HIGH,
|
|
}
|
|
thinking_config = types.ThinkingConfig(
|
|
thinking_level=level_map[thinking_level.upper()]
|
|
)
|
|
elif thinking_budget is not None:
|
|
thinking_config = types.ThinkingConfig(thinking_budget=int(thinking_budget))
|
|
|
|
config_params = {
|
|
"system_instruction": system_prompt,
|
|
"response_mime_type": "application/json",
|
|
"response_schema": GLOBAL_VERIFY_SCHEMA,
|
|
}
|
|
if thinking_config:
|
|
config_params["thinking_config"] = thinking_config
|
|
|
|
response = await client.models.generate_content(
|
|
model=model_name,
|
|
contents=user_prompt,
|
|
config=types.GenerateContentConfig(**config_params),
|
|
)
|
|
return json.loads(response.text)
|
|
|
|
|
|
async def global_verify_openai_async(
|
|
user_prompt: str,
|
|
system_prompt: str,
|
|
model_name: str,
|
|
client,
|
|
reasoning_effort: str = None,
|
|
) -> dict:
|
|
request_params = {
|
|
"model": model_name,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
],
|
|
"response_format": {
|
|
"type": "json_schema",
|
|
"json_schema": {
|
|
"name": "global_verify",
|
|
"strict": True,
|
|
"schema": GLOBAL_VERIFY_SCHEMA,
|
|
},
|
|
},
|
|
}
|
|
if (
|
|
reasoning_effort
|
|
and reasoning_effort.lower() in ("low", "medium", "high")
|
|
and is_openai_reasoning_model(model_name)
|
|
):
|
|
request_params["reasoning_effort"] = reasoning_effort.lower()
|
|
|
|
response = await client.chat.completions.create(**request_params)
|
|
return json.loads(response.choices[0].message.content)
|
|
|
|
|
|
async def audit_block_gemini_async(
|
|
user_prompt: str,
|
|
system_prompt: str,
|
|
model_name: str,
|
|
client,
|
|
thinking_level: str = None,
|
|
thinking_budget: int = None,
|
|
) -> dict:
|
|
"""
|
|
Audit a text block using Google Gemini with strict JSON mode (async version).
|
|
|
|
Args:
|
|
user_prompt: User prompt to audit
|
|
system_prompt: Cached system prompt with rules and instructions
|
|
model_name: Gemini model to use
|
|
client: Gemini async client instance (client.aio)
|
|
thinking_level: Thinking level for Gemini 3 models (MINIMAL, LOW, MEDIUM, HIGH)
|
|
thinking_budget: Thinking token budget for Gemini 2.5 models (integer)
|
|
|
|
Returns:
|
|
Audit result dictionary
|
|
"""
|
|
# Build thinking config based on model and parameters
|
|
thinking_config = None
|
|
|
|
if thinking_level and thinking_level.upper() in (
|
|
"MINIMAL",
|
|
"LOW",
|
|
"MEDIUM",
|
|
"HIGH",
|
|
):
|
|
# For Gemini 3 models
|
|
level_map = {
|
|
"MINIMAL": types.ThinkingLevel.MINIMAL,
|
|
"LOW": types.ThinkingLevel.LOW,
|
|
"MEDIUM": types.ThinkingLevel.MEDIUM,
|
|
"HIGH": types.ThinkingLevel.HIGH,
|
|
}
|
|
thinking_config = types.ThinkingConfig(
|
|
thinking_level=level_map[thinking_level.upper()]
|
|
)
|
|
elif thinking_budget is not None:
|
|
# For Gemini 2.5 models
|
|
thinking_config = types.ThinkingConfig(thinking_budget=int(thinking_budget))
|
|
|
|
config_params = {
|
|
"system_instruction": system_prompt,
|
|
"response_mime_type": "application/json",
|
|
"response_schema": AUDIT_RESULT_SCHEMA,
|
|
}
|
|
|
|
# Only add thinking_config if it's configured
|
|
if thinking_config:
|
|
config_params["thinking_config"] = thinking_config
|
|
|
|
response = await client.models.generate_content(
|
|
model=model_name,
|
|
contents=user_prompt,
|
|
config=types.GenerateContentConfig(**config_params),
|
|
)
|
|
|
|
# With structured output, response is guaranteed to be valid JSON
|
|
result = json.loads(response.text)
|
|
return result
|
|
|
|
|
|
async def audit_block_openai_async(
|
|
user_prompt: str,
|
|
system_prompt: str,
|
|
model_name: str,
|
|
client,
|
|
reasoning_effort: str = None,
|
|
) -> dict:
|
|
"""
|
|
Audit a text block using OpenAI with strict JSON mode (async version).
|
|
|
|
Args:
|
|
user_prompt: User prompt to audit
|
|
system_prompt: Cached system prompt with rules and instructions
|
|
model_name: OpenAI model to use
|
|
client: AsyncOpenAI client instance
|
|
reasoning_effort: Reasoning effort for o-series models (low, medium, high)
|
|
|
|
Returns:
|
|
Audit result dictionary
|
|
"""
|
|
request_params = {
|
|
"model": model_name,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
],
|
|
"response_format": {
|
|
"type": "json_schema",
|
|
"json_schema": {
|
|
"name": "audit_result",
|
|
"strict": True,
|
|
"schema": AUDIT_RESULT_SCHEMA,
|
|
},
|
|
},
|
|
}
|
|
|
|
# Add reasoning_effort only for o-series models that support it
|
|
if (
|
|
reasoning_effort
|
|
and reasoning_effort.lower() in ("low", "medium", "high")
|
|
and is_openai_reasoning_model(model_name)
|
|
):
|
|
request_params["reasoning_effort"] = reasoning_effort.lower()
|
|
|
|
response = await client.chat.completions.create(**request_params)
|
|
|
|
# With structured output, response is guaranteed to be valid JSON
|
|
result = json.loads(response.choices[0].message.content)
|
|
return result
|