Skip to content

Commit da35c4e

Browse files
authored
Merge pull request #43 from redis/claude/issue-42-20250801-1705
feat: Make OpenAI and Anthropic API base URLs configurable
2 parents 7ab7205 + be7f078 commit da35c4e

File tree

4 files changed

+30
-13
lines changed

4 files changed

+30
-13
lines changed

agent_memory_server/config.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ class Settings(BaseSettings):
5252
long_term_memory: bool = True
5353
openai_api_key: str | None = None
5454
anthropic_api_key: str | None = None
55+
openai_api_base: str | None = None
56+
anthropic_api_base: str | None = None
5557
generation_model: str = "gpt-4o"
5658
embedding_model: str = "text-embedding-3-small"
5759
port: int = 8000

agent_memory_server/llms.py

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
from openai import AsyncOpenAI
1010
from pydantic import BaseModel
1111

12+
from agent_memory_server.config import settings
13+
1214

1315
logger = logging.getLogger(__name__)
1416

@@ -203,14 +205,21 @@ def total_tokens(self) -> int:
203205
class AnthropicClientWrapper:
204206
"""Wrapper for Anthropic client"""
205207

206-
def __init__(self, api_key: str | None = None):
208+
def __init__(self, api_key: str | None = None, base_url: str | None = None):
207209
"""Initialize the Anthropic client"""
208210
anthropic_api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
211+
anthropic_api_base = base_url or os.environ.get("ANTHROPIC_API_BASE")
209212

210213
if not anthropic_api_key:
211214
raise ValueError("Anthropic API key is required")
212215

213-
self.client = anthropic.AsyncAnthropic(api_key=anthropic_api_key)
216+
if anthropic_api_base:
217+
self.client = anthropic.AsyncAnthropic(
218+
api_key=anthropic_api_key,
219+
base_url=anthropic_api_base,
220+
)
221+
else:
222+
self.client = anthropic.AsyncAnthropic(api_key=anthropic_api_key)
214223

215224
async def create_chat_completion(
216225
self,
@@ -397,9 +406,15 @@ async def get_model_client(
397406
model_config = get_model_config(model_name)
398407

399408
if model_config.provider == ModelProvider.OPENAI:
400-
model = OpenAIClientWrapper(api_key=os.environ.get("OPENAI_API_KEY"))
401-
if model_config.provider == ModelProvider.ANTHROPIC:
402-
model = AnthropicClientWrapper(api_key=os.environ.get("ANTHROPIC_API_KEY"))
409+
model = OpenAIClientWrapper(
410+
api_key=settings.openai_api_key,
411+
base_url=settings.openai_api_base,
412+
)
413+
elif model_config.provider == ModelProvider.ANTHROPIC:
414+
model = AnthropicClientWrapper(
415+
api_key=settings.anthropic_api_key,
416+
base_url=settings.anthropic_api_base,
417+
)
403418

404419
if model:
405420
_model_clients[model_name] = model

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ dependencies = [
3434
"sentence-transformers>=3.4.1",
3535
"structlog>=25.2.0",
3636
"tiktoken>=0.5.1",
37-
"transformers<=4.50.3,>=4.30.0",
37+
"transformers>=4.51.1",
3838
"uvicorn>=0.24.0",
3939
"sniffio>=1.3.1",
4040
"click>=8.1.0",

uv.lock

Lines changed: 7 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)