# =============================================================================
# Gemma 4 CPU-first stack — base (CPU) compose file.
#
#   Services:
#     llama-swap  : OpenAI endpoint that hot-swaps Gemma 4 GGUF profiles
#     tei-embed   : HF Text-Embeddings-Inference — Persian embedder (Hakim/bge-m3)
#     tei-rerank  : HF Text-Embeddings-Inference — reranker (bge-reranker-v2-m3)
#     openwebui   : chat UI + RAG orchestrator + native-audio Pipe host
#
#   Usage (CPU):     docker compose up -d
#   Usage (GPU):     docker compose -f docker-compose.yml -f docker-compose.gpu.yml up -d
#
#   All tunables come from .env (copy one of profiles/*.env → .env first).
# =============================================================================

x-common: &common
  restart: unless-stopped
  networks: [stack]

services:
  # ---------------------------------------------------------------------------
  # LLM gateway: llama-swap renders its config from .env, then loads Gemma 4
  # llama-server on demand (one model resident at a time, idle-unloaded by ttl).
  # ---------------------------------------------------------------------------
  llama-swap:
    <<: *common
    image: ${LLAMASWAP_IMAGE:-ghcr.io/mostlygeek/llama-swap:cpu}
    container_name: llama-swap
    entrypoint: ["/bin/sh", "/config/entrypoint.sh"]
    environment:
      # Rendered into config.yaml by entrypoint.sh (sed), per-machine:
      NGL: ${NGL:-0}              # GPU layers to offload (0 = pure CPU)
      THREADS: ${THREADS:-4}      # physical cores, not hyperthreads
      CTX: ${CTX:-8192}           # context window
      DEFAULT_MODEL: ${DEFAULT_MODEL:-gemma-e4b}  # preloaded at startup (no cold first message)
      HF_TOKEN: ${HF_TOKEN:-}     # for gated Gemma 4 pulls via -hf
      HF_HOME: /models/hf-cache
      LLAMA_CACHE: /models/llama-cache   # persist GGUFs fetched by -hf across restarts
    volumes:
      - ./llama-swap:/config
      - ./models:/models
    ports:
      - "${LLAMASWAP_PORT:-8080}:8080"
    healthcheck:
      test: ["CMD-SHELL", "curl -fsS http://localhost:8080/v1/models || exit 1"]
      interval: 30s
      timeout: 5s
      retries: 3
      start_period: 40s

  # ---------------------------------------------------------------------------
  # RAG embedder (Persian). Default Hakim (FaMTEB #1); fallback bge-m3.
  # Container listens on :80. OpenAI-compatible /v1/embeddings + native /embed.
  # ---------------------------------------------------------------------------
  tei-embed:
    <<: *common
    image: ${TEI_IMAGE:-ghcr.io/huggingface/text-embeddings-inference:cpu-1.9}
    container_name: tei-embed
    command: >-
      --model-id ${TEI_EMBED_MODEL:-BAAI/bge-m3}
      --pooling ${TEI_EMBED_POOLING:-cls}
      --dtype ${TEI_DTYPE:-float16}
      --max-batch-tokens 4096
      --max-client-batch-size 8
      --auto-truncate
    environment:
      HF_TOKEN: ${HF_TOKEN:-}
    volumes:
      - ./models/tei:/data
    ports:
      - "${TEI_EMBED_PORT:-8081}:80"
    healthcheck:
      test: ["CMD-SHELL", "curl -fsS http://localhost:80/health || exit 1"]
      interval: 30s
      timeout: 5s
      retries: 5
      start_period: 60s

  # ---------------------------------------------------------------------------
  # RAG reranker (cross-encoder). bge-reranker-v2-m3 covers Persian well.
  # Exposes /rerank (Open WebUI external reranker points here).
  # ---------------------------------------------------------------------------
  tei-rerank:
    <<: *common
    image: ${TEI_IMAGE:-ghcr.io/huggingface/text-embeddings-inference:cpu-1.9}
    container_name: tei-rerank
    command: >-
      --model-id ${TEI_RERANK_MODEL:-BAAI/bge-reranker-v2-m3}
      --dtype ${TEI_DTYPE:-float16}
      --max-batch-tokens 4096
      --auto-truncate
    environment:
      HF_TOKEN: ${HF_TOKEN:-}
    volumes:
      - ./models/tei:/data
    ports:
      - "${TEI_RERANK_PORT:-8082}:80"
    healthcheck:
      test: ["CMD-SHELL", "curl -fsS http://localhost:80/health || exit 1"]
      interval: 30s
      timeout: 5s
      retries: 5
      start_period: 60s

  # ---------------------------------------------------------------------------
  # Front-end + RAG orchestrator. Talks to llama-swap (chat/vision) and the
  # two TEI services (embed/rerank). Native audio handled by the imported Pipe.
  # ---------------------------------------------------------------------------
  openwebui:
    <<: *common
    image: ghcr.io/open-webui/open-webui:main
    container_name: openwebui
    environment:
      # --- chat backend -------------------------------------------------------
      OPENAI_API_BASE_URL: http://llama-swap:8080/v1
      OPENAI_API_KEY: ${OPENAI_API_KEY:-sk-local}
      ENABLE_OLLAMA_API: "false"
      DEFAULT_MODELS: ${DEFAULT_MODEL:-gemma-e4b}
      # --- RAG embeddings via TEI (set RAG_* explicitly; OWUI does not inherit
      #     OPENAI_* — see Open WebUI issues #8697 / #22084) ------------------
      RAG_EMBEDDING_ENGINE: openai
      RAG_OPENAI_API_BASE_URL: http://tei-embed:80/v1
      RAG_OPENAI_API_KEY: ${TEI_API_KEY:-x}
      RAG_EMBEDDING_MODEL: ${TEI_EMBED_MODEL:-MCINext/Hakim}
      RAG_EMBEDDING_BATCH_SIZE: "16"
      # --- hybrid retrieval + external reranker via TEI ----------------------
      ENABLE_RAG_HYBRID_SEARCH: "true"
      RAG_RERANKING_ENGINE: external
      RAG_EXTERNAL_RERANKER_URL: http://tei-rerank:80/rerank
      RAG_EXTERNAL_RERANKER_API_KEY: ${TEI_API_KEY:-x}
      RAG_TOP_K: "8"
      RAG_TOP_K_RERANKER: "4"
      # --- native audio is the Pipe's job; turn STT off ----------------------
      AUDIO_STT_ENGINE: ""
      # --- kill hidden background LLM calls (title/tags/follow-up/autocomplete
      #     each fire an extra chat completion on the same slow CPU backend) ---
      ENABLE_TITLE_GENERATION: "false"
      ENABLE_TAGS_GENERATION: "false"
      ENABLE_FOLLOW_UP_GENERATION: "false"
      ENABLE_AUTOCOMPLETE_GENERATION: "false"
      ENABLE_RETRIEVAL_QUERY_GENERATION: "false"
      ENABLE_SEARCH_QUERY_GENERATION: "false"
      WEBUI_NAME: ${WEBUI_NAME:-Gemma 4 Local}
    volumes:
      - ./openwebui/data:/app/backend/data
      # functions/ is mounted for convenience; import the Pipe via Admin → Functions
      - ./openwebui/functions:/app/backend/data/functions-src:ro
    ports:
      - "${WEBUI_PORT:-3000}:8080"
    depends_on:
      - llama-swap
      - tei-embed
      - tei-rerank

networks:
  stack:
    driver: bridge