wehub-resource-sync
803 lines
34 KiB
Python
803 lines
34 KiB
Python
"""LLM client and prompt building for SimpleQA evaluation.
|
||
|
||
Supports:
|
||
- Google Gemini (Vertex AI and standard API)
|
||
- OpenAI-compatible APIs (vLLM, etc.)
|
||
"""
|
||
|
||
import asyncio
|
||
import base64
|
||
import logging
|
||
import os
|
||
|
||
# Try to import Google GenAI for Gemini support
|
||
try:
|
||
import google.genai as genai
|
||
from google.genai.types import (
|
||
GenerateContentConfig,
|
||
Part,
|
||
Blob,
|
||
HttpOptions,
|
||
Content,
|
||
)
|
||
|
||
GEMINI_AVAILABLE = True
|
||
except ImportError:
|
||
GEMINI_AVAILABLE = False
|
||
genai = None
|
||
GenerateContentConfig = None
|
||
Part = None
|
||
Blob = None
|
||
HttpOptions = None
|
||
Content = None
|
||
|
||
from .retrieval import RetrievalResult
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# System Prompts
|
||
SYSTEM_PROMPT_NAIVE = """You are a research assistant who answers questions.
|
||
Use <think></think> tags to show your reasoning if needed.
|
||
Answer the question directly and concisely.
|
||
"""
|
||
|
||
SYSTEM_PROMPT_EVIDENCE_QA = """You are a research assistant who answers questions based on provided evidence.
|
||
Use <think></think> tags to show your reasoning if needed.
|
||
Answer the question directly and concisely based ONLY on the provided evidence.
|
||
"""
|
||
|
||
SYSTEM_PROMPT_SCREENSHOT = SYSTEM_PROMPT_EVIDENCE_QA
|
||
|
||
SYSTEM_PROMPT_TEXT_RAG = SYSTEM_PROMPT_EVIDENCE_QA
|
||
|
||
SYSTEM_PROMPT_VECTOR = SYSTEM_PROMPT_EVIDENCE_QA
|
||
|
||
SYSTEM_PROMPT_SHORT_ANSWER = """Answer the question with as few words as possible. Give only the answer, no explanation.
|
||
"""
|
||
|
||
SYSTEM_PROMPT_REACT = """You are a research assistant who answers questions using a search tool.
|
||
You will be provided with retrieved Wikipedia screenshot tiles as evidence.
|
||
|
||
IMPORTANT: Try your best to answer with the evidence you have. Only search again if the evidence is clearly about a WRONG topic and does not contain the answer at all.
|
||
|
||
To search for different evidence, output ONLY: <search>your refined search query</search>
|
||
Otherwise, answer the question directly and concisely.
|
||
|
||
Rules:
|
||
- READ the evidence images carefully — the answer is often there even if not obvious.
|
||
- If the images show the relevant Wikipedia article, answer from them. Do NOT search again.
|
||
- Only use <search> if the retrieved tiles are about a completely unrelated topic.
|
||
- Do NOT repeat the same search query — use different keywords.
|
||
- Use <think></think> tags to show your reasoning if needed.
|
||
"""
|
||
|
||
SYSTEM_PROMPT_REACT_V2 = """You are a research assistant who answers questions using a search tool.
|
||
You will be provided with retrieved Wikipedia screenshot tiles as evidence.
|
||
|
||
You have two actions:
|
||
1. **Answer**: If you can find or infer the answer from the evidence, respond with your answer directly.
|
||
2. **Search**: If the evidence does NOT contain the answer, output: <search>new search query</search>
|
||
|
||
CRITICAL rules:
|
||
- ALWAYS try to answer first. Only search if the evidence is about the WRONG topic entirely.
|
||
- Each search query MUST use DIFFERENT keywords than all previous queries. Think about synonyms, related entities, or the answer's broader topic.
|
||
- If you've already searched 2+ times without finding the answer, make your BEST GUESS based on whatever partial evidence you have. Do not give up.
|
||
- Never output an empty answer. If unsure, state your best guess with a caveat.
|
||
- Use <think></think> tags for reasoning.
|
||
"""
|
||
|
||
SYSTEM_PROMPT_REACT_MULTIHOP = """You are a research assistant who answers multi-hop questions using a search tool.
|
||
You will be provided with retrieved Wikipedia screenshot tiles as evidence.
|
||
|
||
Multi-hop questions require information from MULTIPLE Wikipedia pages. For example:
|
||
- "Where did X's father die?" → First find who X's father is, then search for the father's death place.
|
||
- "Which film came out first, A or B?" → Search for film A's release date, then film B's release date.
|
||
|
||
Strategy:
|
||
1. Read the evidence carefully. Extract any INTERMEDIATE facts (names, dates, locations) that help answer the question.
|
||
2. If you found an intermediate fact but still need more info, search for the next entity: <search>entity name topic</search>
|
||
3. Only give your final answer when you have ALL the pieces needed.
|
||
|
||
Rules:
|
||
- For multi-hop questions, you will usually need 2-3 searches. This is EXPECTED — do not try to answer with just the first search.
|
||
- In <think> tags, ALWAYS record: the specific facts you found (names, dates, places) so you don't lose them.
|
||
- Extract specific entity names from evidence tiles to use as search queries.
|
||
- Each search query MUST use DIFFERENT keywords. Be specific: use full names, dates, or titles you found.
|
||
- When you have enough info, give a concise final answer.
|
||
"""
|
||
|
||
SYSTEM_PROMPT_PIXEL_QUERY = """You are a research assistant who answers questions based on retrieved visual evidence.
|
||
The first image contains the question you need to answer.
|
||
The remaining images are retrieved evidence that may contain the answer.
|
||
Read the question from the first image, then use the evidence images to answer it.
|
||
Use <think></think> tags to show your reasoning if needed.
|
||
Answer the question directly and concisely.
|
||
"""
|
||
|
||
SYSTEM_PROMPT_MULTIMODAL_QUERY = """You are a research assistant who answers questions based on retrieved visual evidence.
|
||
You will receive: (1) a text question, (2) a query image, and (3) retrieved Wikipedia evidence images.
|
||
Use the query image and evidence images to answer the question.
|
||
Use <think></think> tags to show your reasoning if needed.
|
||
Answer the question directly and concisely.
|
||
"""
|
||
|
||
|
||
def _build_fewshot_turns(demos: list[dict], encode_image_fn) -> list[dict]:
|
||
"""Build a list of (user, assistant) message turns for in-context few-shot.
|
||
|
||
Each demo becomes: user={Q text + demo image} → assistant={answer}. The
|
||
chat-tuned model treats these as prior conversation turns rather than
|
||
mixing them with the current question's evidence — this is the canonical
|
||
few-shot format for instruction-tuned chat models.
|
||
"""
|
||
turns: list[dict] = []
|
||
for demo in demos:
|
||
user_content: list[dict] = [
|
||
{"type": "text", "text": f"Question: {demo['question']}"},
|
||
]
|
||
img_path = demo.get("image_path")
|
||
if img_path and encode_image_fn and os.path.exists(img_path):
|
||
try:
|
||
b64 = encode_image_fn(img_path)
|
||
if b64:
|
||
user_content.append(
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": f"data:image/png;base64,{b64}"},
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to encode few-shot image {img_path}: {e}")
|
||
turns.append({"role": "user", "content": user_content})
|
||
turns.append({"role": "assistant", "content": demo["answer"]})
|
||
return turns
|
||
|
||
|
||
def build_messages(
|
||
query: str,
|
||
retrieval_result: RetrievalResult,
|
||
encode_image_fn=None,
|
||
additional_instructions: str | None = None,
|
||
few_shot_demos: list[dict] | None = None,
|
||
) -> list[dict]:
|
||
"""Build messages for LLM based on retrieval result.
|
||
|
||
When ``retrieval_result.pixel_query_path`` is set the query is sent as an
|
||
image. Two modes:
|
||
- **Multimodal** (retrieval_type contains "multimodal"): text question + query image + retrieved tiles.
|
||
- **Pixel query** (rendered question as image): first image = question, then retrieved tiles.
|
||
"""
|
||
# ---- Multimodal / pixel-query mode: text + raw species/landmark photo + retrieved tiles ----
|
||
# query_image_path = raw species/landmark photo (for generation, always).
|
||
# pixel_query_path = rendered card or raw photo (for retrieval only; ignored here).
|
||
# Falls back to pixel_query_path if query_image_path is not set (backward compat).
|
||
gen_image_path = (
|
||
retrieval_result.query_image_path or retrieval_result.pixel_query_path
|
||
)
|
||
if gen_image_path and encode_image_fn:
|
||
system_prompt = SYSTEM_PROMPT_MULTIMODAL_QUERY
|
||
# Decide evidence_note based on what retrieval actually returned. Three cases:
|
||
# (a) retrieved images (screenshot retrieval) — evidence is image tiles after the query
|
||
# (b) retrieved text (text retrieval) — evidence is rendered as text after the query
|
||
# (c) no retrieval — query image only
|
||
# Until 2026-04-29 this branch silently dropped retrieval_result.text whenever the
|
||
# query image was set, turning every "EVQA + text retrieval" cell into an effective
|
||
# naive run. Fixed by adding the text-passages block alongside the multimodal preamble.
|
||
if retrieval_result.images:
|
||
evidence_note = "The first image is the query image. The following images are retrieved Wikipedia evidence. Answer the question based on the evidence."
|
||
elif retrieval_result.text:
|
||
evidence_note = "The image is the query image. Below is retrieved Wikipedia evidence (text). Answer the question based on the evidence and the image."
|
||
else:
|
||
evidence_note = "The first image is the query image. Answer the question based on the image (no additional evidence was retrieved)."
|
||
text_parts = [
|
||
f"Question: {query}",
|
||
"",
|
||
evidence_note,
|
||
]
|
||
if retrieval_result.text:
|
||
# Option 1: no URL header in multimodal branch either. Reader gets the
|
||
# chunks and the query image, no metadata leak.
|
||
text_parts.extend(
|
||
[
|
||
"",
|
||
retrieval_result.text,
|
||
]
|
||
)
|
||
if additional_instructions:
|
||
text_parts.append("")
|
||
text_parts.append(additional_instructions)
|
||
user_content: list[dict] = [
|
||
{"type": "text", "text": "\n".join(text_parts)},
|
||
]
|
||
|
||
# Add raw species/landmark photo
|
||
if os.path.exists(gen_image_path):
|
||
try:
|
||
img_base64 = encode_image_fn(gen_image_path)
|
||
if img_base64:
|
||
user_content.append(
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": f"data:image/png;base64,{img_base64}"},
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to encode query image {gen_image_path}: {e}")
|
||
user_content.append(
|
||
{"type": "text", "text": f"(Image unavailable) Query: {query}"}
|
||
)
|
||
else:
|
||
logger.warning(f"Query image not found: {gen_image_path}")
|
||
user_content.append({"type": "text", "text": f"Query: {query}"})
|
||
|
||
# Add retrieved tiles
|
||
if retrieval_result.images:
|
||
for img_path, score in retrieval_result.images:
|
||
if os.path.exists(img_path):
|
||
try:
|
||
img_base64 = encode_image_fn(img_path)
|
||
if img_base64:
|
||
user_content.append(
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/png;base64,{img_base64}"
|
||
},
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to encode image {img_path}: {e}")
|
||
|
||
return [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": user_content},
|
||
]
|
||
|
||
# ---- Original modes --------------------------------------------------
|
||
# Select system prompt based on retrieval type
|
||
if retrieval_result.base64_image:
|
||
system_prompt = SYSTEM_PROMPT_SCREENSHOT
|
||
user_content = [
|
||
{"type": "text", "text": query},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/png;base64,{retrieval_result.base64_image}"
|
||
},
|
||
},
|
||
]
|
||
elif (
|
||
retrieval_result.retrieval_type == "text_api+rendered"
|
||
and retrieval_result.images
|
||
and encode_image_fn
|
||
):
|
||
# Text retrieval rendered as images. Mirror the text-RAG framing so
|
||
# evidence comes first and the reader sees an explicit "Question:"
|
||
# suffix — same structure as the text→text branch below, only the
|
||
# evidence modality differs.
|
||
system_prompt = SYSTEM_PROMPT_TEXT_RAG
|
||
user_content = []
|
||
for img_path, score in retrieval_result.images:
|
||
if os.path.exists(img_path):
|
||
try:
|
||
img_base64 = encode_image_fn(img_path)
|
||
if img_base64:
|
||
user_content.append(
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/png;base64,{img_base64}"
|
||
},
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to encode image {img_path}: {e}")
|
||
user_content.append({"type": "text", "text": f"Question: {query}"})
|
||
elif retrieval_result.images and encode_image_fn:
|
||
system_prompt = SYSTEM_PROMPT_VECTOR
|
||
user_content = [{"type": "text", "text": query}]
|
||
# Encode and add retrieved images
|
||
for img_path, score in retrieval_result.images:
|
||
if os.path.exists(img_path):
|
||
try:
|
||
img_base64 = encode_image_fn(img_path)
|
||
if img_base64:
|
||
user_content.append(
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/png;base64,{img_base64}"
|
||
},
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to encode image {img_path}: {e}")
|
||
elif retrieval_result.text:
|
||
system_prompt = SYSTEM_PROMPT_TEXT_RAG
|
||
# Option 1 (2026-04-29): no `Context from {urls}:` wrapper. URL leak gave
|
||
# text retrieval an unfair advantage on entity-answering tasks. Reader sees
|
||
# only the retrieved chunks and the question. URL still recorded in the
|
||
# JSONL via retrieval_result.source_url for logging/grading.
|
||
user_content = f"""{retrieval_result.text}
|
||
|
||
Question: {query}"""
|
||
else:
|
||
# Naive mode
|
||
system_prompt = SYSTEM_PROMPT_NAIVE
|
||
user_content = query
|
||
|
||
# Append additional instructions (e.g. short-answer prompt for EM-eval tasks)
|
||
if additional_instructions:
|
||
if isinstance(user_content, str):
|
||
user_content = user_content + "\n\n" + additional_instructions
|
||
else:
|
||
# list of content blocks — append as text
|
||
user_content.append({"type": "text", "text": additional_instructions})
|
||
|
||
# Few-shot as prior user/assistant turns (canonical chat few-shot format)
|
||
if few_shot_demos and encode_image_fn:
|
||
fewshot_turns = _build_fewshot_turns(few_shot_demos, encode_image_fn)
|
||
else:
|
||
fewshot_turns = []
|
||
|
||
return [
|
||
{"role": "system", "content": system_prompt},
|
||
*fewshot_turns,
|
||
{"role": "user", "content": user_content},
|
||
]
|
||
|
||
|
||
def _encode_images_to_content(
|
||
images: list[tuple[str, float]], encode_image_fn
|
||
) -> list[dict]:
|
||
"""Encode image paths to base64 content blocks."""
|
||
content = []
|
||
for img_path, score in images:
|
||
if os.path.exists(img_path):
|
||
try:
|
||
img_base64 = encode_image_fn(img_path)
|
||
if img_base64:
|
||
content.append(
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": f"data:image/png;base64,{img_base64}"},
|
||
}
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to encode image {img_path}: {e}")
|
||
return content
|
||
|
||
|
||
def build_react_messages(
|
||
query: str,
|
||
retrieval_results: list[RetrievalResult],
|
||
assistant_responses: list[str],
|
||
encode_image_fn=None,
|
||
prompt_version: str = "v1",
|
||
is_last_turn: bool = False,
|
||
previous_queries: list[str] | None = None,
|
||
) -> list[dict]:
|
||
"""Build multi-turn messages for ReAct retrieval loop.
|
||
|
||
Args:
|
||
query: Original question text.
|
||
retrieval_results: List of RetrievalResult from each round.
|
||
assistant_responses: List of assistant responses from previous rounds.
|
||
encode_image_fn: Function to encode images to base64.
|
||
prompt_version: "v1" (original) or "v2" (improved).
|
||
is_last_turn: If True, add force-answer instruction.
|
||
previous_queries: List of previous search queries (for v2, to avoid repetition).
|
||
|
||
Returns:
|
||
Messages list for the LLM.
|
||
"""
|
||
_prompt_map = {
|
||
"v1": SYSTEM_PROMPT_REACT,
|
||
"v2": SYSTEM_PROMPT_REACT_V2,
|
||
"multihop": SYSTEM_PROMPT_REACT_MULTIHOP,
|
||
}
|
||
system_prompt = _prompt_map.get(prompt_version, SYSTEM_PROMPT_REACT_V2)
|
||
messages = [{"role": "system", "content": system_prompt}]
|
||
|
||
for turn_idx, retrieval_result in enumerate(retrieval_results):
|
||
# Build user message with evidence images
|
||
if turn_idx == 0:
|
||
user_content: list[dict] = [
|
||
{
|
||
"type": "text",
|
||
"text": f"Question: {query}\n\nHere are retrieved Wikipedia evidence tiles:",
|
||
}
|
||
]
|
||
else:
|
||
text = "Here are new search results for your query:"
|
||
# Remind model of previous queries to avoid repetition (v2 and multihop)
|
||
if prompt_version in ("v2", "multihop") and previous_queries:
|
||
used = previous_queries[:turn_idx]
|
||
if used:
|
||
text += f"\n⚠️ You already searched: {used}. Do NOT repeat these. Use DIFFERENT keywords."
|
||
user_content = [{"type": "text", "text": text}]
|
||
|
||
if retrieval_result.images and encode_image_fn:
|
||
user_content.extend(
|
||
_encode_images_to_content(retrieval_result.images, encode_image_fn)
|
||
)
|
||
|
||
if not retrieval_result.has_content:
|
||
user_content.append(
|
||
{"type": "text", "text": "(No results found for this search.)"}
|
||
)
|
||
|
||
# On last turn, inject force-answer instruction
|
||
if is_last_turn and turn_idx == len(retrieval_results) - 1:
|
||
user_content.append(
|
||
{
|
||
"type": "text",
|
||
"text": (
|
||
"\n⚠️ This is your FINAL turn. You MUST provide an answer now — do NOT search again. "
|
||
"Give your best answer based on ALL evidence seen so far. If uncertain, make your best guess."
|
||
),
|
||
}
|
||
)
|
||
|
||
messages.append({"role": "user", "content": user_content})
|
||
|
||
# Add assistant response if we have one for this turn
|
||
if turn_idx < len(assistant_responses):
|
||
messages.append(
|
||
{"role": "assistant", "content": assistant_responses[turn_idx]}
|
||
)
|
||
|
||
return messages
|
||
|
||
|
||
class LLMClient:
|
||
"""Simplified async LLM client for Gemini using Vertex AI."""
|
||
|
||
def __init__(
|
||
self,
|
||
model: str,
|
||
api_base: str = "http://localhost:8000/v1",
|
||
api_key: str = "dummy",
|
||
temperature: float = 0.0,
|
||
max_tokens: int = 16384,
|
||
timeout: float = 120.0,
|
||
max_context_tokens: int | None = None,
|
||
enable_thinking: bool | None = None,
|
||
force_openai_compat: bool = False,
|
||
):
|
||
self.model = model
|
||
self.temperature = temperature
|
||
self.max_tokens = max_tokens
|
||
self.timeout = timeout
|
||
self.max_context_tokens = max_context_tokens
|
||
self.enable_thinking = enable_thinking
|
||
print(f"context length model: {max_context_tokens}")
|
||
|
||
# Gemini routes to Google GenAI SDK unless forced to OpenAI-compatible
|
||
# (aggregators like OpenRouter / Commonstack expose Gemini via OAI-compat).
|
||
self.is_gemini = ("gemini" in model.lower()) and not force_openai_compat
|
||
|
||
if self.is_gemini:
|
||
if not GEMINI_AVAILABLE:
|
||
raise ImportError(
|
||
"google-genai package is required for Gemini models. Install with: pip install google-genai"
|
||
)
|
||
|
||
# Use Vertex AI if GEMINI_API_KEY is set and GOOGLE_GENAI_USE_VERTEXAI is true
|
||
vertex_api_key = os.getenv("GEMINI_API_KEY")
|
||
use_vertex = os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true"
|
||
if vertex_api_key and use_vertex:
|
||
logger.info(f"Using Vertex AI for Gemini model: {model}")
|
||
# Ensure GOOGLE_API_KEY is not set when using Vertex AI (it causes conflicts)
|
||
if "GOOGLE_API_KEY" in os.environ:
|
||
logger.warning(
|
||
"GOOGLE_API_KEY is set but using Vertex AI. Unsetting GOOGLE_API_KEY to avoid conflicts."
|
||
)
|
||
del os.environ["GOOGLE_API_KEY"]
|
||
os.environ["GEMINI_API_KEY"] = vertex_api_key
|
||
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "true"
|
||
self.gemini_client = genai.Client(
|
||
http_options=HttpOptions(api_version="v1")
|
||
)
|
||
else:
|
||
# Use standard Gemini API
|
||
logger.info(f"Using standard Gemini API for model: {model}")
|
||
api_key = api_key if api_key != "dummy" else os.getenv("GOOGLE_API_KEY")
|
||
if not api_key:
|
||
raise ValueError(
|
||
"GOOGLE_API_KEY or GEMINI_API_KEY environment variable is required for Gemini models"
|
||
)
|
||
self.gemini_client = genai.Client(api_key=api_key)
|
||
else:
|
||
# Use OpenAI-compatible API
|
||
from openai import AsyncOpenAI
|
||
|
||
logger.info(f"Using OpenAI-compatible API: {api_base}")
|
||
self.client = AsyncOpenAI(
|
||
api_key=api_key,
|
||
base_url=api_base,
|
||
timeout=timeout,
|
||
max_retries=0,
|
||
)
|
||
self.gemini_client = None
|
||
|
||
async def generate(
|
||
self, messages: list[dict], max_retries: int = 3, connection_retries: int = 12
|
||
) -> tuple[str, dict]:
|
||
"""Generate response from messages with retry on timeout/connection errors.
|
||
|
||
Args:
|
||
max_retries: Retry count for timeout errors.
|
||
connection_retries: Retry count for connection errors (server restart).
|
||
12 retries × 10s = ~2 min window for server to come back.
|
||
|
||
Returns:
|
||
Tuple of (generated_text, usage_dict).
|
||
"""
|
||
# Check and truncate if needed
|
||
if hasattr(self, "max_context_tokens") and self.max_context_tokens:
|
||
estimated_tokens = self._estimate_tokens(messages)
|
||
if estimated_tokens > self.max_context_tokens - self.max_tokens:
|
||
logger.warning(
|
||
f"Estimated {estimated_tokens} tokens exceeds limit, truncating..."
|
||
)
|
||
messages = self._truncate_messages(messages, self.max_context_tokens)
|
||
|
||
conn_attempts = 0
|
||
timeout_attempts = 0
|
||
while True:
|
||
try:
|
||
if self.is_gemini:
|
||
return await self._generate_gemini(messages)
|
||
else:
|
||
return await self._generate_openai(messages)
|
||
except asyncio.TimeoutError:
|
||
timeout_attempts += 1
|
||
if timeout_attempts >= max_retries:
|
||
raise
|
||
wait_time = 2**timeout_attempts # 2, 4, 8 seconds
|
||
logger.warning(
|
||
f"Timeout on attempt {timeout_attempts}/{max_retries}, retrying in {wait_time}s..."
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
except Exception as e:
|
||
error_str = str(e).lower()
|
||
if "timeout" in error_str or "timed out" in error_str:
|
||
timeout_attempts += 1
|
||
if timeout_attempts >= max_retries:
|
||
raise
|
||
wait_time = 2**timeout_attempts
|
||
logger.warning(
|
||
f"Timeout on attempt {timeout_attempts}/{max_retries}, retrying in {wait_time}s..."
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
elif "connection" in error_str or "connect" in error_str:
|
||
conn_attempts += 1
|
||
if conn_attempts >= connection_retries:
|
||
raise
|
||
wait_time = 10 # fixed 10s — server restart takes ~30-60s
|
||
logger.warning(
|
||
f"Connection error ({conn_attempts}/{connection_retries}), retrying in {wait_time}s..."
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
elif (
|
||
"429" in error_str
|
||
or "rate_limit" in error_str
|
||
or "rate limit" in error_str
|
||
):
|
||
# Provider rate limit — exponential backoff with jitter
|
||
timeout_attempts += 1
|
||
if timeout_attempts >= max_retries + 3: # extra patience for 429
|
||
raise
|
||
import random
|
||
|
||
wait_time = min(60, 5 * (2**timeout_attempts)) + random.uniform(
|
||
0, 3
|
||
)
|
||
logger.warning(
|
||
f"429 rate-limit (attempt {timeout_attempts}), backing off {wait_time:.1f}s..."
|
||
)
|
||
await asyncio.sleep(wait_time)
|
||
else:
|
||
raise
|
||
|
||
async def _generate_gemini(self, messages: list[dict]) -> tuple[str, dict]:
|
||
"""Generate using Gemini API."""
|
||
# Extract system prompt and user content
|
||
system_prompt = None
|
||
user_content = None
|
||
|
||
for msg in messages:
|
||
if msg.get("role") == "system":
|
||
system_prompt = msg.get("content", "")
|
||
elif msg.get("role") == "user":
|
||
user_content = msg.get("content", "")
|
||
|
||
# Build parts for Gemini
|
||
parts = []
|
||
|
||
# Add system prompt to the beginning of user message if present
|
||
if system_prompt:
|
||
parts.append(Part(text=f"{system_prompt}\n\n"))
|
||
|
||
# Process user content
|
||
if isinstance(user_content, str):
|
||
# Simple text
|
||
if parts:
|
||
parts[0] = Part(text=parts[0].text + user_content)
|
||
else:
|
||
parts.append(Part(text=user_content))
|
||
elif isinstance(user_content, list):
|
||
# Multi-modal content
|
||
for item in user_content:
|
||
if item.get("type") == "text":
|
||
text = item.get("text", "")
|
||
if (
|
||
parts
|
||
and isinstance(parts[0], Part)
|
||
and hasattr(parts[0], "text")
|
||
):
|
||
# Append to existing text part
|
||
parts[0] = Part(text=parts[0].text + text)
|
||
else:
|
||
parts.append(Part(text=text))
|
||
elif item.get("type") == "image_url":
|
||
# Extract base64 image
|
||
image_url = item.get("image_url", {}).get("url", "")
|
||
if image_url.startswith("data:image"):
|
||
try:
|
||
header, data = image_url.split(",", 1)
|
||
mime_type = header.split(";")[0].split(":")[1]
|
||
image_bytes = base64.b64decode(data)
|
||
parts.append(
|
||
Part(
|
||
inline_data=Blob(
|
||
mime_type=mime_type, data=image_bytes
|
||
)
|
||
)
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Failed to process image: {e}")
|
||
raise
|
||
|
||
# Create content
|
||
content = Content(role="user", parts=parts)
|
||
|
||
# Call API in executor to avoid blocking
|
||
loop = asyncio.get_event_loop()
|
||
|
||
def _call_api():
|
||
try:
|
||
response = self.gemini_client.models.generate_content(
|
||
model=self.model,
|
||
contents=[content],
|
||
config=GenerateContentConfig(
|
||
temperature=self.temperature, max_output_tokens=self.max_tokens
|
||
),
|
||
)
|
||
return response
|
||
except Exception as e:
|
||
logger.error(f"Gemini API error: {e}")
|
||
raise
|
||
|
||
response = await loop.run_in_executor(None, _call_api)
|
||
|
||
# Extract text
|
||
text = response.text if hasattr(response, "text") and response.text else ""
|
||
|
||
# Extract usage
|
||
usage = {}
|
||
if hasattr(response, "usage_metadata") and response.usage_metadata:
|
||
usage_meta = response.usage_metadata
|
||
usage = {
|
||
"prompt_tokens": getattr(usage_meta, "prompt_token_count", 0),
|
||
"completion_tokens": getattr(usage_meta, "candidates_token_count", 0),
|
||
"total_tokens": getattr(usage_meta, "total_token_count", 0),
|
||
}
|
||
|
||
return text, usage
|
||
|
||
async def _generate_openai(self, messages: list[dict]) -> tuple[str, dict]:
|
||
"""Generate using OpenAI-compatible API."""
|
||
kwargs = dict(
|
||
model=self.model,
|
||
messages=messages,
|
||
max_tokens=self.max_tokens,
|
||
timeout=self.timeout,
|
||
)
|
||
# Some modern reasoning models deprecate `temperature` (Claude Opus 4.7+, some GPT-5 variants).
|
||
# Only send it when we actually want to override the default.
|
||
model_lower = self.model.lower()
|
||
drops_temperature = any(
|
||
x in model_lower for x in ("opus-4-7", "opus-4-8", "gpt-5.4-pro")
|
||
)
|
||
if not drops_temperature:
|
||
kwargs["temperature"] = self.temperature
|
||
if self.enable_thinking is not None:
|
||
kwargs["extra_body"] = {
|
||
"chat_template_kwargs": {"enable_thinking": self.enable_thinking}
|
||
}
|
||
response = await self.client.chat.completions.create(**kwargs)
|
||
|
||
generated_text = response.choices[0].message.content
|
||
|
||
usage = {}
|
||
if response.usage:
|
||
usage = {
|
||
"prompt_tokens": response.usage.prompt_tokens,
|
||
"completion_tokens": response.usage.completion_tokens,
|
||
"total_tokens": response.usage.total_tokens,
|
||
}
|
||
|
||
return generated_text, usage
|
||
|
||
def _estimate_tokens(self, messages: list[dict]) -> int:
|
||
"""Estimate token count from messages (rough: ~4 chars per token)."""
|
||
total_chars = 0
|
||
for msg in messages:
|
||
content = msg.get("content", "")
|
||
if isinstance(content, str):
|
||
total_chars += len(content)
|
||
elif isinstance(content, list):
|
||
for item in content:
|
||
if isinstance(item, dict):
|
||
if item.get("type") == "text":
|
||
total_chars += len(item.get("text", ""))
|
||
elif item.get("type") == "image_url":
|
||
# Rough estimate for image tokens
|
||
total_chars += 1000 * 4 # ~1000 tokens per image
|
||
return total_chars // 4
|
||
|
||
def _truncate_messages(self, messages: list[dict], max_tokens: int) -> list[dict]:
|
||
"""Truncate text content in messages to fit within token limit."""
|
||
# Reserve tokens for response
|
||
available_tokens = max_tokens - self.max_tokens - 500 # buffer
|
||
max_chars = available_tokens * 4
|
||
|
||
truncated = []
|
||
total_chars = 0
|
||
|
||
for msg in messages:
|
||
new_msg = msg.copy()
|
||
content = msg.get("content", "")
|
||
|
||
if isinstance(content, str):
|
||
if total_chars + len(content) > max_chars:
|
||
remaining = max(0, max_chars - total_chars)
|
||
new_msg["content"] = (
|
||
content[:remaining]
|
||
+ "\n\n[Content truncated due to context limit]"
|
||
)
|
||
logger.warning(
|
||
f"Truncated message content from {len(content)} to {remaining} chars"
|
||
)
|
||
total_chars += len(new_msg["content"])
|
||
elif isinstance(content, list):
|
||
new_content = []
|
||
for item in content:
|
||
if isinstance(item, dict) and item.get("type") == "text":
|
||
text = item.get("text", "")
|
||
if total_chars + len(text) > max_chars:
|
||
remaining = max(0, max_chars - total_chars)
|
||
new_item = item.copy()
|
||
new_item["text"] = (
|
||
text[:remaining]
|
||
+ "\n\n[Content truncated due to context limit]"
|
||
)
|
||
new_content.append(new_item)
|
||
logger.warning(
|
||
f"Truncated text content from {len(text)} to {remaining} chars"
|
||
)
|
||
total_chars += remaining
|
||
else:
|
||
new_content.append(item)
|
||
total_chars += len(text)
|
||
else:
|
||
new_content.append(item)
|
||
if isinstance(item, dict) and item.get("type") == "image_url":
|
||
total_chars += 1000 * 4 # image token estimate
|
||
new_msg["content"] = new_content
|
||
truncated.append(new_msg)
|
||
|
||
return truncated
|