Building Event-Driven Async AI Pipelines with Gemini API — Pub/Sub, Webhooks, and Queue Integration for Production

# publisher.py
import json
import uuid
from datetime import datetime
from fastapi import FastAPI, BackgroundTasks
from google.cloud import pubsub_v1
from pydantic import BaseModel
 
app = FastAPI()
publisher = pubsub_v1.PublisherClient()
 
PROJECT_ID = "your-project-id"
TOPIC_ID = "gemini-tasks"
topic_path = publisher.topic_path(PROJECT_ID, TOPIC_ID)
 
class ProcessRequest(BaseModel):
    content: str
    task_type: str  # "summarize" / "classify" / "analyze"
    user_id: str
    callback_url: str | None = None
 
@app.post("/process")
async def submit_processing_task(request: ProcessRequest):
    """
    Submits an async processing task to Pub/Sub.
    Returns task_id immediately — processing happens in background.
    """
    task_id = str(uuid.uuid4())
 
    message_data = {
        "task_id": task_id,
        "content": request.content,
        "task_type": request.task_type,
        "user_id": request.user_id,
        "callback_url": request.callback_url,
        "created_at": datetime.utcnow().isoformat(),
    }
 
    # Publish to Pub/Sub
    future = publisher.publish(
        topic_path,
        data=json.dumps(message_data).encode("utf-8"),
        # Attributes enable subscription-level filtering
        task_type=request.task_type,
        user_id=request.user_id,
    )
 
    future.add_done_callback(
        lambda f: print(f"Published: {f.result()}")
    )
 
    return {
        "task_id": task_id,
        "status": "queued",
        "estimated_completion": "30-120 seconds",
        "status_url": f"/tasks/{task_id}/status",
    }
 
@app.get("/tasks/{task_id}/status")
async def get_task_status(task_id: str):
    """Poll task processing status."""
    result = await fetch_task_result(task_id)
    return result

# worker.py
import json
import time
import httpx
from google.cloud import pubsub_v1, firestore
import google.generativeai as genai
 
genai.configure(api_key="YOUR_GEMINI_API_KEY")
model = genai.GenerativeModel("gemini-2.5-pro")
 
subscriber = pubsub_v1.SubscriberClient()
db = firestore.Client()
 
PROJECT_ID = "your-project-id"
SUBSCRIPTION_ID = "gemini-tasks-sub"
subscription_path = subscriber.subscription_path(PROJECT_ID, SUBSCRIPTION_ID)
 
def build_prompt(task_type: str, content: str) -> str:
    """Build a task-appropriate prompt for Gemini."""
    prompts = {
        "summarize": f"""Summarize the following content concisely.
Provide 3-5 key bullet points, followed by a one-paragraph summary.
 
Content:
{content}""",
        "classify": f"""Classify the following text into one of these categories:
[Technology, Business, Entertainment, Sports, Politics, Other]
 
Respond in JSON format only:
{{"category": "CategoryName", "confidence": 0.0-1.0, "reason": "explanation"}}
 
Text:
{content}""",
        "analyze": f"""Perform a detailed analysis of the following content:
1. Identify main themes
2. Sentiment analysis (Positive/Negative/Neutral)
3. Extract top 10 keywords
4. Suggest improvements (if applicable)
 
Content:
{content}""",
    }
    return prompts.get(task_type, f"Process the following:\n{content}")
 
def process_message(message: pubsub_v1.subscriber.message.Message) -> None:
    """
    Receive a Pub/Sub message and process it with Gemini API.
    """
    try:
        data = json.loads(message.data.decode("utf-8"))
        task_id = data["task_id"]
 
        print(f"Processing task: {task_id}, type: {data['task_type']}")
 
        # Update status in Firestore
        task_ref = db.collection("tasks").document(task_id)
        task_ref.set({
            "status": "processing",
            "started_at": firestore.SERVER_TIMESTAMP,
        }, merge=True)
 
        # Call Gemini API
        prompt = build_prompt(data["task_type"], data["content"])
 
        response = model.generate_content(
            prompt,
            generation_config={
                "temperature": 0.3,
                "max_output_tokens": 2048,
            }
        )
 
        result_text = response.text
 
        # Save result to Firestore
        task_ref.set({
            "status": "completed",
            "result": result_text,
            "completed_at": firestore.SERVER_TIMESTAMP,
            "model": "gemini-2.5-pro",
        }, merge=True)
 
        # Send webhook callback if configured
        if data.get("callback_url"):
            send_webhook_callback(
                data["callback_url"],
                task_id,
                result_text
            )
 
        # Ack the message on success
        message.ack()
        print(f"Task {task_id} completed successfully")
 
    except Exception as e:
        print(f"Error processing task: {e}")
 
        if "task_id" in data:
            task_ref = db.collection("tasks").document(data["task_id"])
            task_ref.set({
                "status": "failed",
                "error": str(e),
                "failed_at": firestore.SERVER_TIMESTAMP,
            }, merge=True)
 
        # Nack triggers automatic Pub/Sub retry
        message.nack()
 
def send_webhook_callback(url: str, task_id: str, result: str) -> None:
    """Notify the caller of task completion via webhook."""
    try:
        with httpx.Client(timeout=10.0) as client:
            client.post(url, json={
                "task_id": task_id,
                "status": "completed",
                "result": result,
            })
    except Exception as e:
        print(f"Webhook callback failed: {e}")
 
# Start the subscriber
streaming_pull_future = subscriber.subscribe(
    subscription_path,
    callback=process_message,
    flow_control=pubsub_v1.types.FlowControl(max_messages=10),
)
 
print(f"Listening for messages on {subscription_path}")
with subscriber:
    try:
        streaming_pull_future.result(timeout=None)
    except KeyboardInterrupt:
        streaming_pull_future.cancel()

# webhook_receiver.py
import hmac
import hashlib
import json
from fastapi import FastAPI, Request, HTTPException, Header
from typing import Annotated
 
app = FastAPI()
 
WEBHOOK_SECRET = "your-webhook-secret-key"
 
def verify_webhook_signature(payload: bytes, signature: str) -> bool:
    """
    Verify webhook signature to prevent spoofing attacks.
    Uses HMAC-SHA256 — the industry standard for webhook security.
    """
    expected = hmac.new(
        WEBHOOK_SECRET.encode(),
        payload,
        hashlib.sha256
    ).hexdigest()
    return hmac.compare_digest(f"sha256={expected}", signature)
 
@app.post("/webhook/gemini-result")
async def receive_gemini_result(
    request: Request,
    x_signature: Annotated[str | None, Header()] = None,
):
    """
    Receive webhook notification of Gemini processing completion.
    """
    payload = await request.body()
 
    # Signature verification is mandatory for security
    if x_signature and not verify_webhook_signature(payload, x_signature):
        raise HTTPException(status_code=401, detail="Invalid signature")
 
    data = json.loads(payload)
    task_id = data.get("task_id")
    status = data.get("status")
    result = data.get("result")
 
    if status == "completed":
        await handle_completed_result(task_id, result)
    elif status == "failed":
        await handle_failed_task(task_id, data.get("error"))
 
    # Always return 200 immediately — webhook golden rule
    return {"received": True}
 
async def handle_completed_result(task_id: str, result: str):
    """Business logic for completed tasks."""
    # Update database, send notifications, trigger downstream workflows
    print(f"Task {task_id} completed. Result length: {len(result)}")

# tasks.py — ARQ worker definitions
import google.generativeai as genai
from arq import create_pool
from arq.connections import RedisSettings
from typing import Any
 
genai.configure(api_key="YOUR_GEMINI_API_KEY")
model = genai.GenerativeModel("gemini-2.5-flash")
 
async def process_with_gemini(ctx: dict, content: str, task_type: str) -> dict:
    """
    ARQ worker function: calls Gemini API and returns the result.
 
    Args:
        ctx: ARQ context (includes Redis connection)
        content: Content to process
        task_type: Task type (summarize/classify/analyze)
 
    Returns:
        Dict with processing results
 
    Expected output example:
        {"status": "success", "result": "...", "tokens_used": 450}
    """
    response = model.generate_content(
        f"Task: {task_type}\n\nContent:\n{content}",
        generation_config={"temperature": 0.3}
    )
 
    return {
        "status": "success",
        "result": response.text,
        "tokens_used": response.usage_metadata.total_token_count,
    }
 
class WorkerSettings:
    """ARQ worker configuration."""
    functions = [process_with_gemini]
    redis_settings = RedisSettings(host="localhost", port=6379)
 
    # Concurrency limit — tune based on Gemini API rate limits
    max_jobs = 5
 
    # Job timeout in seconds
    job_timeout = 120
 
    # Retry settings for transient failures
    retry_jobs = True
    max_tries = 3
 
async def enqueue_gemini_task(content: str, task_type: str) -> str:
    """
    Enqueue a Gemini processing job to the Redis queue.
    """
    redis = await create_pool(RedisSettings())
 
    job = await redis.enqueue_job(
        "process_with_gemini",
        content=content,
        task_type=task_type,
        _defer_by=0,
        _queue_name="default",
    )
 
    return job.job_id

# cloud-run-worker.yaml
apiVersion: serving.knative.dev/v1
kind: Service
metadata:
  name: gemini-worker
spec:
  template:
    metadata:
      annotations:
        # Max instances — based on Gemini API rate limits
        autoscaling.knative.dev/maxScale: "20"
        autoscaling.knative.dev/minScale: "1"
        # Always-allocated CPU for background processing
        run.googleapis.com/cpu-throttling: "false"
    spec:
      containers:
        - image: gcr.io/your-project/gemini-worker:latest
          env:
            - name: GEMINI_API_KEY
              valueFrom:
                secretKeyRef:
                  name: gemini-secrets
                  key: api-key
          resources:
            limits:
              cpu: "2"
              memory: "2Gi"

# rate_limiter.py
import asyncio
import time
from collections import deque
 
class GeminiRateLimiter:
    """
    Sliding window rate limiter for Gemini API compliance.
 
    Gemini 2.5 Pro limits (approximate):
    - RPM: 2 (Tier 1) → 1,000 (Tier 3)
    - TPM: 32,000 (Tier 1) → 4,000,000 (Tier 3)
    """
 
    def __init__(self, max_requests_per_minute: int = 60):
        self.max_rpm = max_requests_per_minute
        self.requests = deque()
        self._lock = asyncio.Lock()
 
    async def acquire(self):
        """Acquire a request slot while respecting rate limits."""
        async with self._lock:
            now = time.time()
            window_start = now - 60
 
            while self.requests and self.requests[0] < window_start:
                self.requests.popleft()
 
            if len(self.requests) >= self.max_rpm:
                sleep_time = 60 - (now - self.requests[0])
                if sleep_time > 0:
                    await asyncio.sleep(sleep_time)
 
            self.requests.append(time.time())
 
    async def __aenter__(self):
        await self.acquire()
        return self
 
    async def __aexit__(self, *args):
        pass
 
limiter = GeminiRateLimiter(max_requests_per_minute=50)
 
async def call_gemini_with_rate_limit(prompt: str) -> str:
    async with limiter:
        response = await model.generate_content_async(prompt)
        return response.text

# retry_handler.py
import asyncio
import random
from functools import wraps
from typing import Callable, TypeVar
 
T = TypeVar("T")
 
def exponential_backoff_retry(
    max_retries: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    exceptions: tuple = (Exception,),
):
    """
    Retry decorator with exponential backoff and jitter.
 
    Handles transient Gemini API errors (429, 503, etc.) gracefully.
    Jitter prevents thundering herd problems when multiple workers
    hit a rate limit simultaneously.
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            last_exception = None
 
            for attempt in range(max_retries + 1):
                try:
                    return await func(*args, **kwargs)
                except exceptions as e:
                    last_exception = e
 
                    if attempt == max_retries:
                        raise
 
                    delay = min(
                        base_delay * (2 ** attempt) + random.uniform(0, 1),
                        max_delay
                    )
 
                    print(f"Attempt {attempt + 1} failed: {e}. "
                          f"Retrying in {delay:.2f}s...")
                    await asyncio.sleep(delay)
 
            raise last_exception
 
        return wrapper
    return decorator
 
@exponential_backoff_retry(
    max_retries=3,
    base_delay=2.0,
    exceptions=(Exception,)
)
async def reliable_gemini_call(prompt: str) -> str:
    """Gemini API call with automatic retry on failure."""
    response = await model.generate_content_async(prompt)
    return response.text

# dead_letter_handler.py
from google.cloud import pubsub_v1, firestore
from datetime import datetime
import json
 
DLQ_TOPIC = "gemini-tasks-dlq"
publisher = pubsub_v1.PublisherClient()
 
async def send_to_dead_letter_queue(
    original_message: dict,
    error: str,
    attempt_count: int
):
    """
    Route a failed message to the Dead Letter Queue.
    """
    dlq_message = {
        **original_message,
        "failure_info": {
            "error": error,
            "attempt_count": attempt_count,
            "moved_to_dlq_at": datetime.utcnow().isoformat(),
        }
    }
 
    topic_path = publisher.topic_path("your-project-id", DLQ_TOPIC)
    publisher.publish(
        topic_path,
        data=json.dumps(dlq_message).encode("utf-8"),
        original_task_id=original_message.get("task_id", "unknown"),
    )
 
    # Alert the team (e.g., Slack notification)
    await send_alert_to_slack(
        f"🚨 DLQ: task_id={original_message.get('task_id')} "
        f"failed after {attempt_count} attempts. Error: {error}"
    )
 
    print(f"Message sent to DLQ: {original_message.get('task_id')}")

# tracing.py
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
 
provider = TracerProvider()
exporter = CloudTraceSpanExporter(project_id="your-project-id")
processor = BatchSpanProcessor(exporter)
provider.add_span_processor(processor)
trace.set_tracer_provider(provider)
 
tracer = trace.get_tracer(__name__)
 
async def traced_gemini_processing(task_id: str, content: str) -> str:
    """
    Gemini API call with distributed tracing.
    Visualizes end-to-end latency in Cloud Trace.
    """
    with tracer.start_as_current_span("gemini-api-call") as span:
        span.set_attribute("task.id", task_id)
        span.set_attribute("content.length", len(content))
        span.set_attribute("model", "gemini-2.5-pro")
 
        try:
            response = await model.generate_content_async(content)
 
            span.set_attribute("tokens.total",
                               response.usage_metadata.total_token_count)
            span.set_attribute("status", "success")
 
            return response.text
 
        except Exception as e:
            span.set_attribute("status", "error")
            span.set_attribute("error.message", str(e))
            span.record_exception(e)
            raise

# metrics.py
from google.cloud import monitoring_v3
import time
 
client = monitoring_v3.MetricServiceClient()
project_name = f"projects/your-project-id"
 
def record_gemini_latency(latency_seconds: float, task_type: str):
    """
    Record Gemini API call latency to Cloud Monitoring.
    Enables Grafana dashboards and automated alerting.
    """
    series = monitoring_v3.TimeSeries()
    series.metric.type = "custom.googleapis.com/gemini/api_latency"
    series.metric.labels["task_type"] = task_type
 
    now = time.time()
    seconds = int(now)
    nanos = int((now - seconds) * 10 ** 9)
    interval = monitoring_v3.TimeInterval(
        {"end_time": {"seconds": seconds, "nanos": nanos}}
    )
 
    point = monitoring_v3.Point(
        {"interval": interval, "value": {"double_value": latency_seconds}}
    )
    series.points = [point]
 
    client.create_time_series(
        request={"name": project_name, "time_series": [series]}
    )

# cache_middleware.py
import hashlib
import json
import redis.asyncio as redis
from typing import Optional
 
redis_client = redis.from_url("redis://localhost:6379")
CACHE_TTL = 3600  # 1 hour
 
def generate_cache_key(content: str, task_type: str) -> str:
    """Generate a deterministic cache key from content + task type."""
    payload = f"{task_type}:{content}"
    return f"gemini:cache:{hashlib.sha256(payload.encode()).hexdigest()}"
 
async def get_cached_result(content: str, task_type: str) -> Optional[str]:
    """Retrieve a cached result if available."""
    key = generate_cache_key(content, task_type)
    cached = await redis_client.get(key)
    return cached.decode() if cached else None
 
async def cache_result(content: str, task_type: str, result: str):
    """Store a result in cache."""
    key = generate_cache_key(content, task_type)
    await redis_client.setex(key, CACHE_TTL, result.encode())
 
async def get_or_generate(content: str, task_type: str) -> dict:
    """
    Cache-first Gemini API call pattern.
 
    Expected output (cache hit):
    {"result": "...", "cache_hit": True, "tokens_saved": 450}
    """
    cached = await get_cached_result(content, task_type)
    if cached:
        return {"result": cached, "cache_hit": True, "tokens_saved": 450}
 
    response = await model.generate_content_async(
        f"Task: {task_type}\n\n{content}"
    )
    result = response.text
 
    await cache_result(content, task_type, result)
 
    return {
        "result": result,
        "cache_hit": False,
        "tokens_used": response.usage_metadata.total_token_count,
    }

# pipeline_stages.py
import json
from google.cloud import pubsub_v1
import google.generativeai as genai
 
genai.configure(api_key="YOUR_GEMINI_API_KEY")
model = genai.GenerativeModel("gemini-2.5-pro")
publisher = pubsub_v1.PublisherClient()
PROJECT_ID = "your-project-id"
 
def stage_summarize(message: pubsub_v1.subscriber.message.Message) -> None:
    """
    Stage 1: Summarize extracted document text.
    Publishes results to the classify-stage topic.
    """
    data = json.loads(message.data.decode("utf-8"))
    task_id = data["task_id"]
    raw_text = data["extracted_text"]
 
    response = model.generate_content(
        f"Summarize this document in 3-5 bullet points:\n\n{raw_text}",
        generation_config={"temperature": 0.2, "max_output_tokens": 512}
    )
 
    # Forward results to the next stage
    next_stage_message = {
        "task_id": task_id,
        "summary": response.text,
        "original_text": raw_text,
        "stage": "summarized",
    }
 
    topic_path = publisher.topic_path(PROJECT_ID, "gemini-classify-stage")
    publisher.publish(
        topic_path,
        data=json.dumps(next_stage_message).encode("utf-8"),
    )
 
    message.ack()
    print(f"Stage 1 (summarize) complete for task: {task_id}")
 
def stage_classify_and_finalize(
    message: pubsub_v1.subscriber.message.Message
) -> None:
    """
    Stage 2: Classify summary and generate action items.
    Publishes final results to the results topic.
    """
    data = json.loads(message.data.decode("utf-8"))
    task_id = data["task_id"]
    summary = data["summary"]
 
    # Batch both classification and action item generation in one call
    response = model.generate_content(
        f"""Analyze this summary and respond in JSON only:
{{
  "category": "one of: Contract/Report/Email/Technical/Other",
  "priority": "High/Medium/Low",
  "action_items": ["action1", "action2", "action3"],
  "deadline_mentioned": true or false
}}
 
Summary:
{summary}""",
        generation_config={"temperature": 0.1, "max_output_tokens": 256}
    )
 
    final_result = {
        "task_id": task_id,
        "summary": summary,
        "analysis": response.text,
        "pipeline_completed": True,
    }
 
    topic_path = publisher.topic_path(PROJECT_ID, "gemini-results")
    publisher.publish(
        topic_path,
        data=json.dumps(final_result).encode("utf-8"),
    )
 
    message.ack()
    print(f"Stage 2 (classify+finalize) complete for task: {task_id}")

# fan_out_publisher.py
import asyncio
import json
import uuid
from google.cloud import pubsub_v1
 
publisher = pubsub_v1.PublisherClient()
PROJECT_ID = "your-project-id"
 
async def fan_out_analysis(content: str, user_id: str) -> str:
    """
    Dispatch the same content to multiple specialized workers simultaneously.
 
    This pattern is useful when you need:
    - Sentiment analysis + keyword extraction + category classification
    all running in parallel for the same document.
 
    Expected behavior: all three results arrive within the time
    of the slowest single analysis (not 3x the time).
    """
    parent_task_id = str(uuid.uuid4())
 
    # Define parallel subtasks
    subtasks = [
        {"type": "sentiment", "topic": "gemini-sentiment-worker"},
        {"type": "keywords", "topic": "gemini-keyword-worker"},
        {"type": "category", "topic": "gemini-category-worker"},
    ]
 
    for subtask in subtasks:
        message = {
            "parent_task_id": parent_task_id,
            "subtask_type": subtask["type"],
            "content": content,
            "user_id": user_id,
            "total_subtasks": len(subtasks),
        }
 
        topic_path = publisher.topic_path(PROJECT_ID, subtask["topic"])
        publisher.publish(
            topic_path,
            data=json.dumps(message).encode("utf-8"),
            parent_task_id=parent_task_id,
        )
 
    print(f"Fan-out dispatched {len(subtasks)} subtasks for: {parent_task_id}")
    return parent_task_id

# main.tf — Event-driven Gemini pipeline infrastructure
 
# Pub/Sub topics
resource "google_pubsub_topic" "gemini_tasks" {
  name = "gemini-tasks"
  project = var.project_id
 
  # Retain undelivered messages for 7 days
  message_retention_duration = "604800s"
}
 
resource "google_pubsub_topic" "gemini_results" {
  name    = "gemini-results"
  project = var.project_id
}
 
resource "google_pubsub_topic" "gemini_dlq" {
  name    = "gemini-tasks-dlq"
  project = var.project_id
}
 
# Subscription with dead letter policy
resource "google_pubsub_subscription" "gemini_tasks_sub" {
  name  = "gemini-tasks-sub"
  topic = google_pubsub_topic.gemini_tasks.name
  project = var.project_id
 
  # Message delivery timeout
  ack_deadline_seconds = 120
 
  # Retry policy with backoff
  retry_policy {
    minimum_backoff = "10s"
    maximum_backoff = "600s"
  }
 
  # Route to DLQ after 5 failed attempts
  dead_letter_policy {
    dead_letter_topic     = google_pubsub_topic.gemini_dlq.id
    max_delivery_attempts = 5
  }
}
 
# Cloud Run worker service
resource "google_cloud_run_v2_service" "gemini_worker" {
  name     = "gemini-worker"
  location = var.region
  project  = var.project_id
 
  template {
    scaling {
      min_instance_count = 1
      max_instance_count = 20
    }
 
    containers {
      image = "gcr.io/${var.project_id}/gemini-worker:latest"
 
      resources {
        limits = {
          cpu    = "2"
          memory = "2Gi"
        }
        cpu_idle = false  # Always-on CPU for background processing
      }
 
      env {
        name = "GEMINI_API_KEY"
        value_source {
          secret_key_ref {
            secret  = google_secret_manager_secret.gemini_key.secret_id
            version = "latest"
          }
        }
      }
    }
  }
}

# Install and start the Pub/Sub emulator
gcloud components install pubsub-emulator
gcloud beta emulators pubsub start --project=local-project
 
# In a separate terminal, export the emulator endpoint
export PUBSUB_EMULATOR_HOST=localhost:8085
 
# Run your worker — it automatically connects to the emulator
python worker.py

# test_pipeline.py — integration test using the emulator
import os
import json
import pytest
from google.cloud import pubsub_v1
 
# The emulator is active when this env var is set
assert os.getenv("PUBSUB_EMULATOR_HOST"), "Start Pub/Sub emulator first"
 
PROJECT_ID = "local-project"
TOPIC_ID = "gemini-tasks-test"
SUBSCRIPTION_ID = "gemini-tasks-test-sub"
 
@pytest.fixture
def setup_pubsub():
    """Create test topic and subscription for each test run."""
    publisher = pubsub_v1.PublisherClient()
    subscriber = pubsub_v1.SubscriberClient()
 
    topic_path = publisher.topic_path(PROJECT_ID, TOPIC_ID)
    subscription_path = subscriber.subscription_path(
        PROJECT_ID, SUBSCRIPTION_ID
    )
 
    publisher.create_topic(request={"name": topic_path})
    subscriber.create_subscription(
        request={"name": subscription_path, "topic": topic_path}
    )
 
    yield publisher, subscriber, topic_path, subscription_path
 
    # Cleanup after test
    subscriber.delete_subscription(
        request={"subscription": subscription_path}
    )
    publisher.delete_topic(request={"topic": topic_path})
 
def test_message_roundtrip(setup_pubsub):
    """Verify messages are published and received correctly."""
    publisher, subscriber, topic_path, sub_path = setup_pubsub
 
    test_payload = {"task_id": "test-123", "content": "Hello", "task_type": "summarize"}
    publisher.publish(topic_path, data=json.dumps(test_payload).encode())
 
    response = subscriber.pull(
        request={"subscription": sub_path, "max_messages": 1}
    )
 
    assert len(response.received_messages) == 1
    received = json.loads(response.received_messages[0].message.data)
    assert received["task_id"] == "test-123"
    print("✅ Message roundtrip test passed")

from google.cloud import firestore
 
def claim_task(task_id: str) -> bool:
    """Return True if we successfully claimed the task, False if it's already taken."""
    db = firestore.Client()
    task_ref = db.collection("tasks").document(task_id)
 
    @firestore.transactional
    def attempt(tx):
        snap = task_ref.get(transaction=tx)
        if snap.exists and snap.get("status") in ("processing", "completed"):
            return False
        tx.set(task_ref, {
            "status": "processing",
            "claimed_at": firestore.SERVER_TIMESTAMP,
        }, merge=True)
        return True
 
    return attempt(db.transaction())

# slack_dlq_notifier.py — Slack notification with a human gate for DLQ replay
import json
from slack_sdk.webhook import WebhookClient
 
SLACK_URL = "https://hooks.slack.com/services/YOUR/WEBHOOK"
 
def notify_dlq_to_human(dlq_message: dict):
    webhook = WebhookClient(SLACK_URL)
    body = dlq_message.get("failure_info", {})
    webhook.send(
        text=f"DLQ: {dlq_message.get('task_id')}",
        blocks=[
            {"type": "header", "text": {"type": "plain_text",
                "text": "🚨 Gemini pipeline DLQ"}},
            {"type": "section", "fields": [
                {"type": "mrkdwn", "text": f"*task_id*\n{dlq_message.get('task_id')}"},
                {"type": "mrkdwn", "text": f"*error*\n```{body.get('error', '')[:200]}```"},
                {"type": "mrkdwn", "text": f"*attempts*\n{body.get('attempt_count')}"},
            ]},
            {"type": "actions", "elements": [
                {"type": "button", "action_id": "dlq_replay",
                 "text": {"type": "plain_text", "text": "✅ Replay"},
                 "style": "primary",
                 "value": json.dumps({"task_id": dlq_message.get("task_id")})},
                {"type": "button", "action_id": "dlq_discard",
                 "text": {"type": "plain_text", "text": "🗑 Discard"},
                 "style": "danger",
                 "value": json.dumps({"task_id": dlq_message.get("task_id")})},
            ]},
        ]
    )