services/protocol/web_search_tool.py

from __future__ import annotations

import json
import re
from typing import Any

from services.account_service import account_service
from services.openai_backend_api import OpenAIBackendAPI

WEB_SEARCH_TOOL_TYPES = {"web_search", "web_search_preview", "web_search_preview_2025_03_11"}
SEARCH_CHAT_MODEL_PREFIXES = (
    "gpt-4o-search-preview",
    "gpt-4o-mini-search-preview",
    "gpt-5-search-api",
)


def _tool_type(tool: object) -> str:
    return str(tool.get("type") or "").strip() if isinstance(tool, dict) else ""


def has_web_search_tool(body: dict[str, Any]) -> bool:
    tools = body.get("tools")
    if isinstance(tools, list):
        return any(_tool_type(tool) in WEB_SEARCH_TOOL_TYPES for tool in tools)
    tool_choice = body.get("tool_choice")
    return _tool_type(tool_choice) in WEB_SEARCH_TOOL_TYPES


def is_web_search_chat_request(body: dict[str, Any]) -> bool:
    model = str(body.get("model") or "").strip()
    return (
        has_web_search_tool(body)
        or isinstance(body.get("web_search_options"), dict)
        or any(
            model == prefix or model.startswith(f"{prefix}-")
            for prefix in SEARCH_CHAT_MODEL_PREFIXES
        )
    )


def has_unsupported_tools(body: dict[str, Any], allowed_types: set[str]) -> bool:
    tools = body.get("tools")
    if not isinstance(tools, list):
        return False
    return any(_tool_type(tool) not in allowed_types for tool in tools if isinstance(tool, dict))


def message_text(value: object) -> str:
    if isinstance(value, str):
        return value.strip()
    if isinstance(value, list):
        parts: list[str] = []
        for item in value:
            if isinstance(item, str):
                text = item.strip()
            elif isinstance(item, dict):
                text = str(item.get("text") or item.get("input_text") or "").strip()
            else:
                text = ""
            if text:
                parts.append(text)
        return "\n".join(parts).strip()
    return ""


def search_query_from_messages(messages: list[dict[str, Any]]) -> str:
    for message in reversed(messages):
        if str(message.get("role") or "").strip().lower() != "user":
            continue
        text = message_text(message.get("content"))
        if text:
            return text
    return ""


def _readable_annotation_part(parts: list[str]) -> str:
    for part in parts:
        value = part.strip()
        lower = value.lower()
        if value and not (
            lower.startswith(("turn", "source", "sources"))
            or re.fullmatch(r"\d+", value)
        ):
            return value
    return ""


def clean_search_text(text: str) -> str:
    def replace_annotation(match: re.Match[str]) -> str:
        parts = [part.strip() for part in match.group(1).split("\ue202")]
        kind = (parts[0] if parts else "").lower()
        data = parts[1:]
        if kind == "url":
            label = data[0] if data else ""
            url = data[1] if len(data) > 1 else ""
            if label and url.startswith(("http://", "https://")):
                return f"{label} ({url})"
            return label or url
        if kind == "cite":
            return _readable_annotation_part(data)
        if kind in {"image_group", "known_celebrity"}:
            return ""
        readable = _readable_annotation_part(data)
        if readable.startswith(("{", "[")):
            try:
                json.loads(readable)
                return ""
            except json.JSONDecodeError:
                pass
        return readable

    text = re.sub(r"\ue200([^\ue201]*)\ue201", replace_annotation, text)
    text = re.sub(r"\ue200[^\ue201]*$", "", text)
    # Search answers can occasionally contain the web client's presentation
    # protocol without its private \ue200/\ue201 annotation wrapper. These are
    # instructions for ChatGPT's own place card / image carousel renderer, not
    # user-facing answer text and not direct image URLs. The actual cards, when
    # available, are extracted independently from metadata by
    # ``_extract_search_images``.
    decoder = json.JSONDecoder()
    parts: list[str] = []
    cursor = 0
    while cursor < len(text):
        match = re.search(r"[\[{]", text[cursor:])
        if not match:
            parts.append(text[cursor:])
            break
        start = cursor + match.start()
        parts.append(text[cursor:start])
        try:
            value, end = decoder.raw_decode(text[start:])
        except json.JSONDecodeError:
            parts.append(text[start:start + 1])
            cursor = start + 1
            continue
        is_place_card = (
            isinstance(value, list)
            and len(value) >= 1
            and str(value[0]).strip().lower() == "point_of_interest"
        )
        is_image_carousel = (
            isinstance(value, dict)
            and str(value.get("layout") or "").strip().lower() == "carousel"
            and isinstance(value.get("query"), list)
        )
        if is_place_card or is_image_carousel:
            cursor = start + end
            continue
        parts.append(text[start:start + end])
        cursor = start + end
    text = "".join(parts)
    return re.sub(r"\s+([.,;:!?])", r"\1", text).strip()


def normalized_places(result: dict[str, Any]) -> list[dict[str, str]]:
    """Extract the public fields from ChatGPT's inline place-card protocol.

    The protocol is rendered by chatgpt.com's own client, so a compatible
    client must preserve the useful place data separately instead of exposing
    its raw JSON in the answer body.
    """
    text = str(result.get("answer") or "")
    decoder = json.JSONDecoder()
    places: list[dict[str, str]] = []
    seen: set[tuple[str, str]] = set()
    cursor = 0
    while cursor < len(text):
        match = re.search(r"\[", text[cursor:])
        if not match:
            break
        start = cursor + match.start()
        try:
            value, end = decoder.raw_decode(text[start:])
        except json.JSONDecodeError:
            cursor = start + 1
            continue
        cursor = start + end
        if not (isinstance(value, list) and len(value) >= 3 and str(value[0]).strip().lower() == "point_of_interest"):
            continue
        title = str(value[1] or "").strip()
        address = str(value[2] or "").strip()
        key = (title, address)
        if not title or key in seen:
            continue
        seen.add(key)
        places.append({
            "id": f"place-{len(places) + 1}",
            "title": title,
            "address": address,
        })
    return places


def normalized_sources(result: dict[str, Any]) -> list[dict[str, str]]:
    sources = result.get("sources")
    if not isinstance(sources, list):
        return []
    output: list[dict[str, str]] = []
    seen: set[str] = set()
    for item in sources:
        if not isinstance(item, dict):
            continue
        url = str(item.get("url") or "").strip()
        title = str(item.get("title") or "").strip()
        snippet = str(item.get("snippet") or "").strip()
        if not url or url in seen:
            continue
        seen.add(url)
        output.append({"title": title, "url": url, "snippet": snippet})
    return output


def normalized_images(result: dict[str, Any]) -> list[dict[str, Any]]:
    images = result.get("images")
    if not isinstance(images, list):
        return []
    output: list[dict[str, Any]] = []
    seen: set[str] = set()
    for item in images:
        if not isinstance(item, dict):
            continue
        content_url = str(item.get("content_url") or item.get("thumbnail_url") or "").strip()
        if not content_url.startswith(("http://", "https://")) or content_url in seen:
            continue
        seen.add(content_url)
        output.append({
            "id": str(item.get("id") or f"search-image-{len(output) + 1}"),
            "title": str(item.get("title") or "搜索图片").strip() or "搜索图片",
            "content_url": content_url,
            "thumbnail_url": str(item.get("thumbnail_url") or content_url).strip(),
            "source_url": str(item.get("source_url") or "").strip(),
            "query": str(item.get("query") or "").strip(),
            "width": int(item.get("width") or 0),
            "height": int(item.get("height") or 0),
        })
    return output


def text_with_url_citations(result: dict[str, Any]) -> tuple[str, list[dict[str, Any]]]:
    text = clean_search_text(str(result.get("answer") or ""))
    annotations: list[dict[str, Any]] = []
    sources = normalized_sources(result)
    if sources:
        text = text.rstrip()
        if text:
            text += "\n\n"
        text += "Sources:\n"
        for index, source in enumerate(sources, start=1):
            title = source["title"] or source["url"]
            line_prefix = f"{index}. {title}"
            text += line_prefix
            if source["url"]:
                if source["title"]:
                    text += " - "
                start = len(text)
                text += source["url"]
                annotations.append({
                    "type": "url_citation",
                    "start_index": start,
                    "end_index": len(text),
                    "url": source["url"],
                    "title": source["title"] or source["url"],
                })
            text += "\n"
    return text.strip(), annotations


def run_web_search(query: str) -> dict[str, Any]:
    token = account_service.get_text_access_token()
    backend = OpenAIBackendAPI(token)
    try:
        result = backend.search(query)
    finally:
        backend.close()
    account_service.mark_text_used(token)
    return result