From 8585eeb7e158bb6cfad0295f573d40274a579ee1 Mon Sep 17 00:00:00 2001 From: arsam Date: Fri, 18 Sep 2026 22:41:16 +0330 Subject: [PATCH] Add AI provider integration service --- services/ai_client.py | 805 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 805 insertions(+) create mode 100644 services/ai_client.py diff --git a/services/ai_client.py b/services/ai_client.py new file mode 100644 index 0000000..071e207 --- /dev/null +++ b/services/ai_client.py @@ -0,0 +1,805 @@ +""" +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 ." + + 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)