Building Production AI Data Pipelines with Gemini API and Apache Airflow: A

# plugins/gemini_operators.py
from __future__ import annotations
 
import json
import time
import logging
from typing import Any, Sequence
 
import google.generativeai as genai
from airflow.models import BaseOperator
from airflow.utils.context import Context
from airflow.models import Variable
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
 
logger = logging.getLogger(__name__)
 
class GeminiGenerateOperator(BaseOperator):
    """
    Custom Airflow Operator that calls the Gemini API for text generation.
 
    :param prompt_template: Prompt template string (Jinja2-compatible)
    :param input_key: XCom key to pull input data from the previous task
    :param output_key: XCom key to push results to downstream tasks
    :param model_name: Gemini model name to use
    :param max_output_tokens: Maximum number of output tokens
    :param temperature: Sampling temperature (lower = more deterministic)
    """
 
    template_fields: Sequence[str] = ("prompt_template",)
 
    def __init__(
        self,
        *,
        prompt_template: str,
        input_key: str = "data",
        output_key: str = "result",
        model_name: str | None = None,
        max_output_tokens: int | None = None,
        temperature: float = 0.2,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.prompt_template = prompt_template
        self.input_key = input_key
        self.output_key = output_key
        self.model_name = model_name or Variable.get("gemini_model", "gemini-2.5-pro")
        self.max_output_tokens = max_output_tokens or int(
            Variable.get("gemini_max_tokens", 8192)
        )
        self.temperature = temperature
 
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=2, min=4, max=60),
        retry=retry_if_exception_type((Exception,)),
        reraise=True,
    )
    def _call_gemini(self, prompt: str) -> dict[str, Any]:
        """Call the Gemini API with exponential backoff on failure."""
        api_key = Variable.get("gemini_api_key")
        genai.configure(api_key=api_key)
 
        model = genai.GenerativeModel(
            model_name=self.model_name,
            generation_config=genai.types.GenerationConfig(
                temperature=self.temperature,
                max_output_tokens=self.max_output_tokens,
            ),
        )
 
        start_time = time.time()
        response = model.generate_content(prompt)
        elapsed = time.time() - start_time
 
        result = {
            "text": response.text,
            "model": self.model_name,
            "prompt_tokens": response.usage_metadata.prompt_token_count,
            "output_tokens": response.usage_metadata.candidates_token_count,
            "elapsed_seconds": round(elapsed, 2),
        }
 
        logger.info(
            f"Gemini API call complete | "
            f"Model: {self.model_name} | "
            f"Prompt tokens: {result['prompt_tokens']} | "
            f"Output tokens: {result['output_tokens']} | "
            f"Latency: {elapsed:.2f}s"
        )
        return result
 
    def execute(self, context: Context) -> dict[str, Any]:
        task_instance = context["task_instance"]
        input_data = task_instance.xcom_pull(key=self.input_key) or ""
 
        # Airflow expands Jinja2 template fields before execute() is called
        prompt = self.prompt_template.format(input=input_data)
 
        logger.info(f"Calling Gemini API | Prompt length: {len(prompt)} characters")
        result = self._call_gemini(prompt)
 
        task_instance.xcom_push(key=self.output_key, value=result)
        return result

# dags/document_processing_pipeline.py
from __future__ import annotations
 
from datetime import datetime, timedelta
from airflow.decorators import dag, task
from airflow.providers.google.cloud.operators.gcs import GCSListObjectsOperator
from airflow.utils.trigger_rule import TriggerRule
import json
import logging
 
logger = logging.getLogger(__name__)
 
DEFAULT_ARGS = {
    "owner": "data-team",
    "depends_on_past": False,
    "retries": 3,
    "retry_delay": timedelta(minutes=5),
    "retry_exponential_backoff": True,
    "max_retry_delay": timedelta(minutes=30),
    "email_on_failure": True,
    "email": ["alert@yourcompany.com"],
}
 
@dag(
    dag_id="gemini_document_processing",
    default_args=DEFAULT_ARGS,
    description="Analyze GCS documents with Gemini and save to BigQuery",
    schedule="0 2 * * *",  # Daily at 02:00 UTC
    start_date=datetime(2026, 4, 1),
    catchup=False,
    max_active_runs=1,
    tags=["gemini", "nlp", "production"],
)
def document_processing_pipeline():
 
    # Step 1: List files for today's date partition
    list_files = GCSListObjectsOperator(
        task_id="list_documents",
        bucket="your-input-bucket",
        prefix="documents/{{ ds }}/",
        gcp_conn_id="google_cloud_default",
    )
 
    @task
    def process_documents_batch(file_list: list[str]) -> list[dict]:
        """Process documents with rate-limit-aware batching."""
        import google.generativeai as genai
        from airflow.models import Variable
        from google.cloud import storage
        import time
 
        api_key = Variable.get("gemini_api_key")
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel("gemini-2.5-pro")
 
        results = []
        batch_size = 10  # Pause every 10 requests to respect RPM limits
 
        for i, file_path in enumerate(file_list or []):
            try:
                client = storage.Client()
                bucket = client.bucket("your-input-bucket")
                blob = bucket.blob(file_path)
                content = blob.download_as_text()
 
                prompt = f"""Analyze the document below and return a JSON response.
 
Document:
{content[:50000]}
 
Output format:
{{
  "summary": "Summary in under 200 words",
  "category": "Document category",
  "key_entities": ["entity1", "entity2"],
  "sentiment": "positive/neutral/negative",
  "action_items": ["action1", "action2"]
}}"""
 
                response = model.generate_content(prompt)
                parsed = json.loads(response.text)
 
                results.append({
                    "file_path": file_path,
                    "status": "success",
                    **parsed,
                    "processed_at": datetime.utcnow().isoformat(),
                })
 
            except json.JSONDecodeError as e:
                logger.warning(f"JSON parse error: {file_path} | {e}")
                results.append({"file_path": file_path, "status": "parse_error"})
            except Exception as e:
                logger.error(f"Processing error: {file_path} | {e}")
                results.append({"file_path": file_path, "status": "error", "error": str(e)})
 
            # Brief pause every 10 requests to avoid rate limit errors
            if (i + 1) % batch_size == 0:
                time.sleep(1)
 
        return results
 
    @task
    def save_to_bigquery(results: list[dict]) -> dict:
        """Write successful results to BigQuery."""
        from google.cloud import bigquery
 
        client = bigquery.Client()
        table_id = "your-project.your_dataset.document_analysis"
 
        success_records = [r for r in results if r.get("status") == "success"]
 
        if success_records:
            errors = client.insert_rows_json(table_id, success_records)
            if errors:
                raise ValueError(f"BigQuery write error: {errors}")
 
        summary = {
            "total": len(results),
            "success": len(success_records),
            "error": len(results) - len(success_records),
        }
        logger.info(f"BigQuery write complete: {summary}")
        return summary
 
    @task(trigger_rule=TriggerRule.ALL_DONE)
    def send_completion_report(summary: dict) -> None:
        """Send a Slack notification when the pipeline finishes."""
        import requests
        from airflow.models import Variable
 
        webhook_url = Variable.get("slack_webhook_url", default_var=None)
        if not webhook_url:
            logger.info("Slack webhook not configured — skipping notification")
            return
 
        message = (
            f"✅ Document processing complete\n"
            f"Total: {summary['total']}\n"
            f"Success: {summary['success']}\n"
            f"Errors: {summary['error']}"
        )
        requests.post(webhook_url, json={"text": message})
 
    files = list_files.output
    processed = process_documents_batch(files)
    summary = save_to_bigquery(processed)
    send_completion_report(summary)
 
document_processing_pipeline()

# plugins/rate_limiter.py
import time
import threading
from collections import deque
 
class TokenBucketRateLimiter:
    """
    Thread-safe token bucket rate limiter for Gemini API calls.
    Enforces both RPM (requests per minute) and TPM (tokens per minute) limits.
    """
 
    def __init__(self, rpm: int = 60, tpm: int = 1_000_000):
        self.rpm = rpm
        self.tpm = tpm
        self._request_times: deque = deque()
        self._token_times: deque = deque()
        self._lock = threading.Lock()
 
    def acquire(self, estimated_tokens: int = 1000) -> None:
        """
        Block until it is safe to make an API call.
 
        :param estimated_tokens: Estimated token count (prompt + expected output)
        """
        with self._lock:
            now = time.time()
            minute_ago = now - 60
 
            while self._request_times and self._request_times[0] < minute_ago:
                self._request_times.popleft()
            while self._token_times and self._token_times[0][0] < minute_ago:
                self._token_times.popleft()
 
            # Enforce RPM limit
            if len(self._request_times) >= self.rpm:
                sleep_time = 60 - (now - self._request_times[0])
                if sleep_time > 0:
                    time.sleep(sleep_time)
 
            # Enforce TPM limit
            current_tpm = sum(t[1] for t in self._token_times)
            if current_tpm + estimated_tokens > self.tpm:
                sleep_time = 60 - (now - self._token_times[0][0])
                if sleep_time > 0:
                    time.sleep(sleep_time)
 
            self._request_times.append(time.time())
            self._token_times.append((time.time(), estimated_tokens))

# plugins/gemini_sensor.py
from airflow.sensors.base import BaseSensorOperator
from airflow.utils.context import Context
import google.generativeai as genai
from airflow.models import Variable
 
class GeminiAPIHealthSensor(BaseSensorOperator):
    """
    Pokes the Gemini API with a lightweight request before the main DAG runs.
    Blocks execution until the API is confirmed healthy.
    """
 
    def poke(self, context: Context) -> bool:
        try:
            api_key = Variable.get("gemini_api_key")
            genai.configure(api_key=api_key)
            model = genai.GenerativeModel("gemini-2.5-flash")
            # Use Flash here to keep health-check costs near zero
            response = model.generate_content("ping")
            return bool(response.text)
        except Exception as e:
            self.log.warning(f"Gemini API health check failed: {e}")
            return False

# plugins/cost_tracker.py
from dataclasses import dataclass, field
from typing import ClassVar
import logging
 
logger = logging.getLogger(__name__)
 
@dataclass
class GeminiCostTracker:
    """
    Tracks Gemini API token usage and estimated cost in real time.
 
    Pricing as of April 2026 (subject to change — always verify at ai.google.dev):
    - gemini-2.5-pro: $3.50/1M input tokens, $10.50/1M output tokens
    - gemini-2.5-flash: $0.15/1M input tokens, $0.60/1M output tokens
    """
 
    PRICING: ClassVar[dict] = {
        "gemini-2.5-pro": {"input": 3.50 / 1_000_000, "output": 10.50 / 1_000_000},
        "gemini-2.5-flash": {"input": 0.15 / 1_000_000, "output": 0.60 / 1_000_000},
        "gemini-2.0-flash": {"input": 0.10 / 1_000_000, "output": 0.40 / 1_000_000},
    }
 
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_cost_usd: float = 0.0
    model_usage: dict = field(default_factory=dict)
 
    def add_usage(self, model: str, input_tokens: int, output_tokens: int) -> float:
        """Record usage for a single API call and return the cost for that call."""
        pricing = self.PRICING.get(model, self.PRICING["gemini-2.5-pro"])
        cost = input_tokens * pricing["input"] + output_tokens * pricing["output"]
 
        self.total_input_tokens += input_tokens
        self.total_output_tokens += output_tokens
        self.total_cost_usd += cost
 
        if model not in self.model_usage:
            self.model_usage[model] = {
                "requests": 0, "input_tokens": 0, "output_tokens": 0, "cost_usd": 0.0
            }
 
        self.model_usage[model]["requests"] += 1
        self.model_usage[model]["input_tokens"] += input_tokens
        self.model_usage[model]["output_tokens"] += output_tokens
        self.model_usage[model]["cost_usd"] += cost
 
        return cost
 
    def check_threshold(self, threshold_usd: float) -> bool:
        """Return True if cumulative cost has exceeded the threshold."""
        if self.total_cost_usd > threshold_usd:
            logger.warning(
                f"⚠️ Cost threshold exceeded: ${self.total_cost_usd:.4f} > ${threshold_usd}"
            )
            return True
        return False
 
    def get_report(self) -> dict:
        return {
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "total_tokens": self.total_input_tokens + self.total_output_tokens,
            "total_cost_usd": round(self.total_cost_usd, 6),
            "model_breakdown": self.model_usage,
        }

@task
def check_daily_cost_and_adjust(tracker_report: dict) -> None:
    """
    Automatically downgrade to a cheaper model when the daily budget is exceeded,
    and restore the high-quality model when costs are comfortably within budget.
    """
    from airflow.models import Variable
 
    threshold = float(Variable.get("cost_alert_threshold_usd", 50))
    current_cost = tracker_report.get("total_cost_usd", 0)
 
    if current_cost > threshold:
        Variable.set("gemini_model", "gemini-2.5-flash")
        logger.warning(
            f"Cost ${current_cost:.2f} exceeded threshold ${threshold}. "
            f"Switched to gemini-2.5-flash."
        )
    elif current_cost < threshold * 0.5:
        Variable.set("gemini_model", "gemini-2.5-pro")
        logger.info(
            f"Cost ${current_cost:.2f} well within budget. "
            f"Restored gemini-2.5-pro."
        )

# dags/multilingual_content_pipeline.py
from airflow.decorators import dag, task
from datetime import datetime, timedelta
import logging
 
logger = logging.getLogger(__name__)
 
@dag(
    dag_id="gemini_multilingual_content",
    schedule="0 6 * * 1-5",  # Monday–Friday at 06:00 UTC
    start_date=datetime(2026, 4, 1),
    catchup=False,
    default_args={
        "retries": 2,
        "retry_delay": timedelta(minutes=3),
    },
    tags=["gemini", "content", "multilingual"],
)
def multilingual_content_pipeline():
 
    @task
    def fetch_source_articles() -> list[dict]:
        """Pull draft articles from your CMS or database."""
        # Replace with your actual data source
        return [
            {"id": "001", "title": "AI Technology Trends", "content": "...", "lang": "en"},
            {"id": "002", "title": "Machine Learning Primer", "content": "...", "lang": "en"},
        ]
 
    @task
    def translate_and_enhance(articles: list[dict]) -> list[dict]:
        """Use Gemini to translate and SEO-enhance each article."""
        import google.generativeai as genai
        from airflow.models import Variable
        import json
 
        api_key = Variable.get("gemini_api_key")
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel(Variable.get("gemini_model", "gemini-2.5-pro"))
 
        results = []
        for article in articles:
            prompt = f"""Translate the following English article into natural Japanese and return SEO-optimized versions of both.
 
Original:
Title: {article['title']}
Content: {article['content']}
 
Requirements:
1. Japanese must read naturally — not like a translation
2. Weave in relevant SEO keywords organically
3. Return valid JSON only
 
Output format:
{{
  "ja_title": "Japanese title",
  "ja_content": "Japanese body text",
  "meta_description": "Under 160 characters",
  "keywords": ["keyword1", "keyword2"]
}}"""
 
            response = model.generate_content(prompt)
            translated = json.loads(response.text)
            results.append({**article, **translated})
 
        return results
 
    @task
    def publish_content(enhanced_articles: list[dict]) -> dict:
        """Push finalized articles to your CMS or content store."""
        published_count = len(enhanced_articles)
        logger.info(f"Published {published_count} articles")
        return {"published": published_count}
 
    articles = fetch_source_articles()
    enhanced = translate_and_enhance(articles)
    publish_content(enhanced)
 
multilingual_content_pipeline()

@task
def process_single_file(file_path: str) -> dict:
    """Process one file — Airflow will create one task instance per file."""
    # Implementation omitted for brevity
    return {"file": file_path, "status": "processed"}
 
# Dynamic expansion — Airflow creates one task per element automatically
file_list = ["file1.txt", "file2.txt", "file3.txt"]
processed = process_single_file.expand(file_path=file_list)

@task
def submit_batch_job(requests: list[dict]) -> str:
    """Submit a batch job to the Gemini Batch API."""
    import google.generativeai as genai
    from airflow.models import Variable
 
    api_key = Variable.get("gemini_api_key")
    genai.configure(api_key=api_key)
 
    batch_requests = [
        {
            "model": "models/gemini-2.5-pro",
            "contents": [{"parts": [{"text": req["prompt"]}]}],
        }
        for req in requests
    ]
 
    # Actual SDK call depends on your SDK version
    # batch = genai.batch_generate_content(batch_requests)
    # return batch.name
    return "batch_job_id_placeholder"
 
@task.sensor(poke_interval=60, timeout=3600)
def wait_for_batch_completion(job_id: str) -> bool:
    """Sensor that polls until the batch job finishes."""
    # Check job status and return True when complete
    return True

# Create the pool via CLI
airflow pools set gemini_api_pool 5 "Gemini API concurrency cap (5 slots)"

# Reference the pool in any Gemini-calling task
GeminiGenerateOperator(
    task_id="generate_content",
    pool="gemini_api_pool",
    pool_slots=1,
    # ...
)

# Set via environment variables (e.g., in your Docker Compose or Kubernetes manifest)
# AIRFLOW__SECRETS__BACKEND=airflow.providers.google.cloud.secrets.secret_manager.CloudSecretManagerBackend
# AIRFLOW__SECRETS__BACKEND_KWARGS={"project_id": "your-project-id", "variables_prefix": "airflow-variables"}

# Register the Gemini API key in Secret Manager (CLI example)
# echo -n "YOUR_GEMINI_API_KEY" | \
#   gcloud secrets create airflow-variables-gemini_api_key \
#   --data-file=- \
#   --project=your-project-id

# kubernetes/airflow-serviceaccount.yaml
apiVersion: v1
kind: ServiceAccount
metadata:
  name: airflow-worker
  namespace: airflow
  annotations:
    iam.gke.io/gcp-service-account: airflow-sa@your-project.iam.gserviceaccount.com