Skip to content

Commit 63c3f76

Browse files
abrookinsabrookins
authored
Merge pull request #94 from f4roukb/feat/bedrock-support
feat: Support AWS Bedrock Embedding Models and LLM Models
2 parents b2fe2a4 + 95d68ff commit 63c3f76

39 files changed

+4429
-1603
lines changed

‎.dockerignore‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
# ===========================================
2+
# Git
3+
# ===========================================
4+
.git
5+
.gitignore
6+
.gitattributes
7+
8+
# ===========================================
9+
# Python
10+
# ===========================================
11+
__pycache__
12+
*.py[cod]
13+
*$py.class
14+
*.so
15+
.Python
16+
.venv
17+
venv
18+
env
19+
ENV
20+
.env
21+
.env.*
22+
*.egg
23+
*.egg-info
24+
dist
25+
build
26+
eggs
27+
parts
28+
sdist
29+
develop-eggs
30+
.installed.cfg
31+
lib
32+
lib64
33+
*.manifest
34+
*.spec
35+
pip-log.txt
36+
pip-delete-this-directory.txt
37+
38+
# ===========================================
39+
# Testing
40+
# ===========================================
41+
tests
42+
pytest.ini
43+
.pytest_cache
44+
.coverage
45+
.tox
46+
.nox
47+
htmlcov
48+
coverage.xml
49+
*.cover
50+
.hypothesis
51+
52+
# ===========================================
53+
# IDE & Editors
54+
# ===========================================
55+
.idea
56+
.vscode
57+
*.swp
58+
*.swo
59+
*~
60+
.project
61+
.pydevproject
62+
.settings
63+
64+
# ===========================================
65+
# Documentation
66+
# ===========================================
67+
docs
68+
*.md
69+
!README.md
70+
mkdocs.yml
71+
72+
# ===========================================
73+
# Docker
74+
# ===========================================
75+
Dockerfile*
76+
docker-compose*.yml
77+
.docker
78+
79+
# ===========================================
80+
# CI/CD
81+
# ===========================================
82+
.github
83+
.gitlab-ci.yml
84+
.travis.yml
85+
Jenkinsfile
86+
87+
# ===========================================
88+
# Misc
89+
# ===========================================
90+
*.log
91+
*.tmp
92+
*.temp
93+
*.bak
94+
*.cache
95+
.DS_Store
96+
Thumbs.db
97+
*.png
98+
*.jpg
99+
*.jpeg
100+
*.gif
101+
*.svg
102+
103+
# ===========================================
104+
# Project specific
105+
# ===========================================
106+
manual_oauth_qa
107+
examples
108+
scripts
109+
_typos.toml
110+
CLAUDE.md
111+
LICENSE
112+
113+
# ===========================================
114+
# Linting
115+
# ===========================================
116+
.ruff_cache
117+
.mypy_cache
118+
119+
# ===========================================
120+
# Pre-commit
121+
# ===========================================
122+
.pre-commit-config.yaml
123+
.pre-commit-hooks.yaml
124+
125+
# Note: uv.lock and pyproject.toml are REQUIRED for the build
126+
# Note: agent-memory-client/ is REQUIRED (workspace dependency)

‎.env.example‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,22 @@ WINDOW_SIZE=12
1010
GENERATION_MODEL=gpt-4o-mini
1111
EMBEDDING_MODEL=text-embedding-3-small
1212

13-
# OpenAI API key
13+
# API Keys
1414
OPENAI_API_KEY=your_openai_api_key
15+
ANTHROPIC_API_KEY=your_anthropic_api_key
16+
17+
# Cloud Region
18+
REGION_NAME=your_cloud_region
19+
# AWS Cloud Credentials
20+
AWS_ACCESS_KEY_ID=your_aws_access_key_id
21+
AWS_SECRET_ACCESS_KEY=your_aws_secret_access_key
22+
AWS_SESSION_TOKEN=your_aws_session_token
23+
1524

1625
# OAuth2/JWT Authentication (for production)
1726
# OAUTH2_ISSUER_URL=https://your-auth-provider.com
1827
# OAUTH2_AUDIENCE=your-api-audience
1928
# OAUTH2_JWKS_URL=https://your-auth-provider.com/.well-known/jwks.json # Optional
2029

21-
# Development Mode (DISABLE AUTHENTICATION - DEVELOPMENT ONLY)
22-
DISABLE_AUTH=true
30+
# In development mode, you may disable authentication
31+
DISABLE_AUTH=false

‎Dockerfile‎

Lines changed: 115 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,75 @@
1-
FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim
1+
# ============================================
2+
# BUILDER BASE - Build tools for compilation
3+
# ============================================
4+
FROM python:3.12-slim-bookworm AS builder-base
5+
6+
WORKDIR /app
7+
8+
# Copy uv binary from official image
9+
COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
10+
11+
ENV UV_COMPILE_BYTECODE=1
12+
ENV UV_LINK_MODE=copy
13+
14+
# Install build tools (only needed for compilation)
15+
RUN apt-get update && apt-get install -y --no-install-recommends \
16+
build-essential \
17+
gcc \
18+
g++ \
19+
&& rm -rf /var/lib/apt/lists/*
20+
21+
# ============================================
22+
# BUILDER STANDARD - Compile standard deps
23+
# ============================================
24+
FROM builder-base AS builder-standard
25+
26+
# Create virtual environment explicitly
27+
RUN uv venv .venv
28+
29+
# Copy dependency files first for better layer caching
30+
COPY pyproject.toml uv.lock ./
31+
COPY agent-memory-client ./agent-memory-client
32+
33+
# Install dependencies into the venv (without the project)
34+
RUN --mount=type=cache,target=/root/.cache/uv \
35+
VIRTUAL_ENV=/app/.venv uv sync --frozen --no-install-project --no-dev
36+
37+
# Copy source code
38+
COPY . /app
39+
40+
# Install the project itself
41+
RUN --mount=type=cache,target=/root/.cache/uv \
42+
. .venv/bin/activate && \
43+
uv pip install --no-deps .
44+
45+
# ============================================
46+
# BUILDER AWS - Compile AWS deps
47+
# ============================================
48+
FROM builder-base AS builder-aws
49+
50+
# Create virtual environment explicitly
51+
RUN uv venv .venv
52+
53+
# Copy dependency files first for better layer caching
54+
COPY pyproject.toml uv.lock ./
55+
COPY agent-memory-client ./agent-memory-client
56+
57+
# Install dependencies into the venv (without the project)
58+
RUN --mount=type=cache,target=/root/.cache/uv \
59+
VIRTUAL_ENV=/app/.venv uv sync --frozen --no-install-project --no-dev --extra aws
60+
61+
# Copy source code
62+
COPY . /app
63+
64+
# Install the project itself
65+
RUN --mount=type=cache,target=/root/.cache/uv \
66+
. .venv/bin/activate && \
67+
uv pip install --no-deps .
68+
69+
# ============================================
70+
# RUNTIME BASE - Slim image without build tools
71+
# ============================================
72+
FROM python:3.12-slim-bookworm AS runtime-base
273

374
# OCI labels for Docker Hub and container registries
475
LABEL org.opencontainers.image.title="Redis Agent Memory Server"
@@ -11,46 +82,35 @@ LABEL org.opencontainers.image.licenses="Apache-2.0"
1182

1283
WORKDIR /app
1384

14-
ENV UV_COMPILE_BYTECODE=1
15-
ENV UV_LINK_MODE=copy
16-
17-
# Install system dependencies including build tools
18-
RUN apt-get update && apt-get install -y \
85+
# Install only runtime dependencies (curl for healthcheck)
86+
RUN apt-get update && apt-get install -y --no-install-recommends \
1987
curl \
20-
build-essential \
21-
gcc \
22-
g++ \
2388
&& rm -rf /var/lib/apt/lists/*
2489

25-
RUN --mount=type=cache,target=/root/.cache/uv \
26-
--mount=type=bind,source=uv.lock,target=uv.lock \
27-
--mount=type=bind,source=pyproject.toml,target=pyproject.toml \
28-
--mount=type=bind,source=agent-memory-client,target=agent-memory-client \
29-
uv sync --frozen --no-install-project --no-dev
90+
# Create non-root user for security
91+
RUN groupadd -r agentmemory && useradd -r -g agentmemory agentmemory
3092

31-
ADD . /app
32-
RUN --mount=type=cache,target=/root/.cache/uv \
33-
uv sync --frozen --no-dev
93+
# ============================================
94+
# STANDARD VARIANT - OpenAI/Anthropic only
95+
# ============================================
96+
FROM runtime-base AS standard
3497

35-
# Create non-root user for security
36-
RUN groupadd -r agentmemory && useradd -r -g agentmemory agentmemory && \
37-
chown -R agentmemory:agentmemory /app
98+
# Copy the virtual environment and app from builder
99+
COPY --chown=agentmemory:agentmemory --from=builder-standard /app /app
38100

39101
ENV PATH="/app/.venv/bin:$PATH"
40102

41103
# Switch to non-root user
42104
USER agentmemory
43105

44-
ENTRYPOINT []
45-
46106
EXPOSE 8000
47107

48108
HEALTHCHECK --interval=30s --timeout=10s --start-period=40s --retries=3 \
49109
CMD curl -f http://localhost:8000/v1/health || exit 1
50110

51-
# Disable auth by default for easier local development.
52-
# Override with DISABLE_AUTH=false in production.
53-
ENV DISABLE_AUTH=true
111+
# Enable authentication by default.
112+
# You may override with DISABLE_AUTH=true in development.
113+
ENV DISABLE_AUTH=false
54114

55115
# Default to development mode (no separate worker needed).
56116
# For production, override the command to remove --no-worker and run a separate task-worker container.
@@ -59,3 +119,33 @@ ENV DISABLE_AUTH=true
59119
# Production API: docker run -p 8000:8000 redislabs/agent-memory-server agent-memory api --host 0.0.0.0 --port 8000
60120
# Production Worker: docker run redislabs/agent-memory-server agent-memory task-worker --concurrency 10
61121
CMD ["agent-memory", "api", "--host", "0.0.0.0", "--port", "8000", "--no-worker"]
122+
123+
# ============================================
124+
# AWS VARIANT - Includes AWS Bedrock support
125+
# ============================================
126+
FROM runtime-base AS aws
127+
128+
# Copy the virtual environment and app from builder
129+
COPY --chown=agentmemory:agentmemory --from=builder-aws /app /app
130+
131+
ENV PATH="/app/.venv/bin:$PATH"
132+
133+
# Switch to non-root user
134+
USER agentmemory
135+
136+
EXPOSE 8000
137+
138+
HEALTHCHECK --interval=30s --timeout=10s --start-period=40s --retries=3 \
139+
CMD curl -f http://localhost:8000/v1/health || exit 1
140+
141+
# Enable authentication by default.
142+
# You may override with DISABLE_AUTH=true in development.
143+
ENV DISABLE_AUTH=false
144+
145+
# Default to development mode (no separate worker needed).
146+
# For production, override the command to remove --no-worker and run a separate task-worker container.
147+
# Examples:
148+
# Development: docker run -p 8000:8000 redislabs/agent-memory-server:aws
149+
# Production API: docker run -p 8000:8000 redislabs/agent-memory-server:aws agent-memory api --host 0.0.0.0 --port 8000
150+
# Production Worker: docker run redislabs/agent-memory-server:aws agent-memory task-worker --concurrency 10
151+
CMD ["agent-memory", "api", "--host", "0.0.0.0", "--port", "8000", "--no-worker"]

‎agent_memory_server/_aws/__init__.py‎

Whitespace-only changes.

‎agent_memory_server/_aws/clients.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
"""AWS clients for the Agent Memory Server.
2+
3+
This module contains utilities for creating and managing AWS clients.
4+
"""
5+
6+
from typing import TYPE_CHECKING
7+
8+
from boto3 import Session
9+
10+
from agent_memory_server.config import settings
11+
12+
13+
if TYPE_CHECKING:
14+
from mypy_boto3_bedrock import BedrockClient
15+
from mypy_boto3_bedrock_runtime import BedrockRuntimeClient
16+
17+
18+
def create_aws_session(
19+
region_name: str | None = None, credentials: dict[str, str] | None = None
20+
) -> Session:
21+
"""Create an AWS session.
22+
23+
Args:
24+
credentials (dict[str, str | None]): The AWS credentials to use.
25+
26+
Returns:
27+
An AWS session.
28+
"""
29+
if credentials is None:
30+
credentials = settings.aws_credentials
31+
if region_name is None:
32+
region_name = settings.aws_region
33+
return Session(region_name=region_name, **credentials)
34+
35+
36+
def create_bedrock_client(
37+
region_name: str | None = None,
38+
session: Session | None = None,
39+
) -> "BedrockClient":
40+
"""Create a Bedrock client.
41+
42+
Args:
43+
region_name (str | None): The AWS region to use.\
44+
If not provided, it will be picked up from the environment.
45+
session (Session | None): The AWS session to use.\
46+
If not provided, a new session will be created based on the environment.
47+
"""
48+
if session is None:
49+
session = create_aws_session(region_name=region_name)
50+
if region_name is None:
51+
region_name = settings.aws_region
52+
return session.client("bedrock", region_name=region_name)
53+
54+
55+
def create_bedrock_runtime_client(
56+
region_name: str | None = None,
57+
session: Session | None = None,
58+
) -> "BedrockRuntimeClient":
59+
"""Create a Bedrock runtime client.
60+
61+
Args:
62+
region_name (str | None): The AWS region to use.\
63+
If not provided, it will be picked up from the environment.
64+
session (Session | None): The AWS session to use.\
65+
If not provided, a new session will be created based on the environment.
66+
67+
Returns:
68+
A Bedrock runtime client.
69+
"""
70+
if session is None:
71+
session = create_aws_session(region_name=region_name)
72+
if region_name is None:
73+
region_name = settings.aws_region
74+
return session.client("bedrock-runtime", region_name=region_name)

0 commit comments

Comments
 (0)