借助 Firecrawl 搜索为 Vertex AI 上的 Gemini 提供 Grounding

"""基于 Firecrawl /v2/search 的 Gemini “Grounding with your search API” 适配器。"""
import hmac
import os
import re

import httpx
from flask import Flask, jsonify, request

API_URL = os.environ.get("FIRECRAWL_API_URL", "https://api.firecrawl.dev")
INTEGRATION_TAG = "_vertex_grounding"
LIMIT = int(os.environ.get("RESULT_LIMIT", "10"))
SNIPPET_CHARS = int(os.environ.get("SNIPPET_CHARS", "1200"))
SNIPPET_MODE = os.environ.get("SNIPPET_MODE", "description")
GATEWAY_KEY = os.environ.get("GATEWAY_KEY", "")
ENV_FIRECRAWL_KEY = os.environ.get("FIRECRAWL_API_KEY", "")
TIMEOUT = float(os.environ.get("FIRECRAWL_TIMEOUT", "20"))
MAX_QUERY_CHARS = 500

app = Flask(__name__)
client = httpx.Client(timeout=TIMEOUT)


def resolve_firecrawl_key(supplied):
    """返回要使用的 Firecrawl API 密钥；若调用方无权访问，则返回 None。"""
    if not supplied:
        return None
    if GATEWAY_KEY:
        if ENV_FIRECRAWL_KEY and hmac.compare_digest(supplied, GATEWAY_KEY):
            return ENV_FIRECRAWL_KEY
        return None
    return supplied


def excerpt(text, limit):
    text = re.sub(r"!\[[^\]]*\]\([^)]*\)", " ", text or "")
    text = re.sub(r"\[([^\]]*)\]\([^)]*\)", r"\1", text)
    text = re.sub(r"[#*_>|`]+", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    if len(text) <= limit:
        return text
    return text[:limit].rsplit(" ", 1)[0] + "..."


def to_grounding(item):
    uri = item.get("url")
    if not uri:
        return None
    source = item.get("markdown") if SNIPPET_MODE == "markdown" else None
    snippet = excerpt(source or item.get("description") or item.get("title"), SNIPPET_CHARS)
    return {"snippet": snippet, "uri": uri} if snippet else None


@app.get("/healthz")
def healthz():
    return {"ok": True}


@app.post("/")
@app.post("/search")
@app.post("/v0/search")
def search():
    key = resolve_firecrawl_key(request.args.get("key", ""))
    if not key:
        return jsonify({"error": "Missing or invalid key"}), 401
    body = request.get_json(silent=True)
    query = body.get("query") if isinstance(body, dict) else None
    if not isinstance(query, str) or not query.strip():
        return jsonify({"error": "JSON body with a non-empty 'query' string is required"}), 400

    payload = {"query": query.strip()[:MAX_QUERY_CHARS], "limit": LIMIT, "integration": INTEGRATION_TAG}
    if SNIPPET_MODE == "markdown":
        payload["scrapeOptions"] = {"formats": ["markdown"], "onlyMainContent": True}
    try:
        resp = client.post(
            f"{API_URL}/v2/search",
            json=payload,
            headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
        )
    except httpx.HTTPError:
        return jsonify({"error": "Firecrawl request failed"}), 502
    if resp.status_code in (401, 403):
        return jsonify({"error": "Missing or invalid key"}), 401
    if resp.status_code == 429:
        return jsonify({"error": "Firecrawl rate limit reached"}), 429
    if resp.status_code != 200:
        return jsonify({"error": f"Firecrawl returned {resp.status_code}"}), 502

    web = ((resp.json().get("data") or {}).get("web")) or []
    results = [r for r in map(to_grounding, web) if r]
    return jsonify(results)

flask>=3.0,<4
gunicorn>=22.0
httpx>=0.27,<1

FROM python:3.12-slim
ENV PYTHONUNBUFFERED=1
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY app.py .
RUN useradd --create-home app
USER app
CMD exec gunicorn --bind 0.0.0.0:${PORT:-8080} --workers 1 --threads 8 --timeout 60 app:app

from google import genai
from google.genai.types import ExternalApi, GenerateContentConfig, Retrieval, Tool

client = genai.Client(vertexai=True, project="your-project-id", location="global")

tool = Tool(
    retrieval=Retrieval(
        external_api=ExternalApi(
            api_spec="SIMPLE_SEARCH",
            endpoint="https://YOUR_SERVICE_URL/v0/search",
            api_auth={"apiKeyConfig": {"apiKeyString": "YOUR_KEY"}},
        )
    )
)

response = client.models.generate_content(
    model="gemini-2.5-flash",
    contents="Who won Super Bowl LX?",
    config=GenerateContentConfig(tools=[tool]),
)

print(response.text)
print(response.candidates[0].grounding_metadata.grounding_chunks)

PROJECT_ID=your-project-id
LOCATION=global
MODEL_ID=gemini-2.5-flash

cat > request.json <<'EOF'
{
  "contents": [{
    "role": "user",
    "parts": [{ "text": "Who won Super Bowl LX?" }]
  }],
  "tools": [{
    "retrieval": {
      "externalApi": {
        "api_spec": "SIMPLE_SEARCH",
        "endpoint": "https://YOUR_SERVICE_URL/v0/search",
        "apiAuth": {
          "apiKeyConfig": { "apiKeyString": "YOUR_KEY" }
        }
      }
    }
  }]
}
EOF

curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  -d `request.json \
  "https://aiplatform.googleapis.com/v1beta1/projects/$PROJECT_ID/locations/$LOCATION/publishers/google/models/$MODEL_ID:generateContent"