# ---------------------------------------------------------------
# Docker image for the Tool Router (based on the official LiteLLM image)
# ---------------------------------------------------------------
# Why use the official base?
# • It already contains the core LiteLLM binaries and entrypoint
# • Smaller image size (no extra OS packages we don’t need)
# • Keeps the same security update cadence as the official image
#
# What we add:
#   • Our Tool Router source code (tool_router_logger.py, config.yaml, etc.)
#   • Extra Python dependencies required for the Jaccard/Embedding logic
#       – numpy
#       – scikit‑learn
#   • A tiny wrapper entrypoint that runs `litellm --config /app/config.yaml`
# ---------------------------------------------------------------

# 1. Use the official LiteLLM image (based on python:3.12-slim)
FROM litellm/litellm:latest AS base

# 2. Set environment variables that control the router
ENV \
  # Port the proxy will expose (match docker‑compose)
  PORT=4000 \
  # Router defaults – can be overridden at runtime with env vars
  TOOL_ROUTER_SIMILARITY_MODE=jaccard \
  TOOL_ROUTER_TOP_K=3 \
  TOOL_ROUTER_MIN_SIMILARITY=0.15 \
  TOOL_ROUTER_ALWAYS_INCLUDE=get_time,help,list_tools \
  TOOL_ROUTER_EMBEDDING_MODEL=text-embedding-3-small

# 3. Create a non-root user (optional but recommended)
RUN useradd -m litellm && \
    mkdir -p /app && \
    chown litellm:litellm /app

# 4. Switch to that user for the rest of the steps
USER litellm
WORKDIR /app

# 5. Copy only the files that belong to the Tool Router
#    (keeps the image lean – we don’t copy .git, docs, etc.)
COPY --chown=litellm:litellm \
    tool_router_logger.py \
    config.yaml \
    test_tool_router.py \
    requirements.txt \
    ./

# 6. Install the extra Python deps that the router needs:
#       – numpy   (for cosine similarity)
#       – scikit‑learn (for the Jaccard fallback)
RUN pip install --no-cache-dir -r requirements.txt

# 7. Expose the port used by LiteLLM
EXPOSE 4000

# 8. Healthcheck - LiteLLM returns a JSON "health" endpoint
HEALTHCHECK --interval=30s --timeout=5s --start-period=20s --retries=3 \
  CMD curl -fsS "http://127.0.0.1:${PORT}/health" || exit 1

# 9. Default command - run the proxy with our config file
#     The official image’s entrypoint is preserved, we just pass the config.
CMD ["sh", "-c", "litellm --config /app/config.yaml --port ${PORT} --host 0.0.0.0"]