806 lines
29 KiB
Python
806 lines
29 KiB
Python
"""
|
|
AI client:
|
|
• Text + tool calls → native google-genai (Google) OR OpenAI client (OpenRouter)
|
|
• Vision → google-genai native always
|
|
• Image generation → google-genai Imagen 4 Fast always
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
|
|
from openai import AsyncOpenAI
|
|
from openai import APIStatusError
|
|
from google import genai
|
|
from google.genai import types as gt
|
|
|
|
from config import (
|
|
GOOGLE_API_KEY, GOOGLE_AI_BASE_URL,
|
|
OPENROUTER_API_KEY, OPENROUTER_BASE_URL,
|
|
DEFAULT_MODEL, DEFAULT_OR_MODEL, IMAGEN_MODEL,
|
|
MAX_API_CONTEXT,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ── Clients ───────────────────────────────────────────────────
|
|
|
|
_or_client = AsyncOpenAI(
|
|
api_key=OPENROUTER_API_KEY or "sk-no-key",
|
|
base_url=OPENROUTER_BASE_URL,
|
|
default_headers={
|
|
"HTTP-Referer": "https://github.com/puppy-bot",
|
|
"X-Title": "Puppy Bot",
|
|
},
|
|
)
|
|
_gc = genai.Client(api_key=GOOGLE_API_KEY)
|
|
|
|
_current_provider: str = "google"
|
|
_google_model: str = DEFAULT_MODEL
|
|
_or_model: str = DEFAULT_OR_MODEL
|
|
_last_fallback: str = ""
|
|
|
|
GOOGLE_TEXT_MODELS = [
|
|
"gemini-2.5-flash-lite",
|
|
"gemini-2.5-flash",
|
|
"gemini-2.5-pro",
|
|
"gemini-2.0-flash",
|
|
"gemini-2.0-flash-lite",
|
|
]
|
|
|
|
CHEAP_IMAGEN_MODEL = "imagen-4.0-fast-generate-001"
|
|
IMAGEN_FALLBACK_MODELS = [CHEAP_IMAGEN_MODEL]
|
|
|
|
|
|
# ── Provider / model control ──────────────────────────────────
|
|
|
|
def set_provider(provider: str):
|
|
global _current_provider
|
|
if provider not in ("google", "openrouter"):
|
|
raise ValueError(f"Unknown provider: {provider}")
|
|
_current_provider = provider
|
|
_autosave()
|
|
|
|
|
|
def get_provider() -> str:
|
|
return _current_provider
|
|
|
|
|
|
def get_available_providers() -> list[str]:
|
|
return ["google", "openrouter"]
|
|
|
|
|
|
def set_model(model: str, provider: str | None = None):
|
|
global _google_model, _or_model
|
|
p = provider or _current_provider
|
|
if p == "google":
|
|
_google_model = model
|
|
else:
|
|
_or_model = model
|
|
_autosave()
|
|
|
|
|
|
def get_model() -> str:
|
|
return _google_model if _current_provider == "google" else _or_model
|
|
|
|
|
|
def get_model_for_provider(provider: str) -> str:
|
|
if provider == "google":
|
|
return _google_model
|
|
if provider == "openrouter":
|
|
return _or_model
|
|
raise ValueError(f"Unknown provider: {provider}")
|
|
|
|
|
|
def provider_status() -> dict:
|
|
return {
|
|
"provider": _current_provider,
|
|
"model": get_model(),
|
|
"google_model": _google_model,
|
|
"openrouter_model": _or_model,
|
|
"google_configured": bool(GOOGLE_API_KEY),
|
|
"openrouter_configured": bool(OPENROUTER_API_KEY),
|
|
"vision_provider": "google",
|
|
"image_provider": "google",
|
|
"last_fallback": _last_fallback or "none",
|
|
}
|
|
|
|
|
|
def import_provider_settings(data: dict):
|
|
global _current_provider, _google_model, _or_model
|
|
provider = data.get("provider")
|
|
if provider in ("google", "openrouter"):
|
|
_current_provider = provider
|
|
if isinstance(data.get("google_model"), str):
|
|
_google_model = data["google_model"]
|
|
if isinstance(data.get("openrouter_model"), str):
|
|
_or_model = data["openrouter_model"]
|
|
|
|
|
|
def export_provider_settings() -> dict:
|
|
return {
|
|
"provider": _current_provider,
|
|
"google_model": _google_model,
|
|
"openrouter_model": _or_model,
|
|
}
|
|
|
|
|
|
def _autosave():
|
|
try:
|
|
from services import settings as settings_store
|
|
settings_store.save()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def validate_model(model: str, provider: str | None = None) -> tuple[bool, str]:
|
|
p = provider or _current_provider
|
|
if p == "google":
|
|
if model in GOOGLE_TEXT_MODELS:
|
|
return True, ""
|
|
return False, "Unknown Google model. Use /models google to see supported examples."
|
|
|
|
if p == "openrouter":
|
|
models = await fetch_or_models()
|
|
if not models:
|
|
return True, "Could not verify OpenRouter model list; saved model anyway."
|
|
ids = {m.get("id") for m in models}
|
|
if model in ids:
|
|
return True, ""
|
|
return False, "Model was not found on OpenRouter. Use /models openrouter <keyword>."
|
|
|
|
return False, f"Unknown provider: {p}"
|
|
|
|
|
|
# ── Tool definitions ──────────────────────────────────────────
|
|
|
|
# Native google-genai tool declarations
|
|
_GC_TOOLS = [
|
|
gt.Tool(function_declarations=[
|
|
gt.FunctionDeclaration(
|
|
name="search_web",
|
|
description=(
|
|
"Search the internet for current information. "
|
|
"Use for today's date/time, news, prices, weather, recent events, "
|
|
"or any fact that might have changed."
|
|
),
|
|
parameters=gt.Schema(
|
|
type="object",
|
|
properties={"query": gt.Schema(type="string", description="Search query")},
|
|
required=["query"],
|
|
),
|
|
),
|
|
gt.FunctionDeclaration(
|
|
name="image_search",
|
|
description="Search the web for images. Use when user asks to find, show, or search for images/pictures.",
|
|
parameters=gt.Schema(
|
|
type="object",
|
|
properties={"query": gt.Schema(type="string", description="Image search query")},
|
|
required=["query"],
|
|
),
|
|
),
|
|
])
|
|
]
|
|
|
|
# OpenAI-format tool definitions (for OpenRouter)
|
|
_OAI_TOOLS = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "search_web",
|
|
"description": (
|
|
"Search the internet for current information. "
|
|
"Use for today's date/time, news, prices, weather, recent events."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "image_search",
|
|
"description": "Search the web for images. Use when user asks to find or show images/pictures.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
# ── Search functions ──────────────────────────────────────────
|
|
|
|
async def search_web(query: str, max_results: int = 5) -> str:
|
|
from utils import context
|
|
context.record_event("web_searches")
|
|
|
|
def _sync():
|
|
try:
|
|
from ddgs import DDGS
|
|
results = list(DDGS().text(query, max_results=max_results))
|
|
if not results:
|
|
return "No results found."
|
|
lines = ["Sources:"]
|
|
for idx, r in enumerate(results, start=1):
|
|
title = r.get("title", "Untitled")
|
|
href = r.get("href", "")
|
|
body = (r.get("body", "") or "")[:350]
|
|
lines.append(f"[{idx}] {title}\nURL: {href}\nSnippet: {body}")
|
|
lines.append("When answering, cite sources inline like [1] or [2].")
|
|
return "\n\n".join(lines)
|
|
except Exception as exc:
|
|
return f"Search error: {exc}"
|
|
return await asyncio.to_thread(_sync)
|
|
|
|
|
|
async def image_search(query: str, max_results: int = 6) -> tuple[list[str], str]:
|
|
from utils import context
|
|
context.record_event("image_searches")
|
|
|
|
def _sync():
|
|
try:
|
|
from ddgs import DDGS
|
|
results = list(DDGS().images(query, max_results=max_results))
|
|
urls = [r["image"] for r in results if r.get("image")]
|
|
return urls, f"Found {len(urls)} images for '{query}'."
|
|
except Exception as exc:
|
|
return [], f"Image search error: {exc}"
|
|
return await asyncio.to_thread(_sync)
|
|
|
|
|
|
# ── OpenRouter model listing ──────────────────────────────────
|
|
|
|
_or_models_cache: list[dict] = []
|
|
_or_models_fetched_at: float = 0.0
|
|
_OR_CACHE_TTL = 3600 # 1 hour
|
|
|
|
|
|
async def fetch_or_models(free_only: bool = False) -> list[dict]:
|
|
"""Fetch OpenRouter model list, cached for 1 hour. Returns list of model dicts."""
|
|
global _or_models_cache, _or_models_fetched_at
|
|
|
|
if _or_models_cache and (time.time() - _or_models_fetched_at) < _OR_CACHE_TTL:
|
|
models = _or_models_cache
|
|
else:
|
|
import httpx
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
resp = await client.get(
|
|
"https://openrouter.ai/api/v1/models",
|
|
headers={"Authorization": f"Bearer {OPENROUTER_API_KEY}"},
|
|
)
|
|
resp.raise_for_status()
|
|
models = resp.json().get("data", [])
|
|
_or_models_cache = models
|
|
_or_models_fetched_at = time.time()
|
|
except Exception as exc:
|
|
logger.warning("Failed to fetch OR models: %s", exc)
|
|
return []
|
|
|
|
if free_only:
|
|
models = [
|
|
m for m in models
|
|
if str(m.get("pricing", {}).get("prompt", "1")) == "0"
|
|
]
|
|
|
|
# Sort by name
|
|
models.sort(key=lambda m: m.get("id", ""))
|
|
return models
|
|
|
|
|
|
# ── Core chat: Google (native genai) ─────────────────────────
|
|
|
|
def _oai_to_gc_contents(messages: list[dict]) -> list[gt.Content]:
|
|
"""Convert OpenAI-format message history to google-genai Content list."""
|
|
contents = []
|
|
for m in messages:
|
|
role = m["role"]
|
|
if role == "system":
|
|
continue # handled via system_instruction
|
|
gc_role = "model" if role == "assistant" else "user"
|
|
content = m["content"]
|
|
if isinstance(content, str) and content:
|
|
contents.append(gt.Content(role=gc_role, parts=[gt.Part.from_text(text=content)]))
|
|
return contents
|
|
|
|
|
|
async def _chat_google(
|
|
messages: list[dict],
|
|
system_prompt: str,
|
|
max_tokens: int,
|
|
search_enabled: bool,
|
|
) -> tuple[str, list[str]]:
|
|
trimmed = messages[-MAX_API_CONTEXT:]
|
|
contents = _oai_to_gc_contents(trimmed)
|
|
|
|
config = gt.GenerateContentConfig(
|
|
system_instruction=system_prompt,
|
|
tools=_GC_TOOLS if search_enabled else None,
|
|
max_output_tokens=max_tokens,
|
|
)
|
|
|
|
found_image_urls: list[str] = []
|
|
|
|
while True:
|
|
response = await asyncio.to_thread(
|
|
lambda c=contents: _gc.models.generate_content(
|
|
model=_google_model,
|
|
contents=c,
|
|
config=config,
|
|
)
|
|
)
|
|
|
|
candidate = response.candidates[0]
|
|
func_calls = [p.function_call for p in candidate.content.parts if p.function_call]
|
|
|
|
if not func_calls:
|
|
return response.text or "", found_image_urls
|
|
|
|
# Append model's turn (with function calls) to history
|
|
contents.append(candidate.content)
|
|
|
|
# Execute all tool calls and collect responses
|
|
tool_response_parts: list[gt.Part] = []
|
|
for fc in func_calls:
|
|
args = dict(fc.args)
|
|
name = fc.name
|
|
if name == "search_web":
|
|
result = await search_web(args.get("query", ""))
|
|
logger.info("search_web: %s", args.get("query"))
|
|
elif name == "image_search":
|
|
urls, summary = await image_search(args.get("query", ""))
|
|
found_image_urls.extend(urls)
|
|
result = summary
|
|
logger.info("image_search: %s → %d urls", args.get("query"), len(urls))
|
|
else:
|
|
result = f"Unknown tool: {name}"
|
|
|
|
tool_response_parts.append(
|
|
gt.Part.from_function_response(name=name, response={"result": result})
|
|
)
|
|
|
|
contents.append(gt.Content(role="user", parts=tool_response_parts))
|
|
|
|
|
|
# ── Core chat: OpenRouter (OpenAI client) ─────────────────────
|
|
|
|
def _or_friendly_error(exc: Exception) -> str:
|
|
"""Convert OpenRouter API errors to actionable user-facing messages."""
|
|
msg = str(exc)
|
|
if isinstance(exc, APIStatusError):
|
|
code = exc.status_code
|
|
try:
|
|
body = exc.response.json()
|
|
err = body.get("error", {})
|
|
raw = err.get("metadata", {}).get("raw", "")
|
|
provider = err.get("metadata", {}).get("provider_name", "")
|
|
except Exception:
|
|
err, raw, provider = {}, "", ""
|
|
|
|
if code == 502:
|
|
hint = f" (provider: {provider})" if provider else ""
|
|
return (
|
|
f"❌ OpenRouter provider error{hint} — the upstream model is down or unavailable.\n"
|
|
f"Try a different model with /setmodel or switch provider with /setprovider google."
|
|
)
|
|
if code == 404 or "not found" in msg.lower():
|
|
return f"❌ Model not found on OpenRouter: `{_or_model}`\nCheck available models with /models."
|
|
if code == 429:
|
|
return "❌ OpenRouter rate limit reached. Try again in a moment."
|
|
if code == 402:
|
|
return "❌ OpenRouter credits exhausted. Top up at openrouter.ai."
|
|
if raw:
|
|
return f"❌ OpenRouter error ({code}): {raw[:200]}"
|
|
return f"❌ OpenRouter error: {msg[:300]}"
|
|
|
|
|
|
async def _chat_openrouter(
|
|
messages: list[dict],
|
|
system_prompt: str,
|
|
max_tokens: int,
|
|
search_enabled: bool,
|
|
) -> tuple[str, list[str]]:
|
|
trimmed = messages[-MAX_API_CONTEXT:]
|
|
full = [{"role": "system", "content": system_prompt}] + trimmed
|
|
|
|
tools = _OAI_TOOLS if search_enabled else []
|
|
kwargs: dict = {"model": _or_model, "messages": full, "max_tokens": max_tokens}
|
|
if tools:
|
|
kwargs["tools"] = tools
|
|
kwargs["tool_choice"] = "auto"
|
|
|
|
try:
|
|
response = await _or_client.chat.completions.create(**kwargs)
|
|
except APIStatusError as exc:
|
|
# Retry without tools if the model doesn't support them
|
|
if tools and exc.status_code in (400, 422):
|
|
logger.warning("OpenRouter rejected tools (%s), retrying without", exc.status_code)
|
|
kwargs.pop("tools", None)
|
|
kwargs.pop("tool_choice", None)
|
|
try:
|
|
response = await _or_client.chat.completions.create(**kwargs)
|
|
except APIStatusError as exc2:
|
|
raise RuntimeError(_or_friendly_error(exc2)) from exc2
|
|
else:
|
|
raise RuntimeError(_or_friendly_error(exc)) from exc
|
|
except Exception as exc:
|
|
raise RuntimeError(f"❌ OpenRouter error: {exc}") from exc
|
|
|
|
found_image_urls: list[str] = []
|
|
|
|
while response.choices[0].finish_reason == "tool_calls":
|
|
choice = response.choices[0]
|
|
tool_calls = choice.message.tool_calls
|
|
|
|
full.append({
|
|
"role": "assistant",
|
|
"content": choice.message.content,
|
|
"tool_calls": [
|
|
{"id": tc.id, "type": "function",
|
|
"function": {"name": tc.function.name, "arguments": tc.function.arguments}}
|
|
for tc in tool_calls
|
|
],
|
|
})
|
|
|
|
for tc in tool_calls:
|
|
args = json.loads(tc.function.arguments)
|
|
if tc.function.name == "search_web":
|
|
result = await search_web(args.get("query", ""))
|
|
logger.info("search_web: %s", args.get("query"))
|
|
elif tc.function.name == "image_search":
|
|
urls, summary = await image_search(args.get("query", ""))
|
|
found_image_urls.extend(urls)
|
|
result = summary
|
|
logger.info("image_search: %s → %d urls", args.get("query"), len(urls))
|
|
else:
|
|
result = f"Unknown tool: {tc.function.name}"
|
|
|
|
full.append({
|
|
"role": "tool",
|
|
"tool_call_id": tc.id,
|
|
"content": result,
|
|
"name": tc.function.name,
|
|
})
|
|
|
|
try:
|
|
response = await _or_client.chat.completions.create(
|
|
model=_or_model,
|
|
messages=full,
|
|
max_tokens=max_tokens,
|
|
)
|
|
except APIStatusError as exc:
|
|
raise RuntimeError(_or_friendly_error(exc)) from exc
|
|
|
|
return response.choices[0].message.content or "", found_image_urls
|
|
|
|
|
|
# ── Public chat dispatcher ────────────────────────────────────
|
|
|
|
def _sanitize(text: str) -> str:
|
|
from utils.formatter import strip_raw_tool_calls
|
|
return strip_raw_tool_calls(text)
|
|
|
|
|
|
def _limit_reply(text: str, max_tokens: int) -> str:
|
|
"""Silent hard cap after generation; providers can overshoot max_tokens."""
|
|
if not text:
|
|
return text
|
|
|
|
# Conservative approximation: one token is usually around 3-4 chars in English
|
|
# and varies by language. This keeps Telegram output visibly inside the user's cap.
|
|
max_chars = max(120, int(max_tokens * 3.5))
|
|
if len(text) <= max_chars:
|
|
return text
|
|
|
|
clipped = text[:max_chars].rstrip()
|
|
split_at = max(clipped.rfind("\n"), clipped.rfind(". "), clipped.rfind("؟ "), clipped.rfind("! "), clipped.rfind(" "))
|
|
if split_at > max_chars * 0.65:
|
|
clipped = clipped[:split_at].rstrip()
|
|
|
|
return clipped
|
|
|
|
|
|
def _with_limit_instruction(system_prompt: str, max_tokens: int) -> str:
|
|
return _with_runtime_instruction(
|
|
system_prompt,
|
|
max_tokens,
|
|
_current_provider,
|
|
get_model(),
|
|
)
|
|
|
|
|
|
def _with_runtime_instruction(
|
|
system_prompt: str,
|
|
max_tokens: int,
|
|
provider: str,
|
|
model: str,
|
|
) -> str:
|
|
creator_hint = _model_creator_hint(model)
|
|
return (
|
|
f"{system_prompt}\n\n"
|
|
f"RUNTIME METADATA: This bot is currently routing your replies through provider `{provider}` "
|
|
f"with model ID `{model}`. "
|
|
f"{creator_hint}"
|
|
"If the user asks what model/provider is being used, include this exact provider and model ID. "
|
|
"If the user asks who created/trained you, answer from the model family/developer when you know it; "
|
|
"do not infer creator from the hosting provider alone. If uncertain, say you only know the active model ID. "
|
|
"You may still answer naturally as the underlying model; this metadata only corrects runtime identity.\n\n"
|
|
f"REPLY LIMIT: You have at most {max_tokens} output tokens for the final answer. "
|
|
f"Plan briefly inside that budget and finish naturally before token {max_tokens}. "
|
|
"Do not mention the token limit, truncation, or that the answer was cut."
|
|
)
|
|
|
|
|
|
def _model_creator_hint(model: str) -> str:
|
|
lower = model.lower()
|
|
hints = [
|
|
("meta-llama/", "Model family/developer hint: Meta AI. "),
|
|
("llama", "Model family/developer hint: Meta AI. "),
|
|
("anthropic/", "Model family/developer hint: Anthropic. "),
|
|
("claude", "Model family/developer hint: Anthropic. "),
|
|
("openai/", "Model family/developer hint: OpenAI. "),
|
|
("gpt-", "Model family/developer hint: OpenAI. "),
|
|
("google/", "Model family/developer hint: Google. "),
|
|
("gemini", "Model family/developer hint: Google. "),
|
|
("mistralai/", "Model family/developer hint: Mistral AI. "),
|
|
("deepseek/", "Model family/developer hint: DeepSeek. "),
|
|
("qwen/", "Model family/developer hint: Alibaba/Qwen. "),
|
|
("x-ai/", "Model family/developer hint: xAI. "),
|
|
]
|
|
for needle, hint in hints:
|
|
if needle in lower:
|
|
return hint
|
|
return ""
|
|
|
|
|
|
async def chat(
|
|
messages: list[dict],
|
|
system_prompt: str,
|
|
max_tokens: int,
|
|
search_enabled: bool = True,
|
|
) -> tuple[str, list[str]]:
|
|
global _last_fallback
|
|
primary = _current_provider
|
|
fallback = "openrouter" if primary == "google" else "google"
|
|
primary_prompt = _with_runtime_instruction(
|
|
system_prompt,
|
|
max_tokens,
|
|
primary,
|
|
get_model_for_provider(primary),
|
|
)
|
|
fallback_prompt = _with_runtime_instruction(
|
|
system_prompt,
|
|
max_tokens,
|
|
fallback,
|
|
get_model_for_provider(fallback),
|
|
)
|
|
|
|
try:
|
|
if primary == "google":
|
|
reply, urls = await _chat_google(messages, primary_prompt, max_tokens, search_enabled)
|
|
else:
|
|
reply, urls = await _chat_openrouter(messages, primary_prompt, max_tokens, search_enabled)
|
|
except Exception as primary_exc:
|
|
if fallback == "openrouter" and not OPENROUTER_API_KEY:
|
|
raise
|
|
if fallback == "google" and not GOOGLE_API_KEY:
|
|
raise
|
|
|
|
from utils import context
|
|
context.record_event("provider_fallbacks")
|
|
logger.warning("Primary provider %s failed; trying %s: %s", primary, fallback, primary_exc)
|
|
try:
|
|
if fallback == "google":
|
|
reply, urls = await _chat_google(messages, fallback_prompt, max_tokens, search_enabled)
|
|
else:
|
|
reply, urls = await _chat_openrouter(messages, fallback_prompt, max_tokens, search_enabled)
|
|
_last_fallback = f"{primary} -> {fallback}"
|
|
except Exception as fallback_exc:
|
|
raise RuntimeError(
|
|
f"Primary provider failed: {primary_exc}\nFallback provider failed: {fallback_exc}"
|
|
) from fallback_exc
|
|
return _limit_reply(_sanitize(reply), max_tokens), urls
|
|
|
|
|
|
# ── Vision (native google-genai) ─────────────────────────────
|
|
|
|
async def chat_with_image(
|
|
messages: list[dict],
|
|
system_prompt: str,
|
|
max_tokens: int,
|
|
image_bytes: bytes,
|
|
mime_type: str = "image/jpeg",
|
|
caption: str = "",
|
|
analysis_mode: str = "auto",
|
|
) -> str:
|
|
def _sync():
|
|
history_lines = []
|
|
for m in messages[-6:]:
|
|
c = m["content"]
|
|
history_lines.append(
|
|
f"{m['role'].capitalize()}: {c if isinstance(c, str) else '[media]'}"
|
|
)
|
|
ctx = "\n".join(history_lines)
|
|
|
|
prompt = _with_limit_instruction(system_prompt, max_tokens)
|
|
if ctx:
|
|
prompt += f"\n\nRecent conversation:\n{ctx}"
|
|
mode_prompts = {
|
|
"auto": (
|
|
"Analyze the image carefully. Identify visible subjects, important details, "
|
|
"text/OCR, context, and answer the user's caption if provided."
|
|
),
|
|
"describe": "Describe the image clearly and include the most important visual details.",
|
|
"ocr": "Extract all readable text from the image, preserving line breaks when possible, then summarize it.",
|
|
"objects": "List the main objects, people, places, and notable visual attributes in the image.",
|
|
"question": "Answer the user's question about the image using only visible evidence.",
|
|
}
|
|
instruction = mode_prompts.get(analysis_mode, mode_prompts["auto"])
|
|
if caption:
|
|
prompt += f"\n\nImage task: {instruction}\nUser caption/question: {caption}"
|
|
else:
|
|
prompt += f"\n\nImage task: {instruction}"
|
|
|
|
resp = _gc.models.generate_content(
|
|
model=_google_model,
|
|
contents=[
|
|
gt.Content(role="user", parts=[
|
|
gt.Part.from_bytes(data=image_bytes, mime_type=mime_type),
|
|
gt.Part.from_text(text=prompt),
|
|
])
|
|
],
|
|
config=gt.GenerateContentConfig(max_output_tokens=max_tokens),
|
|
)
|
|
return resp.text or ""
|
|
return _limit_reply(_sanitize(await asyncio.to_thread(_sync)), max_tokens)
|
|
|
|
|
|
async def chat_with_images(
|
|
messages: list[dict],
|
|
system_prompt: str,
|
|
max_tokens: int,
|
|
images: list[tuple[bytes, str]],
|
|
caption: str = "",
|
|
analysis_mode: str = "auto",
|
|
) -> str:
|
|
def _sync():
|
|
history_lines = []
|
|
for m in messages[-6:]:
|
|
c = m["content"]
|
|
history_lines.append(
|
|
f"{m['role'].capitalize()}: {c if isinstance(c, str) else '[media]'}"
|
|
)
|
|
ctx = "\n".join(history_lines)
|
|
|
|
prompt = _with_limit_instruction(system_prompt, max_tokens)
|
|
if ctx:
|
|
prompt += f"\n\nRecent conversation:\n{ctx}"
|
|
prompt += (
|
|
f"\n\nAnalyze these {len(images)} images together. Compare them when useful, "
|
|
"identify relationships, extract visible text, and answer the user's caption/question."
|
|
)
|
|
if analysis_mode == "ocr":
|
|
prompt += " Focus on OCR and preserve readable text by image."
|
|
if caption:
|
|
prompt += f"\nUser caption/question: {caption}"
|
|
|
|
parts = []
|
|
for idx, (image_bytes, mime_type) in enumerate(images, start=1):
|
|
parts.append(gt.Part.from_text(text=f"Image {idx}:"))
|
|
parts.append(gt.Part.from_bytes(data=image_bytes, mime_type=mime_type))
|
|
parts.append(gt.Part.from_text(text=prompt))
|
|
|
|
resp = _gc.models.generate_content(
|
|
model=_google_model,
|
|
contents=[gt.Content(role="user", parts=parts)],
|
|
config=gt.GenerateContentConfig(max_output_tokens=max_tokens),
|
|
)
|
|
return resp.text or ""
|
|
|
|
return _limit_reply(_sanitize(await asyncio.to_thread(_sync)), max_tokens)
|
|
|
|
|
|
# ── Image generation (cheap Imagen 4 Fast — always Google) ───
|
|
|
|
async def generate_image(prompt: str) -> bytes:
|
|
# Keep image generation on the cheapest Imagen path. Do not fall back to
|
|
# standard/ultra Imagen here; those cost more per image.
|
|
models = list(dict.fromkeys([CHEAP_IMAGEN_MODEL, *IMAGEN_FALLBACK_MODELS]))
|
|
errors: list[str] = []
|
|
|
|
def _empty_reason(resp) -> str:
|
|
for attr in ("filtered_reason", "prompt_feedback", "safety_ratings"):
|
|
value = getattr(resp, attr, None)
|
|
if value:
|
|
return f"{attr}: {value}"
|
|
try:
|
|
dumped = resp.model_dump(exclude_none=True)
|
|
if dumped:
|
|
return str(dumped)[:500]
|
|
except Exception:
|
|
pass
|
|
return "empty generated_images"
|
|
|
|
def _sync():
|
|
for model in models:
|
|
try:
|
|
resp = _gc.models.generate_images(
|
|
model=model,
|
|
prompt=prompt,
|
|
config=gt.GenerateImagesConfig(number_of_images=1),
|
|
)
|
|
imgs = resp.generated_images or []
|
|
for item in imgs:
|
|
image = getattr(item, "image", None)
|
|
image_bytes = getattr(image, "image_bytes", None)
|
|
if image_bytes:
|
|
return image_bytes
|
|
errors.append(f"{model}: {_empty_reason(resp)}")
|
|
except Exception as exc:
|
|
errors.append(f"{model}: {exc}")
|
|
|
|
detail = "; ".join(errors[-3:]) if errors else "unknown empty response"
|
|
raise RuntimeError(
|
|
"No image returned by cheap Imagen 4 Fast. The prompt may have been "
|
|
f"filtered or the model returned empty output. Details: {detail}"
|
|
)
|
|
return await asyncio.to_thread(_sync)
|
|
|
|
|
|
# ── Inline quick reply ────────────────────────────────────────
|
|
|
|
async def quick_reply(query: str, system_prompt: str, max_tokens: int = 250) -> str:
|
|
google_prompt = _with_runtime_instruction(
|
|
system_prompt,
|
|
max_tokens,
|
|
"google",
|
|
get_model_for_provider("google"),
|
|
)
|
|
openrouter_prompt = _with_runtime_instruction(
|
|
system_prompt,
|
|
max_tokens,
|
|
"openrouter",
|
|
get_model_for_provider("openrouter"),
|
|
)
|
|
|
|
async def _quick_google() -> str:
|
|
def _sync():
|
|
resp = _gc.models.generate_content(
|
|
model=_google_model,
|
|
contents=[gt.Content(role="user", parts=[gt.Part.from_text(text=query)])],
|
|
config=gt.GenerateContentConfig(
|
|
system_instruction=google_prompt,
|
|
max_output_tokens=max_tokens,
|
|
),
|
|
)
|
|
return resp.text or ""
|
|
return await asyncio.to_thread(_sync)
|
|
|
|
async def _quick_openrouter() -> str:
|
|
response = await _or_client.chat.completions.create(
|
|
model=_or_model,
|
|
messages=[
|
|
{"role": "system", "content": openrouter_prompt},
|
|
{"role": "user", "content": query},
|
|
],
|
|
max_tokens=max_tokens,
|
|
)
|
|
return response.choices[0].message.content or ""
|
|
|
|
try:
|
|
reply = await (_quick_google() if _current_provider == "google" else _quick_openrouter())
|
|
except Exception:
|
|
if _current_provider == "google" and OPENROUTER_API_KEY:
|
|
reply = await _quick_openrouter()
|
|
elif _current_provider == "openrouter" and GOOGLE_API_KEY:
|
|
reply = await _quick_google()
|
|
else:
|
|
raise
|
|
return _limit_reply(_sanitize(reply), max_tokens)
|