"""Assistant module for Togen LLM chat (DVI-529). Provides a pluggable provider registry (ASSISTANT_PROVIDERS), settings persistence for global admin config and per-user connection overrides, an availability gate (is_assistant_available), a chat proxy, and a vision-based field-extraction helper used by the Claude-vision OCR backend. Architecture note (D1): ASSISTANT_PROVIDERS is intentionally open-ended so a future full-MCP provider can be registered here without touching routes or UI. The dispatch branch in chat() keys on provider id; MCP providers would add their own branch. """ from __future__ import annotations import json import os import re import stat import threading import uuid as _uuid from datetime import datetime, timezone from pathlib import Path from typing import Any # --------------------------------------------------------------------------- # Provider registry # --------------------------------------------------------------------------- ASSISTANT_PROVIDERS: dict[str, dict] = { "anthropic": { "id": "anthropic", "label": "Claude (Anthropic API)", "models": [ {"id": "claude-opus-4-7", "label": "Claude Opus 4.7"}, {"id": "claude-sonnet-4-6", "label": "Claude Sonnet 4.6"}, {"id": "claude-haiku-4-5-20251001", "label": "Claude Haiku 4.5"}, ], "default_model": "claude-sonnet-4-6", "docs_url": "https://console.anthropic.com/settings/keys", }, "google": { "id": "google", "label": "Gemini (Google AI)", "models": [ {"id": "gemini-2.5-pro", "label": "Gemini 2.5 Pro"}, {"id": "gemini-2.5-flash", "label": "Gemini 2.5 Flash"}, {"id": "gemini-2.0-flash", "label": "Gemini 2.0 Flash"}, ], "default_model": "gemini-2.5-flash", "docs_url": "https://aistudio.google.com/app/apikey", }, "openai": { "id": "openai", "label": "OpenAI (GPT)", "models": [ {"id": "gpt-4.1", "label": "GPT-4.1"}, {"id": "gpt-4o", "label": "GPT-4o"}, {"id": "o4-mini", "label": "o4-mini"}, ], "default_model": "gpt-4o", "docs_url": "https://platform.openai.com/api-keys", }, } DEFAULT_PROVIDER = "anthropic" DEFAULT_MODEL = "claude-sonnet-4-6" def _provider_for_model(model_id: str) -> str: """Return the provider id that owns model_id, defaulting to DEFAULT_PROVIDER.""" for prov_id, prov in ASSISTANT_PROVIDERS.items(): for m in prov.get("models", []): if m["id"] == model_id: return prov_id return DEFAULT_PROVIDER def _is_known_model(model_id: str) -> bool: """Return True if model_id is a registered model of any provider.""" for prov in ASSISTANT_PROVIDERS.values(): for m in prov.get("models", []): if m["id"] == model_id: return True return False # --------------------------------------------------------------------------- # Settings files # --------------------------------------------------------------------------- ASSISTANT_SETTINGS_FILE = Path( os.environ.get( "ASSISTANT_SETTINGS_FILE", Path(__file__).resolve().parent / "assistant_settings.json", ) ) DEFAULT_ASSISTANT_SETTINGS: dict[str, Any] = { "enabled": False, "accessible_to_users": False, "default_context_md": "", # org-level api_key is optional; absent means org key not configured # api_keys: {anthropic: "...", google: "..."} — per-provider (legacy api_key → anthropic) } def load_assistant_settings() -> dict: """Return global assistant settings, falling back to defaults.""" settings = dict(DEFAULT_ASSISTANT_SETTINGS) if ASSISTANT_SETTINGS_FILE.is_file(): try: saved = json.loads(ASSISTANT_SETTINGS_FILE.read_text()) if isinstance(saved, dict): if isinstance(saved.get("enabled"), bool): settings["enabled"] = saved["enabled"] if isinstance(saved.get("accessible_to_users"), bool): settings["accessible_to_users"] = saved["accessible_to_users"] if isinstance(saved.get("default_context_md"), str): settings["default_context_md"] = saved["default_context_md"] if saved.get("api_key"): settings["api_key"] = saved["api_key"] if isinstance(saved.get("api_keys"), dict): settings["api_keys"] = saved["api_keys"] if isinstance(saved.get("provider_enabled"), dict): settings["provider_enabled"] = saved["provider_enabled"] if saved.get("default_model"): settings["default_model"] = saved["default_model"] if saved.get("ocr_model"): settings["ocr_model"] = saved["ocr_model"] except (json.JSONDecodeError, OSError): pass return settings def save_assistant_settings(settings: dict) -> None: """Persist global assistant settings. File is 0600 (key-bearing).""" ASSISTANT_SETTINGS_FILE.write_text(json.dumps(settings, indent=2)) ASSISTANT_SETTINGS_FILE.chmod(stat.S_IRUSR | stat.S_IWUSR) # --------------------------------------------------------------------------- # Per-user connection helpers # --------------------------------------------------------------------------- def get_user_assistant(user_settings: dict) -> dict: """Extract the assistant sub-dict from user settings (may be empty).""" return dict(user_settings.get("assistant") or {}) def set_user_assistant(user_settings: dict, assistant: dict) -> dict: """Return a copy of user_settings with the assistant sub-dict updated.""" updated = dict(user_settings) updated["assistant"] = assistant return updated # --------------------------------------------------------------------------- # Availability gate (plan §3.6) # --------------------------------------------------------------------------- def is_assistant_available(user_email: str, is_admin: bool, user_settings: dict) -> bool: """Return True if the assistant is available for this user. Admins: always available when the feature is enabled. Normal users: available when they have their own key (any provider), OR when the org key is configured for any provider and accessible_to_users is True. """ if is_admin: return True global_cfg = load_assistant_settings() if not global_cfg.get("enabled"): return False ua = get_user_assistant(user_settings) user_keys = ua.get("keys") or {} if any(user_keys.values()) or ua.get("api_key"): return True if global_cfg.get("accessible_to_users"): org_keys = global_cfg.get("api_keys") or {} if any(org_keys.values()) or global_cfg.get("api_key"): return True return False def is_assistant_connected(user_email: str, is_admin: bool, user_settings: dict) -> bool: """Return True only when the assistant is both entitled AND has a usable key. `is_assistant_available` is an *entitlement* gate — admins pass it unconditionally and org-key access passes it on the flag alone. That is not enough to actually run a request: without an effective key, chat fails with "No API key configured". This stricter check is what the topbar indicator (green vs red) and the chat gate use, so an entitled-but-unconfigured user (e.g. an admin who never set a key) sees red and the "configure your key" prompt instead of an input box that silently errors. """ if not is_assistant_available(user_email, is_admin, user_settings): return False return _has_any_effective_key(user_settings) # --------------------------------------------------------------------------- # Effective key resolution (per-provider; legacy api_key → anthropic) # --------------------------------------------------------------------------- def _effective_key(user_settings: dict, provider_id: str = DEFAULT_PROVIDER) -> str | None: """Return the best API key for provider_id (user key beats org key). Checks new per-provider `keys` dict first; falls back to legacy single `api_key` field when provider_id is 'anthropic' for back-compat. """ ua = get_user_assistant(user_settings) user_keys = ua.get("keys") or {} if user_keys.get(provider_id): return user_keys[provider_id] if provider_id == "anthropic" and ua.get("api_key"): return ua["api_key"] global_cfg = load_assistant_settings() org_keys = global_cfg.get("api_keys") or {} if org_keys.get(provider_id): return org_keys[provider_id] if provider_id == "anthropic" and global_cfg.get("api_key"): return global_cfg["api_key"] return None def _has_any_effective_key(user_settings: dict) -> bool: """Return True if there is at least one usable key for any registered provider.""" return any(_effective_key(user_settings, pid) for pid in ASSISTANT_PROVIDERS) def _effective_context_md(user_settings: dict) -> str: """Return the context markdown to use (user override beats admin default).""" ua = get_user_assistant(user_settings) if ua.get("context_md") is not None: return ua["context_md"] global_cfg = load_assistant_settings() return global_cfg.get("default_context_md") or "" # --------------------------------------------------------------------------- # Chat proxy (non-streaming, D4) # --------------------------------------------------------------------------- def _resolve_chat_model(user_settings: dict, model_override: str | None = None) -> tuple[str, str, str | None]: """Return (model_id, provider_id, key) for a plain-/agentic-chat request. A per-invocation `model_override` (when it names a registered model) beats the user's saved default_model; otherwise falls back to the caller default. The provider — and therefore the key — is always derived from the resolved model, so selecting e.g. a Gemini model routes to the Google key. """ model = None if model_override: candidate = str(model_override).strip() if _is_known_model(candidate): model = candidate if not model: ua = get_user_assistant(user_settings) model = ua.get("default_model") or ua.get("model") or DEFAULT_MODEL provider_id = _provider_for_model(model) key = _effective_key(user_settings, provider_id) return model, provider_id, key def chat(messages: list[dict], system: str, *, user_settings: dict, model: str | None = None) -> str: """Send messages to the configured provider and return the reply text. Resolves provider from `model` when given (per-invocation override), else from the user's selected default model; dispatches to the appropriate backend (anthropic/google/openai). Raises AssistantUnavailable if no key is configured for the resolved provider. Raises AssistantError on API failure. """ model, provider_id, key = _resolve_chat_model(user_settings, model) if not key: raise AssistantUnavailable("No API key configured.") if provider_id == "google": return _chat_google(messages, system, model=model, key=key) if provider_id == "openai": return _chat_openai(messages, system, model=model, key=key) if provider_id == "anthropic": return _chat_anthropic(messages, system, model=model, key=key) # Registered provider without a chat backend yet. raise AssistantUnavailable( f"Chat is not yet available for provider '{provider_id}'." ) def _chat_anthropic( messages: list[dict], system: str, *, model: str, key: str ) -> str: try: import anthropic as _anthropic except ImportError as exc: raise AssistantUnavailable("anthropic package is not installed.") from exc client = _anthropic.Anthropic(api_key=key) try: response = client.messages.create( model=model, max_tokens=4096, system=system or "", messages=messages, ) return response.content[0].text except _anthropic.APIError as exc: raise AssistantError(str(exc)) from exc def _chat_google( messages: list[dict], system: str, *, model: str, key: str ) -> str: """Send a plain (non-tool) chat request to Gemini and return reply text.""" try: from google import genai as _genai from google.genai import types as _gtypes except ImportError as exc: raise AssistantUnavailable("google-genai package is not installed.") from exc client = _genai.Client(api_key=key) contents = [] for msg in messages: role = msg.get("role", "user") if role == "assistant": role = "model" contents.append( _gtypes.Content(role=role, parts=[_gtypes.Part.from_text(str(msg.get("content", "")))]) ) try: response = client.models.generate_content( model=model, contents=contents, config=_gtypes.GenerateContentConfig( system_instruction=system or None, max_output_tokens=4096, ), ) return response.text except Exception as exc: raise AssistantError(str(exc)) from exc def _chat_openai( messages: list[dict], system: str, *, model: str, key: str ) -> str: """Send a plain (non-tool) chat request to OpenAI Chat Completions and return reply text.""" try: from openai import OpenAI as _OpenAI except ImportError as exc: raise AssistantUnavailable("openai package is not installed.") from exc client = _OpenAI(api_key=key) oa_messages: list[dict] = [] if system: oa_messages.append({"role": "system", "content": system}) for msg in messages: oa_messages.append({ "role": msg.get("role", "user"), "content": str(msg.get("content", "")), }) try: response = client.chat.completions.create( model=model, max_tokens=4096, messages=oa_messages, ) return response.choices[0].message.content or "" except Exception as exc: raise AssistantError(str(exc)) from exc # --------------------------------------------------------------------------- # Vision field extraction (OCR Claude-vision backend) # --------------------------------------------------------------------------- def _resolve_ocr_model() -> tuple[str, str, str | None]: """Resolve (model_id, provider_id, key) for title-block OCR. Uses the admin ``ocr_model`` setting (a bare registered model id) plus the org-level key for that model's provider. Defaults to Anthropic (DEFAULT_MODEL + org anthropic key) when the setting is unset or names an unknown model, so existing OCR behavior is unchanged. """ cfg = load_assistant_settings() raw = str(cfg.get("ocr_model") or "").strip() if raw and _is_known_model(raw): model_id = raw provider_id = _provider_for_model(raw) else: model_id = DEFAULT_MODEL provider_id = DEFAULT_PROVIDER key = _effective_key_system_only(provider_id) return model_id, provider_id, key def extract_fields(image_bytes: bytes, fields: list[str]) -> dict[str, str]: """Use a vision model to extract named fields from an image. image_bytes: raw PNG/JPEG image bytes. fields: list of field names to locate in the image (e.g. ["customer_name", "job_no"]). Returns a dict mapping field names to extracted string values (empty string when not found). Raises AssistantUnavailable or AssistantError on failure. Provider/model come from the admin ``ocr_model`` setting (default Anthropic org key); the request is routed to the matching provider's vision path. """ model_id, provider_id, key = _resolve_ocr_model() if not key: raise AssistantUnavailable( f"No org API key configured for {provider_id} vision OCR." ) fields_list = "\n".join(f"- {f}" for f in fields) if fields == ["full_text"]: prompt = ( "You are an OCR assistant. Transcribe ALL text visible in this image, " "preserving line breaks and layout as closely as possible. " "Return ONLY a JSON object with a single key 'full_text' whose value " "is the full transcribed text. No other keys, no markdown, no commentary." ) else: prompt = ( "You are an OCR assistant. Extract the following fields from the image " "of an engineering drawing title block. Return ONLY a JSON object mapping " f"field names to their string values. Fields to extract:\n{fields_list}\n\n" "If a field is not present or not legible, use an empty string. " "Return only the JSON object, no other text." ) raw = _ocr_vision_call(provider_id, model_id, key, image_bytes, prompt) try: raw = raw.strip() # Strip markdown code fences if the model added them if raw.startswith("```"): raw = raw.split("```")[1] if raw.startswith("json"): raw = raw[4:] result = json.loads(raw) return {f: str(result.get(f, "")) for f in fields} except (json.JSONDecodeError, KeyError, IndexError, AttributeError) as exc: raise AssistantError(f"Unexpected response format from {provider_id} vision: {exc}") from exc # Header keys returned by extract_batch_ticket (DVI-1374). Mirrors the QC # New Batch Ticket form fields; per-material actuals ride in a separate array. _BATCH_TICKET_KEYS = ( "batch_date", "batch_number", "mixture", "yards", "rock_moisture", "sand_moisture", "initials", "notes", ) def extract_batch_ticket(image_bytes: bytes, material_names: list[str]) -> dict: """Vision-extract a concrete QC batch ticket into structured fields. image_bytes: raw PNG/JPEG of one batch-ticket page. material_names: the QC material vocabulary, so the model can match the per-material weighed amounts to known materials. Returns ``{"fields": {batch_date, batch_number, mixture, yards, rock_moisture, sand_moisture, initials, notes}, "materials": [{"name", "actual"}, ...]}`` — best effort; a missing/illegible value is an empty string and materials may be empty. Raises AssistantUnavailable/AssistantError on provider failure (callers treat OCR as best-effort prefill). Provider/model come from the admin ``ocr_model`` setting (default Anthropic org key), reusing the title-block OCR resolution + per-provider vision path. """ model_id, provider_id, key = _resolve_ocr_model() if not key: raise AssistantUnavailable( f"No org API key configured for {provider_id} vision OCR." ) mats = "\n".join(f"- {m}" for m in material_names) or "- (none configured)" prompt = ( "You are reading a concrete Quality-Control batch ticket. Extract the " "batch header and the per-material actual (weighed/dispensed) amounts. " "Return ONLY a JSON object with these keys:\n" " batch_date (YYYY-MM-DD if you can determine it, else the date text),\n" " batch_number (the batch or ticket number),\n" " mixture (the mix design name or code),\n" " yards (cubic yards produced, number only),\n" " rock_moisture (rock/stone moisture percent, number only),\n" " sand_moisture (sand moisture percent, number only),\n" " initials (the operator initials),\n" " notes (any free-text notes),\n" " materials (an array of objects {\"name\": , \"actual\": " "} for each material weighed).\n" "Match each material name to this list where possible:\n" + mats + "\n\n" "Use an empty string (or omit) any field that is absent or not legible. " "Return only the JSON object — no markdown fences, no commentary." ) raw = _ocr_vision_call(provider_id, model_id, key, image_bytes, prompt) try: raw = raw.strip() if raw.startswith("```"): raw = raw.split("```")[1] if raw.startswith("json"): raw = raw[4:] data = json.loads(raw) if not isinstance(data, dict): raise ValueError("expected a JSON object") except (json.JSONDecodeError, ValueError, IndexError, AttributeError) as exc: raise AssistantError( f"Unexpected response format from {provider_id} vision: {exc}" ) from exc fields = {k: str(data.get(k, "") or "") for k in _BATCH_TICKET_KEYS} materials = [] for item in (data.get("materials") or []): if not isinstance(item, dict): continue name = str(item.get("name", "") or "").strip() if not name: continue materials.append({"name": name, "actual": str(item.get("actual", "") or "").strip()}) return {"fields": fields, "materials": materials} def _ocr_vision_call( provider_id: str, model_id: str, key: str, image_bytes: bytes, prompt: str ) -> str: """Send a single image + prompt to the given provider and return raw text. Images are sent at full resolution (no downscale) to preserve fine title-block text for OCR. Mirrors the per-provider vision message shapes used by interpret_images (DVI-953). """ if provider_id == "google": return _ocr_google(image_bytes, prompt, model=model_id, key=key) if provider_id == "openai": return _ocr_openai(image_bytes, prompt, model=model_id, key=key) if provider_id != "anthropic": raise AssistantUnavailable( f"Vision OCR is not configured for provider '{provider_id}'." ) return _ocr_anthropic(image_bytes, prompt, model=model_id, key=key) def _ocr_anthropic(image_bytes: bytes, prompt: str, *, model: str, key: str) -> str: import base64 try: import anthropic as _anthropic except ImportError as exc: raise AssistantUnavailable("anthropic package is not installed.") from exc client = _anthropic.Anthropic(api_key=key) b64 = base64.standard_b64encode(image_bytes).decode() try: response = client.messages.create( model=model, max_tokens=1024, messages=[ { "role": "user", "content": [ { "type": "image", "source": { "type": "base64", "media_type": "image/png", "data": b64, }, }, {"type": "text", "text": prompt}, ], } ], ) return response.content[0].text except _anthropic.APIError as exc: raise AssistantError(str(exc)) from exc def _ocr_openai(image_bytes: bytes, prompt: str, *, model: str, key: str) -> str: import base64 try: from openai import OpenAI as _OpenAI except ImportError as exc: raise AssistantUnavailable("openai package is not installed.") from exc client = _OpenAI(api_key=key) b64 = base64.standard_b64encode(image_bytes).decode() try: response = client.chat.completions.create( model=model, max_tokens=1024, messages=[ { "role": "user", "content": [ { "type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}, }, {"type": "text", "text": prompt}, ], } ], ) return response.choices[0].message.content or "" except Exception as exc: raise AssistantError(str(exc)) from exc def _ocr_google(image_bytes: bytes, prompt: str, *, model: str, key: str) -> str: try: from google import genai as _genai from google.genai import types as _gtypes except ImportError as exc: raise AssistantUnavailable("google-genai package is not installed.") from exc client = _genai.Client(api_key=key) try: response = client.models.generate_content( model=model, contents=[ _gtypes.Content( role="user", parts=[ _gtypes.Part.from_bytes(data=image_bytes, mime_type="image/png"), _gtypes.Part.from_text(prompt), ], ) ], config=_gtypes.GenerateContentConfig(max_output_tokens=1024), ) return response.text except Exception as exc: raise AssistantError(str(exc)) from exc # --------------------------------------------------------------------------- # Vision image interpretation (DVI-867 Phase 1) # --------------------------------------------------------------------------- _INTERPRET_MAX_LONG_EDGE = 1568 # cap before sending to LLM (token budget) _INTERPRET_JPEG_QUALITY = 85 def _downscale_image(image_bytes: bytes) -> bytes: """Re-encode image to JPEG, capped at _INTERPRET_MAX_LONG_EDGE on the long edge. Accepts any PIL-readable format; returns JPEG bytes. Falls back to the original bytes if Pillow is unavailable or the image is already small enough. """ try: from io import BytesIO from PIL import Image img = Image.open(BytesIO(image_bytes)) w, h = img.size long_edge = max(w, h) if long_edge > _INTERPRET_MAX_LONG_EDGE: scale = _INTERPRET_MAX_LONG_EDGE / long_edge img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS) if img.mode not in ("RGB", "L"): img = img.convert("RGB") buf = BytesIO() img.save(buf, format="JPEG", quality=_INTERPRET_JPEG_QUALITY) return buf.getvalue() except Exception: return image_bytes def interpret_image( image_bytes: bytes, prompt: str, history: list[dict] | None = None, *, user_settings: dict, model: str | None = None, ) -> str: """Interpret a single image using the configured Vision assistant model. Delegates to interpret_images with a single-element list; behavior is unchanged. """ return interpret_images([image_bytes], prompt, history, user_settings=user_settings, model=model) def _resolve_vision_model(user_settings: dict, raw_model: str) -> tuple[str, str, str | None]: """Resolve (model_id, provider_id, key) from a ``source:modelId`` string. Shared by the saved Vision default and the per-request override. ``system:`` forces the org key, ``user:`` forces the user's own key, otherwise the default provider/key. Falls back to the default model/provider when unparseable and always tries a last-resort effective key so a valid model never fails purely on prefix mismatch. """ raw_model = raw_model or "" if raw_model.startswith("system:"): model_id = raw_model[len("system:"):] provider_id = _provider_for_model(model_id) key = _effective_key_system_only(provider_id) elif raw_model.startswith("user:"): model_id = raw_model[len("user:"):] provider_id = _provider_for_model(model_id) ua = get_user_assistant(user_settings) user_keys = ua.get("keys") or {} key = user_keys.get(provider_id) or (ua.get("api_key") if provider_id == "anthropic" else None) else: model_id = DEFAULT_MODEL provider_id = DEFAULT_PROVIDER key = _effective_key(user_settings, provider_id) if not model_id: model_id = DEFAULT_MODEL provider_id = DEFAULT_PROVIDER if not key: key = _effective_key(user_settings, provider_id) return model_id, provider_id, key def interpret_images( images: list[bytes], prompt: str, history: list[dict] | None = None, *, user_settings: dict, model: str | None = None, ) -> str: """Interpret an ordered sequence of images using the configured Vision assistant model. images: ordered list of raw image bytes (any PIL-readable format). Each image is preceded by a 'Frame k/N' text label so the model can reference frames. prompt: user question / instruction for this turn. history: prior conversation turns [{role, content}] — images are pinned on the first user turn as context for multi-turn sessions. user_settings: caller's settings dict (for key + model resolution). Raises AssistantUnavailable if no key is configured or provider not supported. Raises AssistantError on API failure. """ import base64 if not images: raise AssistantUnavailable("No images provided for interpretation.") # --- model + key resolution --- # A per-request `model` override (same "source:modelId" form as the saved # selection) beats vision.assistant_model when it names a registered model. raw_model = "" if model: candidate = str(model).strip() if _is_known_model(candidate.split(":", 1)[-1]): raw_model = candidate if not raw_model: raw_model = (user_settings.get("vision") or {}).get("assistant_model", "") model_id, provider_id, key = _resolve_vision_model(user_settings, raw_model) if not key: raise AssistantUnavailable("No API key configured for Vision interpretation.") if provider_id == "google": return _interpret_images_google(images, prompt, history, model=model_id, key=key) if provider_id == "openai": return _interpret_images_openai(images, prompt, history, model=model_id, key=key) if provider_id != "anthropic": raise AssistantUnavailable( f"Vision interpretation is not yet configured for provider '{provider_id}'." ) try: import anthropic as _anthropic except ImportError as exc: raise AssistantUnavailable("anthropic package is not installed.") from exc # --- build interleaved label + image blocks --- n = len(images) frame_content: list[dict] = [] for k, img_bytes in enumerate(images, start=1): scaled = _downscale_image(img_bytes) b64 = base64.standard_b64encode(scaled).decode() if n > 1: frame_content.append({"type": "text", "text": f"Frame {k}/{n}"}) frame_content.append({ "type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": b64}, }) messages: list[dict] = [] if history: first_user_injected = False for turn in history: if turn.get("role") == "user" and not first_user_injected: messages.append({ "role": "user", "content": frame_content + [{"type": "text", "text": str(turn.get("content", ""))}], }) first_user_injected = True else: messages.append({"role": turn["role"], "content": str(turn.get("content", ""))}) messages.append({"role": "user", "content": prompt}) else: messages.append({ "role": "user", "content": frame_content + [{"type": "text", "text": prompt}], }) client = _anthropic.Anthropic(api_key=key) try: response = client.messages.create( model=model_id, max_tokens=2048, messages=messages, ) return response.content[0].text except _anthropic.APIError as exc: raise AssistantError(str(exc)) from exc def _interpret_images_google( images: list[bytes], prompt: str, history: list[dict] | None, *, model: str, key: str, ) -> str: """Interpret an image sequence using Gemini vision (inline image parts).""" try: from google import genai as _genai from google.genai import types as _gtypes except ImportError as exc: raise AssistantUnavailable("google-genai package is not installed.") from exc client = _genai.Client(api_key=key) n = len(images) # Build Frame-labelled image parts for the first user turn frame_parts: list[Any] = [] for k, img_bytes in enumerate(images, start=1): scaled = _downscale_image(img_bytes) if n > 1: frame_parts.append(_gtypes.Part.from_text(f"Frame {k}/{n}")) frame_parts.append(_gtypes.Part.from_bytes(data=scaled, mime_type="image/jpeg")) contents: list[Any] = [] if history: first_user_injected = False for turn in history: role = turn.get("role", "user") if role == "assistant": role = "model" if role == "user" and not first_user_injected: contents.append(_gtypes.Content( role="user", parts=frame_parts + [_gtypes.Part.from_text(str(turn.get("content", "")))], )) first_user_injected = True else: contents.append(_gtypes.Content( role=role, parts=[_gtypes.Part.from_text(str(turn.get("content", "")))], )) contents.append(_gtypes.Content(role="user", parts=[_gtypes.Part.from_text(prompt)])) else: contents.append(_gtypes.Content( role="user", parts=frame_parts + [_gtypes.Part.from_text(prompt)], )) try: response = client.models.generate_content( model=model, contents=contents, config=_gtypes.GenerateContentConfig(max_output_tokens=2048), ) return response.text except Exception as exc: raise AssistantError(str(exc)) from exc def _interpret_images_openai( images: list[bytes], prompt: str, history: list[dict] | None, *, model: str, key: str, ) -> str: """Interpret an image sequence using OpenAI vision (Chat Completions image_url parts).""" import base64 try: from openai import OpenAI as _OpenAI except ImportError as exc: raise AssistantUnavailable("openai package is not installed.") from exc client = _OpenAI(api_key=key) n = len(images) # Build Frame-labelled image parts for the first user turn frame_content: list[dict] = [] for k, img_bytes in enumerate(images, start=1): scaled = _downscale_image(img_bytes) b64 = base64.standard_b64encode(scaled).decode() if n > 1: frame_content.append({"type": "text", "text": f"Frame {k}/{n}"}) frame_content.append({ "type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}, }) messages: list[dict] = [] if history: first_user_injected = False for turn in history: if turn.get("role") == "user" and not first_user_injected: messages.append({ "role": "user", "content": frame_content + [{"type": "text", "text": str(turn.get("content", ""))}], }) first_user_injected = True else: messages.append({"role": turn["role"], "content": str(turn.get("content", ""))}) messages.append({"role": "user", "content": prompt}) else: messages.append({ "role": "user", "content": frame_content + [{"type": "text", "text": prompt}], }) try: response = client.chat.completions.create( model=model, max_tokens=2048, messages=messages, ) return response.choices[0].message.content or "" except Exception as exc: raise AssistantError(str(exc)) from exc def _effective_key_system_only(provider_id: str = DEFAULT_PROVIDER) -> str | None: """Return the org-level key for provider_id (for system: model selections).""" global_cfg = load_assistant_settings() org_keys = global_cfg.get("api_keys") or {} if org_keys.get(provider_id): return org_keys[provider_id] if provider_id == "anthropic" and global_cfg.get("api_key"): return global_cfg["api_key"] return None # --------------------------------------------------------------------------- # Tool dispatch infrastructure (Phase A, DVI-557) # --------------------------------------------------------------------------- # # Architecture: # - TIER0_TOOL_DEFS — registry of read-only Tier-0 tools; each entry carries # the Anthropic tool-use schema plus internal `required_roles` for authz. # - dispatch_tool() — authorize (deny-by-default) → execute → audit. # - chat_agentic() — Anthropic tool-use loop; admin-only in Phase A. # - Audit log — append-only JSONL outside model reach. # # Security invariants: # - Keys never enter model context. # - No tool may read assistant_settings.json / user settings / env. # - Tool results are scrubbed: no filesystem paths, no secrets. # - Every call is audited regardless of authz outcome. # Paths mirroring the same env vars used in app.py — read once at import time. _PRODRUN_DIR = Path(os.environ.get("PRODRUN_DIR", "/var/www/html/prodrunmerge")) _SHARED_DIR = Path(os.environ.get("QR_SHARED_DIR", "/var/www/html/shared")) _WOR_DIR = Path(os.environ.get("WOR_DIR", "/var/www/html/wor")) _APP_DIR = Path(__file__).resolve().parent TOOL_AUDIT_LOG = Path( os.environ.get("TOOL_AUDIT_LOG", str(_APP_DIR / "tool_audit.jsonl")) ) _SAFE_TOOL_ARG_RE = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9 _\-\.]{0,79}$") _audit_lock = threading.Lock() def _audit_append(entry: dict) -> None: """Append one JSON line to the audit log; swallows exceptions.""" try: with _audit_lock: with TOOL_AUDIT_LOG.open("a", encoding="utf-8") as fh: fh.write(json.dumps(entry) + "\n") except Exception: pass # Tier-0 tool definitions. `required_roles`: empty = any authenticated user; # non-empty = user must hold at least one of the listed session role flags. _TIER0_TOOL_DEFS: dict[str, dict] = { "prodrun_status": { "required_roles": ["prod_run", "togen_admin"], "anthropic_def": { "name": "prodrun_status", "description": ( "Get the current status of a production run merge job. " "Returns status, message, and result summary." ), "input_schema": { "type": "object", "properties": { "job_id": { "type": "string", "description": ( "Production run job ID — alphanumeric, spaces, " "dashes, or dots; max 80 characters." ), } }, "required": ["job_id"], }, }, }, "reprint_list_projects": { "required_roles": [], "anthropic_def": { "name": "reprint_list_projects", "description": "List available project folders for reprint or lookup, newest first.", "input_schema": { "type": "object", "properties": {}, "required": [], }, }, }, "wor_history": { "required_roles": [], "anthropic_def": { "name": "wor_history", "description": ( "List new and in-progress Work Order Request PDFs. " "Defaults to 'maintenance' type." ), "input_schema": { "type": "object", "properties": { "type": { "type": "string", "description": "WOR type.", "enum": ["maintenance", "mobilemachine", "safety", "preventive"], } }, "required": [], }, }, }, "manuals_lookup": { "required_roles": [], "anthropic_def": { "name": "manuals_lookup", "description": ( "Look up a manual previously associated with an Equipment, " "Asset, or System name." ), "input_schema": { "type": "object", "properties": { "equipment": { "type": "string", "description": "Equipment, asset, or system name.", } }, "required": ["equipment"], }, }, }, } # Anthropic-format tool list for use in client.messages.create(tools=...) TIER0_ANTHROPIC_TOOLS: list[dict] = [ d["anthropic_def"] for d in _TIER0_TOOL_DEFS.values() ] # --------------------------------------------------------------------------- # Tier-1 (mutating) tool registry (Phase B, DVI-560) # --------------------------------------------------------------------------- # These tools write data. They are only surfaced to togen_admin users (Phase B # is admin-only) and require explicit per-action human confirmation before # execution. Tier-2 routes (togen-admin writes, CSV-schema paths) are not # registered here and cannot be reached through the agentic path. _TIER1_TOOL_DEFS: dict[str, dict] = { "qr_submit": { "required_roles": ["takeoff_workflow", "togen_admin"], "tier": 1, "display_name": "Submit QR Merge Job", "make_display_args": lambda a: {"Folder": a.get("folder", "")}, "anthropic_def": { "name": "qr_submit", "description": ( "Submit a QR code merge job for a project folder. " "Merges shop drawings with carrier sheets and writes results to the shared folder. " "Requires explicit user confirmation before execution." ), "input_schema": { "type": "object", "properties": { "folder": { "type": "string", "description": "Project folder name (alphanumeric, spaces, dashes, dots; max 80 chars).", } }, "required": ["folder"], }, }, }, "duplicate_folder": { "required_roles": ["takeoff_workflow", "togen_admin"], "tier": 1, "display_name": "Duplicate Project Folder", "make_display_args": lambda a: { "Source": a.get("folder", ""), "New name": a.get("new_folder", ""), }, "anthropic_def": { "name": "duplicate_folder", "description": ( "Duplicate a project folder under a new name, copying all documents, CSV, and metadata. " "Requires explicit user confirmation before execution." ), "input_schema": { "type": "object", "properties": { "folder": { "type": "string", "description": "Source project folder name.", }, "new_folder": { "type": "string", "description": "New folder name (must not already exist).", }, }, "required": ["folder", "new_folder"], }, }, }, "prodrun_submit": { "required_roles": ["prod_run", "togen_admin"], "tier": 1, "display_name": "Submit Production Run", "make_display_args": lambda a: {"Folder": a.get("folder", "")}, "anthropic_def": { "name": "prodrun_submit", "description": ( "Submit a Production Run merge job. Starts the merge worker for the given folder. " "Requires explicit user confirmation before execution." ), "input_schema": { "type": "object", "properties": { "folder": { "type": "string", "description": "Production run job folder name.", } }, "required": ["folder"], }, }, }, "scs_save_to_project": { "required_roles": ["takeoff_workflow", "togen_admin"], "tier": 1, "display_name": "Save Carrier Sheet to Project", "make_display_args": lambda a: { "Job": a.get("job_id", ""), "Project folder": a.get("folder_name", ""), "Product code": a.get("product_code", ""), }, "anthropic_def": { "name": "scs_save_to_project", "description": ( "Save a Single Carrier Sheet (from an existing creation session) to a project folder, " "updating the project CSV with the carrier sheet URL. " "Requires explicit user confirmation before execution." ), "input_schema": { "type": "object", "properties": { "job_id": { "type": "string", "description": "SCS job ID from a completed carrier creation session.", }, "folder_name": { "type": "string", "description": "Target project folder name in the shared directory.", }, "product_code": { "type": "string", "description": "Product code to match or append in the project CSV.", }, "append_row": { "type": "boolean", "description": "Append a new CSV row if the product code is not found (default: false).", }, }, "required": ["job_id", "folder_name", "product_code"], }, }, }, "create_carrier": { "required_roles": [], # login_required — any authenticated user; admin-gate is at the agentic layer "tier": 1, "display_name": "Create Carrier Sheet", "make_display_args": lambda a: { "Job number": a.get("job_number", a.get("job_id", "")), "Product code": a.get("product_code", ""), "Customer": a.get("customer_name", ""), }, "anthropic_def": { "name": "create_carrier", "description": ( "Generate a QR carrier sheet for an existing upload session (job_id). " "The shop drawing PDF must already have been uploaded via the Create Carrier UI first. " "Optionally overrides carrier field values. " "Requires explicit user confirmation before execution." ), "input_schema": { "type": "object", "properties": { "job_id": { "type": "string", "description": "Job ID from a completed file-upload session (returned by the UI after uploading).", }, "job_number": {"type": "string", "description": "Job number override."}, "product_code": {"type": "string", "description": "Product code override."}, "customer_name": {"type": "string", "description": "Customer name override."}, "project_name": {"type": "string", "description": "Project name override."}, }, "required": ["job_id"], }, }, }, } # Merged registry: Tier-0 (read-only) + Tier-1 (mutating). Used for authz. _ALL_TOOL_DEFS: dict[str, dict] = {**_TIER0_TOOL_DEFS, **_TIER1_TOOL_DEFS} # Callbacks registered by app.py at startup — one per Tier-1 tool. # Signature: fn(args: dict, user: dict) -> dict _tier1_callbacks: dict[str, Any] = {} def register_tier1_callback(tool_id: str, fn: Any) -> None: """Register the executor for a Tier-1 tool (called by app.py at startup).""" _tier1_callbacks[tool_id] = fn def _authorize_tool(tool_id: str, user: dict) -> bool: """Return True when user holds at least one required role for this tool. Deny-by-default: unknown tool_ids return False. Empty required_roles list means any authenticated user is permitted (mirrors @login_required). Mirrors the same role flags as the underlying Flask route decorators. """ defn = _ALL_TOOL_DEFS.get(tool_id) if defn is None: return False required = defn["required_roles"] if not required: return True return any(user.get(r) for r in required) def has_any_tool_access(user: dict) -> bool: """Return True if the user is authorized to call at least one registered tool. Used by the chat route (Phase C) to decide whether a non-admin user gets agentic mode — if they have access to at least one tool, they use chat_agentic(); otherwise they fall back to plain chat(). """ return any(_authorize_tool(tid, user) for tid in _ALL_TOOL_DEFS) # --------------------------------------------------------------------------- # Tool handler functions — read-only, no mutations # --------------------------------------------------------------------------- def _tool_prodrun_status(args: dict, _user: dict) -> dict: """Return production run job status (safe summary only — no file paths).""" job_id = str(args.get("job_id", "")).strip() if not job_id or not _SAFE_TOOL_ARG_RE.match(job_id) or ".." in job_id: return {"error": "Invalid job_id."} status_path = _PRODRUN_DIR / job_id / "status.json" if not status_path.exists(): return {"status": "pending", "message": "Initializing…"} try: data = json.loads(status_path.read_text()) # Scrub: return only safe summary fields — no file system paths safe: dict[str, Any] = {} for k in ("status", "message", "drawing_merge", "label_merge"): if k in data: safe[k] = data[k] return safe except Exception: return {"error": "Could not read status."} def _tool_reprint_list_projects(_args: dict, _user: dict) -> dict: """Return project folder names from the shared directory.""" folders: list[str] = [] if _SHARED_DIR.is_dir(): try: for entry in sorted( _SHARED_DIR.iterdir(), key=lambda e: e.stat().st_mtime, reverse=True, ): if entry.is_dir() and not entry.name.startswith("."): folders.append(entry.name) except OSError: pass return {"folders": folders} def _tool_wor_history(args: dict, _user: dict) -> dict: """Return new and in-progress WOR PDF names for the requested type.""" wor_type = str(args.get("type", "maintenance")).strip() if wor_type not in {"maintenance", "mobilemachine", "safety", "preventive"}: wor_type = "maintenance" wor_base = _WOR_DIR / wor_type dirs: list[tuple[Path, str]] = [ (wor_base, "new"), (wor_base / "in_progress", "in_progress"), ] files: list[dict] = [] for directory, status in dirs: if not directory.is_dir(): continue for f in directory.glob("*.pdf"): try: mtime = datetime.utcfromtimestamp(f.stat().st_mtime) files.append({ "name": f.name, "date": mtime.strftime("%Y-%m-%dT%H:%M:%SZ"), "status": status, }) except OSError: continue files.sort(key=lambda x: x["date"], reverse=True) return {"files": files} def _tool_manuals_lookup(args: dict, _user: dict) -> dict: """Look up a manual by equipment name; returns name + association date only.""" equipment = str(args.get("equipment", "")).strip() if not equipment: return {"found": False} key = " ".join(equipment.lower().split()) index_file = _APP_DIR / "pm_manuals_index.json" if not index_file.is_file(): return {"found": False} try: data = json.loads(index_file.read_text()) except (json.JSONDecodeError, OSError): return {"found": False} entry = data.get(key) if not entry or not entry.get("filename"): return {"found": False} # Return safe fields only — no server-side paths return { "found": True, "equipment": entry.get("equipment", ""), "original_name": entry.get("original_name", entry["filename"]), "updated_at": entry.get("updated_at", ""), } _TOOL_HANDLERS: dict[str, Any] = { "prodrun_status": _tool_prodrun_status, "reprint_list_projects": _tool_reprint_list_projects, "wor_history": _tool_wor_history, "manuals_lookup": _tool_manuals_lookup, } def dispatch_tool(tool_id: str, args: dict, user: dict) -> dict: """Authorize, execute, and audit one tool call. Returns a result dict suitable for a tool_result message block. Never raises — errors are returned as {"error": "..."} dicts and audited. """ ts = datetime.now(timezone.utc).isoformat() email = (user or {}).get("email", "unknown") authz = _authorize_tool(tool_id, user) audit: dict[str, Any] = { "ts": ts, "user": email, "tool": tool_id, "args": args, "authz": authz, } if not authz: audit["outcome"] = "denied" _audit_append(audit) return {"error": f"Access denied: insufficient permissions for tool '{tool_id}'."} handler = _TOOL_HANDLERS.get(tool_id) if handler is None: audit["outcome"] = "unknown_tool" _audit_append(audit) return {"error": f"Unknown tool: '{tool_id}'."} try: result = handler(args, user) # Distinguish validated-input rejections from genuine successes so # audit reviewers can tell "blocked path traversal" from "returned data". audit["outcome"] = "input_error" if "error" in result else "success" _audit_append(audit) return result except Exception as exc: audit["outcome"] = f"handler_error: {exc}" _audit_append(audit) return {"error": "Tool execution failed."} # --------------------------------------------------------------------------- # Pending-action store (Phase B, DVI-560) # --------------------------------------------------------------------------- # When the agentic loop hits a Tier-1 tool call it stores conversation state # here, raises PendingActionRequired, and waits for user confirmation. Each # token is bound to the requesting user's email and expires after 10 minutes. _pending_actions: dict[str, dict] = {} _pending_lock = threading.Lock() _PENDING_ACTION_TTL = 600 # 10 minutes def _pending_purge_expired() -> None: """Remove expired tokens (caller must hold _pending_lock).""" now = datetime.now(timezone.utc).timestamp() expired = [k for k, v in _pending_actions.items() if now - v["ts"] > _PENDING_ACTION_TTL] for k in expired: del _pending_actions[k] def _pending_create( tool: str, args: dict, user_email: str, working_messages: list, tool_use_id: str, tier0_results: list, system: str, model: str, authorized_tools: Any, provider_id: str, extra_tier1: list, ) -> str: """Store conversation state keyed by a fresh UUID token; return the token. provider_id + working_messages are provider-native so resume_agentic can rebuild the correct client/adapter. extra_tier1 holds any additional Tier-1 calls emitted in the same turn beyond the first; they are auto-cancelled on resume so every tool call in the paused assistant turn gets a result (required by OpenAI/Gemini, and correct for Anthropic too). """ token = str(_uuid.uuid4()) with _pending_lock: _pending_purge_expired() _pending_actions[token] = { "ts": datetime.now(timezone.utc).timestamp(), "tool": tool, "args": args, "user_email": user_email, "working_messages": working_messages, "tool_use_id": tool_use_id, "tier0_results": tier0_results, "system": system, "model": model, "authorized_tools": authorized_tools, "provider_id": provider_id, "extra_tier1": extra_tier1, } return token def _pending_consume(token: str, user_email: str) -> dict | None: """Pop and return a pending action; None if missing, expired, or wrong user. Leaves the entry in place when the email doesn't match so the correct user can still consume it. """ with _pending_lock: _pending_purge_expired() entry = _pending_actions.get(token) if entry is None: return None if entry.get("user_email") != user_email: return None # wrong user — leave token intact del _pending_actions[token] return entry # --------------------------------------------------------------------------- # Provider-agnostic agentic tool-use loop (DVI-954; Phase A + B) # --------------------------------------------------------------------------- # # The loop logic — Tier-0/Tier-1 separation, deny-by-default authz, per-call # audit, and the Tier-1 pending-action confirmation flow — is provider-neutral. # All provider-format specifics live in adapter objects so Claude, Gemini, and # OpenAI reach identical tool-use parity from a single loop. # # Normalized types: # - tool call: {"id": str, "name": str, "input": dict} # - tool result: {"tool_use_id": str, "tool_name": str, "content": } # - _AgenticResult(text, tool_calls, stop_reason, raw) <- adapter.create() # stop_reason is normalized to "tool_use" | "end_turn" | provider-native # (e.g. "max_tokens"/"length"); raw carries the provider-native assistant # content needed to rebuild the assistant turn. # # Working messages are always provider-native so the pending-action store can # persist them (in-process memory) alongside the provider id and resume_agentic # rebuilds the matching client/adapter. _AGENTIC_MAX_ROUNDS = 10 # safety cap: maximum tool-call rounds per request def _content_block_to_dict(block: Any) -> dict: """Convert an Anthropic SDK content block to a plain dict.""" if block.type == "text": return {"type": "text", "text": block.text} if block.type == "tool_use": return { "type": "tool_use", "id": block.id, "name": block.name, "input": dict(block.input), } # Fallback: try model_dump() for future block types, else drop try: return block.model_dump() except Exception: return {"type": block.type} class _AgenticResult: """Normalized result of one adapter.create() call.""" __slots__ = ("text", "tool_calls", "stop_reason", "raw") def __init__(self, text: str, tool_calls: list, stop_reason: str, raw: Any) -> None: self.text = text self.tool_calls = tool_calls self.stop_reason = stop_reason self.raw = raw def _tool_spec(defn: dict) -> dict: """Return the canonical JSON-Schema tool spec for a tool definition. The `anthropic_def` field ({name, description, input_schema}) is the canonical spec reused across all providers; each adapter's build_tools() translates it into the provider-native format. """ return defn["anthropic_def"] def _json_schema_to_gemini(schema: Any) -> Any: """Best-effort convert a JSON-Schema object into a Gemini types.Schema-friendly dict (uppercased `type` values, recursively) so FunctionDeclaration accepts it. Falls back to returning the input unchanged for non-dict nodes. """ if not isinstance(schema, dict): return schema out: dict[str, Any] = {} for k, v in schema.items(): if k == "type" and isinstance(v, str): out["type"] = v.upper() elif k == "properties" and isinstance(v, dict): out["properties"] = {pk: _json_schema_to_gemini(pv) for pk, pv in v.items()} elif k == "items": out["items"] = _json_schema_to_gemini(v) else: out[k] = v return out class _AnthropicAgentAdapter: provider_id = "anthropic" def make_client(self, key: str) -> Any: try: import anthropic as _anthropic except ImportError as exc: raise AssistantUnavailable("anthropic package is not installed.") from exc return _anthropic.Anthropic(api_key=key) def init_messages(self, messages: list[dict]) -> list: return [dict(m) for m in messages] def build_tools(self, specs: list[dict]) -> Any: return [dict(s) for s in specs] def create(self, client, messages, system, model, tools) -> _AgenticResult: try: import anthropic as _anthropic except ImportError as exc: raise AssistantUnavailable("anthropic package is not installed.") from exc try: resp = client.messages.create( model=model, max_tokens=4096, system=system or "", messages=messages, tools=tools, ) except _anthropic.APIError as exc: raise AssistantError(str(exc)) from exc text = "" calls: list = [] for block in resp.content: if block.type == "text": text += block.text elif block.type == "tool_use": calls.append({"id": block.id, "name": block.name, "input": dict(block.input)}) raw = [_content_block_to_dict(b) for b in resp.content] stop = "tool_use" if resp.stop_reason == "tool_use" else ( "end_turn" if resp.stop_reason == "end_turn" else resp.stop_reason ) return _AgenticResult(text, calls, stop, raw) def format_assistant_turn(self, resp: _AgenticResult) -> Any: return {"role": "assistant", "content": resp.raw} def format_tool_results(self, results: list[dict]) -> list: blocks = [ {"type": "tool_result", "tool_use_id": r["tool_use_id"], "content": r["content"]} for r in results ] return [{"role": "user", "content": blocks}] class _OpenAIAgentAdapter: provider_id = "openai" def make_client(self, key: str) -> Any: try: from openai import OpenAI as _OpenAI except ImportError as exc: raise AssistantUnavailable("openai package is not installed.") from exc return _OpenAI(api_key=key) def init_messages(self, messages: list[dict]) -> list: return [{"role": m.get("role", "user"), "content": str(m.get("content", ""))} for m in messages] def build_tools(self, specs: list[dict]) -> Any: return [ { "type": "function", "function": { "name": s["name"], "description": s.get("description", ""), "parameters": s["input_schema"], }, } for s in specs ] def create(self, client, messages, system, model, tools) -> _AgenticResult: oa_messages: list = [] if system: oa_messages.append({"role": "system", "content": system}) oa_messages.extend(messages) try: resp = client.chat.completions.create( model=model, max_tokens=4096, messages=oa_messages, tools=tools or None, ) except Exception as exc: raise AssistantError(str(exc)) from exc choice = resp.choices[0] msg = choice.message text = msg.content or "" calls: list = [] raw_tool_calls: list = [] for tc in (getattr(msg, "tool_calls", None) or []): raw_args = tc.function.arguments or "{}" try: args = json.loads(raw_args) except json.JSONDecodeError: args = {} calls.append({"id": tc.id, "name": tc.function.name, "input": args}) raw_tool_calls.append({ "id": tc.id, "type": "function", "function": {"name": tc.function.name, "arguments": raw_args}, }) finish = choice.finish_reason stop = "tool_use" if (finish == "tool_calls" or calls) else ( "end_turn" if finish == "stop" else finish ) raw = {"content": text or None, "tool_calls": raw_tool_calls} return _AgenticResult(text, calls, stop, raw) def format_assistant_turn(self, resp: _AgenticResult) -> Any: msg: dict = {"role": "assistant", "content": resp.raw["content"]} if resp.raw["tool_calls"]: msg["tool_calls"] = resp.raw["tool_calls"] return msg def format_tool_results(self, results: list[dict]) -> list: return [ {"role": "tool", "tool_call_id": r["tool_use_id"], "content": r["content"]} for r in results ] class _GoogleAgentAdapter: provider_id = "google" def make_client(self, key: str) -> Any: try: from google import genai as _genai except ImportError as exc: raise AssistantUnavailable("google-genai package is not installed.") from exc return _genai.Client(api_key=key) def init_messages(self, messages: list[dict]) -> list: from google.genai import types as _gtypes contents = [] for m in messages: role = m.get("role", "user") if role == "assistant": role = "model" contents.append(_gtypes.Content(role=role, parts=[_gtypes.Part.from_text(str(m.get("content", "")))])) return contents def build_tools(self, specs: list[dict]) -> Any: from google.genai import types as _gtypes decls = [ _gtypes.FunctionDeclaration( name=s["name"], description=s.get("description", ""), parameters=_json_schema_to_gemini(s["input_schema"]), ) for s in specs ] return [_gtypes.Tool(function_declarations=decls)] def create(self, client, messages, system, model, tools) -> _AgenticResult: from google.genai import types as _gtypes try: resp = client.models.generate_content( model=model, contents=messages, config=_gtypes.GenerateContentConfig( system_instruction=system or None, tools=tools, max_output_tokens=4096, ), ) except Exception as exc: raise AssistantError(str(exc)) from exc cand = resp.candidates[0] text = "" calls: list = [] for i, part in enumerate(cand.content.parts): if getattr(part, "text", None): text += part.text fc = getattr(part, "function_call", None) if fc is not None: # Gemini matches tool results by function name, not id; synthesize # a positional id for the normalized layer. calls.append({"id": f"{fc.name}::{i}", "name": fc.name, "input": dict(fc.args or {})}) stop = "tool_use" if calls else "end_turn" return _AgenticResult(text, calls, stop, cand.content) def format_assistant_turn(self, resp: _AgenticResult) -> Any: return resp.raw # native model Content def format_tool_results(self, results: list[dict]) -> list: from google.genai import types as _gtypes parts = [] for r in results: try: payload = json.loads(r["content"]) except (json.JSONDecodeError, TypeError): payload = {"result": r["content"]} if not isinstance(payload, dict): payload = {"result": payload} parts.append(_gtypes.Part.from_function_response(name=r["tool_name"], response=payload)) return [_gtypes.Content(role="user", parts=parts)] _AGENT_ADAPTERS: dict[str, Any] = { "anthropic": _AnthropicAgentAdapter(), "openai": _OpenAIAgentAdapter(), "google": _GoogleAgentAdapter(), } def _get_agentic_adapter(provider_id: str) -> Any: """Return the agentic adapter for provider_id, or raise AssistantUnavailable.""" adapter = _AGENT_ADAPTERS.get(provider_id) if adapter is None: raise AssistantUnavailable( f"Agentic tool-use is not available for provider '{provider_id}'." ) return adapter def _agentic_loop( adapter: Any, client: Any, working_messages: list, system: str, model: str, tools: Any, user: dict, ) -> str: """Run the provider-agnostic agentic tool-calling loop. Executes Tier-0 tools inline. Raises PendingActionRequired when the model calls a Tier-1 (mutating) tool so the caller can request user confirmation. Raises AssistantError on provider API failure. """ user_email = (user or {}).get("email", "") for _round in range(_AGENTIC_MAX_ROUNDS): resp = adapter.create(client, working_messages, system, model, tools) if resp.stop_reason != "tool_use" or not resp.tool_calls: return resp.text # Record the assistant turn (provider-native) before answering tools. working_messages.append(adapter.format_assistant_turn(resp)) tier0_calls: list = [] tier1_calls: list = [] for call in resp.tool_calls: if call["name"] in _TIER1_TOOL_DEFS: tier1_calls.append(call) else: tier0_calls.append(call) # Execute Tier-0 (read-only) tools immediately. tier0_results = [] for call in tier0_calls: result = dispatch_tool(call["name"], call["input"], user) tier0_results.append({ "tool_use_id": call["id"], "tool_name": call["name"], "content": json.dumps(result), }) if tier1_calls: # Pause on the first Tier-1 tool; store state for user confirmation. # Any additional Tier-1 calls in the same turn are recorded so resume # can auto-cancel them (keeps every tool call answered for all providers). t1 = tier1_calls[0] defn = _TIER1_TOOL_DEFS[t1["name"]] display_args = defn["make_display_args"](t1["input"]) extra_tier1 = [ {"tool_use_id": c["id"], "tool_name": c["name"]} for c in tier1_calls[1:] ] token = _pending_create( tool=t1["name"], args=t1["input"], user_email=user_email, working_messages=working_messages, tool_use_id=t1["id"], tier0_results=tier0_results, system=system, model=model, authorized_tools=tools, provider_id=adapter.provider_id, extra_tier1=extra_tier1, ) raise PendingActionRequired( token=token, tool=t1["name"], display_name=defn["display_name"], display_args=display_args, partial_reply=resp.text, ) working_messages.extend(adapter.format_tool_results(tier0_results)) return "\n\n[Response truncated: tool-call limit reached.]" def _build_agentic_context(user: dict, system: str) -> tuple[list, str]: """Build the authorized JSON-Schema tool specs + confirmation preamble. dispatch_tool() still enforces authz per-call; restricting the surfaced tool list just prevents the model from requesting tools the user can't use. Returns (tool_specs, system) where system has the Tier-1 preamble prepended when any mutating tool is available. """ tool_specs = [ _tool_spec(defn) for tool_id, defn in _ALL_TOOL_DEFS.items() if _authorize_tool(tool_id, user) ] tier1_names = [tid for tid in _TIER1_TOOL_DEFS if _authorize_tool(tid, user)] if tier1_names: names_str = ", ".join(tier1_names) preamble = ( "IMPORTANT — mutating tools: when you call any of the following tools " f"({names_str}), the user will see a confirmation prompt before the " "action executes. If they decline you will receive " '{"cancelled": true, "message": "Action cancelled by user."} as the ' "tool result; acknowledge this gracefully. " "Propose at most one mutating action per response turn." ) system = (preamble + "\n\n" + system).strip() if system else preamble return tool_specs, system def chat_agentic( messages: list[dict], system: str, *, user: dict, user_settings: dict, model: str | None = None, ) -> str: """Agentic chat: provider-agnostic tool-use loop with per-tool authz and audit. Resolves the provider from the user's selected default model and runs the tool-use loop on Claude, Gemini, or OpenAI. Tier-0 (read-only) tools execute inline. Tier-1 (mutating) tools pause the loop and raise PendingActionRequired; call resume_agentic() after the user decides. user: session["user"] dict — used for per-tool authorization. user_settings: caller's user-settings dict — used for key resolution. Raises AssistantUnavailable / AssistantError on API failure. Raises PendingActionRequired when a mutating tool awaits user confirmation. """ model, provider_id, key = _resolve_chat_model(user_settings, model) if not key: raise AssistantUnavailable("No API key configured.") adapter = _get_agentic_adapter(provider_id) client = adapter.make_client(key) tool_specs, system = _build_agentic_context(user, system) tools = adapter.build_tools(tool_specs) working_messages = adapter.init_messages(list(messages)) return _agentic_loop(adapter, client, working_messages, system, model, tools, user) def resume_agentic( token: str, confirmed: bool, user: dict, user_settings: dict, ) -> str: """Resume the agentic loop after the user confirms or cancels a Tier-1 action. Retrieves the pending conversation state (including the provider id + provider-native working messages), executes (or skips) the Tier-1 tool, and continues the loop until the model emits a final reply. Raises AssistantUnavailable / AssistantError on API failure. Raises PendingActionRequired if the resumed loop hits another Tier-1 tool. """ user_email = (user or {}).get("email", "") entry = _pending_consume(token, user_email) if entry is None: return "This action has expired or was already handled. Please start a new request." tool = entry["tool"] args = entry["args"] working_msgs = entry["working_messages"] tool_use_id = entry["tool_use_id"] tool_name = tool tier0_results = entry["tier0_results"] system = entry["system"] model = entry["model"] tools = entry["authorized_tools"] provider_id = entry.get("provider_id", DEFAULT_PROVIDER) extra_tier1 = entry.get("extra_tier1") or [] ts = datetime.now(timezone.utc).isoformat() audit: dict[str, Any] = { "ts": ts, "user": user_email, "tool": tool, "args": args, "authz": True, } if not confirmed: audit["outcome"] = "cancelled_by_user" _audit_append(audit) tier1_result: dict = {"cancelled": True, "message": "Action cancelled by user."} else: cb = _tier1_callbacks.get(tool) if cb is None: audit["outcome"] = "no_callback" _audit_append(audit) tier1_result = {"error": f"No handler registered for '{tool}'. Contact your administrator."} else: try: tier1_result = cb(args, user) audit["outcome"] = "input_error" if "error" in tier1_result else "success" _audit_append(audit) except Exception as exc: audit["outcome"] = f"handler_error: {exc}" _audit_append(audit) tier1_result = {"error": "Tool execution failed."} # Answer every tool call from the paused assistant turn: # tier-0 results (dispatched before the pause) + the confirmed/cancelled # tier-1 result + auto-cancelled results for any extra tier-1 calls. combined_results = list(tier0_results) + [{ "tool_use_id": tool_use_id, "tool_name": tool_name, "content": json.dumps(tier1_result), }] for extra in extra_tier1: combined_results.append({ "tool_use_id": extra["tool_use_id"], "tool_name": extra["tool_name"], "content": json.dumps({ "cancelled": True, "message": "Not executed: only one mutating action is processed per turn.", }), }) adapter = _get_agentic_adapter(provider_id) key = _effective_key(user_settings, provider_id) if not key: raise AssistantUnavailable("No API key configured.") client = adapter.make_client(key) resumed_messages = list(working_msgs) resumed_messages.extend(adapter.format_tool_results(combined_results)) return _agentic_loop(adapter, client, resumed_messages, system, model, tools, user) # --------------------------------------------------------------------------- # Key masking helper # --------------------------------------------------------------------------- def mask_key(key: str | None) -> dict: """Return a safe public representation of a key (never the raw value).""" if not key: return {"configured": False, "masked": ""} last4 = key[-4:] if len(key) >= 4 else "****" return {"configured": True, "masked": f"••••{last4}"} def mask_keys(keys_dict: dict | None) -> dict: """Return masked representation of a per-provider keys dict for all registered providers.""" src = keys_dict or {} return {pid: mask_key(src.get(pid)) for pid in ASSISTANT_PROVIDERS} # --------------------------------------------------------------------------- # Per-provider enabled flags (DVI-952) # --------------------------------------------------------------------------- # Two independent toggles gate a provider (enforcement lands in later phases): # - org-level ("available system-wide") stored in assistant_settings.json # - user-level ("active") stored in user settings assistant{} # Both default to True when absent so existing Claude/Gemini setups are # unaffected until an admin/user explicitly turns a provider off. def provider_enabled_map(enabled_dict: dict | None) -> dict: """Return {provider_id: bool} for every registered provider, defaulting True.""" src = enabled_dict or {} return {pid: bool(src.get(pid, True)) for pid in ASSISTANT_PROVIDERS} def org_provider_enabled_map(global_cfg: dict | None = None) -> dict: """Return the org-level enabled map (admin 'available system-wide' toggles).""" cfg = global_cfg if global_cfg is not None else load_assistant_settings() return provider_enabled_map(cfg.get("provider_enabled")) def user_provider_enabled_map(user_settings: dict) -> dict: """Return the user-level enabled map (user 'active' toggles).""" ua = get_user_assistant(user_settings) return provider_enabled_map(ua.get("provider_enabled")) def available_models(user_settings: dict) -> dict: """List the models a user can pick per-invocation for chat. A model is offered only when its provider is org-enabled AND user-enabled AND has an effective key (user or org). Returns {models: [...], default_model}. Each model carries provider id/label so the picker can group by provider. """ org_en = org_provider_enabled_map() user_en = user_provider_enabled_map(user_settings) ua = get_user_assistant(user_settings) default_model = ua.get("default_model") or ua.get("model") or DEFAULT_MODEL models: list[dict] = [] for pid, prov in ASSISTANT_PROVIDERS.items(): if not (org_en.get(pid, True) and user_en.get(pid, True)): continue if not _effective_key(user_settings, pid): continue for m in prov.get("models", []): models.append({ "provider_id": pid, "provider_label": prov["label"], "id": m["id"], "label": m["label"], "default": m["id"] == default_model, }) return {"models": models, "default_model": default_model} def org_ocr_models(global_cfg: dict | None = None) -> dict: """List models an admin can pick for title-block OCR. A model is offered only when its provider has an org-level key configured (OCR always runs on the org key). Returns {models: [...], selected} where ``selected`` is the saved ``ocr_model`` ("" means the Anthropic default). """ cfg = global_cfg if global_cfg is not None else load_assistant_settings() selected = str(cfg.get("ocr_model") or "").strip() models: list[dict] = [] for pid, prov in ASSISTANT_PROVIDERS.items(): if not _effective_key_system_only(pid): continue for m in prov.get("models", []): models.append({ "provider_id": pid, "provider_label": prov["label"], "id": m["id"], "label": m["label"], }) return {"models": models, "selected": selected, "default_model": DEFAULT_MODEL} # --------------------------------------------------------------------------- # Exceptions # --------------------------------------------------------------------------- class AssistantUnavailable(RuntimeError): """Raised when the assistant cannot be used (not configured / gated out).""" class AssistantError(RuntimeError): """Raised on API or provider errors during a chat/vision call.""" class PendingActionRequired(Exception): """Raised by the agentic loop when a Tier-1 tool awaits user confirmation. Attributes: token: UUID string identifying the stored conversation state. tool: Tier-1 tool name. display_name: Human-readable action label for the confirmation card. display_args: Key-value pairs to render in the confirmation card. partial_reply: Any model text emitted before the tool call (may be ""). """ __slots__ = ("token", "tool", "display_name", "display_args", "partial_reply") def __init__( self, token: str, tool: str, display_name: str, display_args: dict, partial_reply: str = "", ) -> None: self.token = token self.tool = tool self.display_name = display_name self.display_args = display_args self.partial_reply = partial_reply