Skip to content

Add ollama support for running biome #76

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
satchelbaldwin wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Add ollama support for running biome #76

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 21 additions & 12 deletions Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,15 +34,6 @@ RUN mkdir -p /tmp/archytas && \
WORKDIR /tmp/archytas
RUN uv pip install --system .

# Install beaker-kernel from GitHub source (changes when you update REF)
ARG BEAKER_KERNEL_REF=dev
ADD https://codeload.github.com/jataware/beaker-kernel/tar.gz/$BEAKER_KERNEL_REF /tmp/beaker-kernel.tar.gz
RUN mkdir -p /tmp/beaker-kernel && \
tar -xzf /tmp/beaker-kernel.tar.gz -C /tmp/beaker-kernel --strip-components=1

WORKDIR /tmp/beaker-kernel
RUN make build && \
uv pip install --system dist/beaker_kernel-*-py3-none-any.whl

# Copy source code and install
COPY --chown=1000:1000 . /jupyter
Expand All @@ -53,6 +44,21 @@ WORKDIR /jupyter
RUN rm -f /jupyter/.beaker.conf /jupyter/.env && \
uv pip install --system -e /jupyter

# install after local package so that the arg ref supercedes the package version
ARG BEAKER_KERNEL_REF=dev
ADD https://codeload.github.com/jataware/beaker-kernel/tar.gz/$BEAKER_KERNEL_REF /tmp/beaker-kernel.tar.gz
RUN mkdir -p /tmp/beaker-kernel && \
tar -xzf /tmp/beaker-kernel.tar.gz -C /tmp/beaker-kernel --strip-components=1

WORKDIR /tmp/beaker-kernel
RUN make beaker_kernel/service/ui/index.html \
&& python -m build --no-isolation . \
&& uv pip install --system dist/beaker_kernel-*-py3-none-any.whl

# install ollama and prepull the 3.1:8b image
RUN curl -fsSL https://ollama.com/install.sh | sh
RUN bash -c "(ollama serve &) ; sleep 15 && ollama pull llama3.1:8b"

# Create runtime directories
RUN mkdir -m 777 /var/run/beaker

Expand All @@ -66,10 +72,13 @@ RUN mkdir -p /jupyter/.local /jupyter/.pqa /jupyter/.config/biopython/Bio/Entrez
chown -R 1000:1000 /jupyter/.local /jupyter/.pqa /jupyter/.config

# Set default server env variables
# ollama added here as well for the single container image
ENV BEAKER_AGENT_USER=jupyter \
BEAKER_SUBKERNEL_USER=jupyter \
BEAKER_RUN_PATH=/var/run/beaker \
BEAKER_APP=biome.app.BiomeApp
BEAKER_APP=biome.app.BiomeApp \
LLM_PROVIDER_IMPORT_PATH=archytas.models.ollama.OllamaModel \
LLM_SERVICE_MODEL=llama3.1:8b

# Service
CMD ["python", "-m", "beaker_kernel.service.server", "--ip", "0.0.0.0"]
# fixed context length - could be via ARG. entrypoint might not be the best place for this, deployment wise
CMD ["bash", "-c", "(OLLAMA_CONTEXT_LENGTH=64000 ollama serve &) ; python -m beaker_kernel.service.server --ip 0.0.0.0"]
33 changes: 17 additions & 16 deletions docker-compose.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,8 @@ services:
dockerfile: Dockerfile
args:
- ARCHYTAS_REF=main # archytas branch/tag
- BEAKER_KERNEL_REF=dev # beaker-kernel branch/tag
# pre-server fixes
- BEAKER_KERNEL_REF=b979ed280edd29d1967e4cc83d427651ac2a26d5
- DOWNLOAD_CHEMBL_DATABASE=false # set to true if you want to download the ChEMBL database (it is several GB and is used in various workflows)
ports:
- "8888:8888"
Expand All @@ -34,21 +35,21 @@ services:
- ./literature_review:/jupyter/_literature_review
networks:
- beaker
rest_api:
build:
context: .
dockerfile: ./rest_api/Dockerfile
volumes:
- ./adhoc_data:/jupyter/integrations
env_file:
- .env
ports:
- "8001:8001"
environment:
- BIOME_INTEGRATIONS_DIR=/integrations
- BIOME_DATA_DIR=/data
networks:
- beaker
# rest_api:
# build:
# context: .
# dockerfile: ./rest_api/Dockerfile
# volumes:
# - ./adhoc_data:/jupyter/integrations
# env_file:
# - .env
# ports:
# - "8001:8001"
# environment:
# - BIOME_INTEGRATIONS_DIR=/integrations
# - BIOME_DATA_DIR=/data
# networks:
# - beaker
# ui:
# image: node:20
# user: "node"
Expand Down
36 changes: 20 additions & 16 deletions src/biome/context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,21 +26,24 @@ class BiomeContext(BeakerContext):
agent_cls: "BaseAgent" = BiomeAgent

def __init__(self, beaker_kernel: "LLMKernel", config: Dict[str, Any]):
from adhoc_api.uaii import gpt_41, o3_mini, claude_37_sonnet, gemini_15_pro
ttl_seconds = 1800
drafter_config_gemini = {**gemini_15_pro, 'ttl_seconds': ttl_seconds, 'api_key': os.environ.get("GEMINI_API_KEY", "")}
drafter_config_anthropic = {**claude_37_sonnet, 'api_key': os.environ.get("ANTHROPIC_API_KEY")}
curator_config = {**o3_mini, 'api_key': os.environ.get("OPENAI_API_KEY")}
gpt_41_config = {**gpt_41, 'api_key': os.environ.get("OPENAI_API_KEY")}

# Initialize adhoc integration
adhoc_integration = BiomeAdhocIntegrations(
drafter_config=[gpt_41_config, drafter_config_anthropic, drafter_config_gemini],
curator_config=curator_config,
contextualizer_config=gpt_41_config,
logger=logger,
display_name="Specialist Agents"
)
# Disable integrations for ollama until adhoc-api UAII is updated
# TODO

# from adhoc_api.uaii import gpt_41, o3_mini, claude_37_sonnet, gemini_15_pro
# ttl_seconds = 1800
# drafter_config_gemini = {**gemini_15_pro, 'ttl_seconds': ttl_seconds, 'api_key': os.environ.get("GEMINI_API_KEY", "")}
# drafter_config_anthropic = {**claude_37_sonnet, 'api_key': os.environ.get("ANTHROPIC_API_KEY")}
# curator_config = {**o3_mini, 'api_key': os.environ.get("OPENAI_API_KEY")}
# gpt_41_config = {**gpt_41, 'api_key': os.environ.get("OPENAI_API_KEY")}

# # Initialize adhoc integration
# adhoc_integration = BiomeAdhocIntegrations(
# drafter_config=[gpt_41_config, drafter_config_anthropic, drafter_config_gemini],
# curator_config=curator_config,
# contextualizer_config=gpt_41_config,
# logger=logger,
# display_name="Specialist Agents"
# )

# Track missing API keys (must be set before super().__init__ for agent access)
self.api_key_map: Dict[str, str] = {
Expand All @@ -66,7 +69,8 @@ def __init__(self, beaker_kernel: "LLMKernel", config: Dict[str, Any]):
beaker_kernel,
self.agent_cls,
config,
integrations=[adhoc_integration]
#integrations=[adhoc_integration]
integrations=[]
)

if not isinstance(self.subkernel, PythonSubkernel):
Expand Down