FROM rust:1.75.0-bookworm AS aicirt

WORKDIR /workspace

RUN rustup target add wasm32-wasi
RUN curl -fsSL https://deb.nodesource.com/setup_18.x | bash -
RUN apt-get install -y nodejs

COPY controllers controllers
COPY aicirt aicirt
COPY scripts scripts
COPY py/pyaici py/pyaici
COPY Cargo.toml Cargo.lock /workspace/

# make sure we rebuild these
RUN rm -rf controllers/jsctrl/samples/dist controllers/jsctrl/ts/dist

RUN grep -v rllm Cargo.toml > Cargo.toml.tmp && mv Cargo.toml.tmp Cargo.toml

RUN --mount=type=cache,target=/usr/local/cargo/git \
    --mount=type=cache,target=/usr/local/cargo/registry \
    --mount=type=cache,target=/workspace/target \
    cargo fetch

ARG tag=latest
ENV BUILD_TAG=$tag
RUN --mount=type=cache,target=/usr/local/cargo/git \
    --mount=type=cache,target=/usr/local/cargo/registry \
    --mount=type=cache,target=/workspace/target \
    SKIP_LLAMA_CPP=1 \
    ./scripts/release.sh && cp -r /workspace/target/dist /workspace/


FROM vllm/vllm-openai as vllm-base

# install pyaici pre-requisites
RUN pip install posix_ipc ujson

# install pyaici
RUN mkdir /tmp/pyaici
COPY py/setup.py /tmp/pyaici/
COPY py/pyaici /tmp/pyaici/pyaici
RUN cd /tmp/pyaici && pip install . && rm -rf /tmp/pyaici

# patch the vllm python files
RUN --mount=source=py/vllm/vllm,target=/tmp/vllm \
    (cd /tmp/vllm && find . -name '*.py' -print0 | tar -cf - --null -T -) | \
    tar -C /usr/local/lib/python3.10/dist-packages/vllm -xf -

# copy the controllers and aicirt
WORKDIR /vllm-workspace
RUN mkdir wasm

RUN --mount=from=aicirt,source=/workspace/dist,target=/tmp/dist \
    cp /tmp/dist/aicirt/aicirt /usr/bin/aicirt && \
    cp /tmp/dist/aici_*.wasm wasm/

RUN ls -l wasm/

ENV RUST_LOG info,tokenizers=error

ENTRYPOINT ["python3", "-m", "pyaici.vllm_server", "--enforce-eager", "--use-v2-block-manager", "--enable-chunked-prefill", "--aici-rt=/usr/bin/aicirt", "-A--restricted", "-A--wasm-timer-resolution-us=10"]

FROM vllm-base AS vllm-guidance

RUN aicirt --module wasm/aici_llguidance_ctrl.wasm --tag llguidance

FROM vllm-base as vllm-general

RUN aicirt --module wasm/aici_llguidance_ctrl.wasm --tag llguidance
RUN aicirt --module wasm/aici_pyctrl.wasm --tag pyctrl --gh-module gh:microsoft/aici/pyctrl
RUN aicirt --module wasm/aici_jsctrl.wasm --tag jsctrl --gh-module gh:microsoft/aici/jsctrl
