Skip to content

Commit b943664

Browse files
abrookinsabrookins
committed
Consolidate MODEL_CONFIGS to fix embedding dimension configuration
Move ModelProvider, ModelConfig, and MODEL_CONFIGS from llms.py to config.py to eliminate the conflicting dual definitions that were causing embedding dimension configuration failures in CI. - Unified model configuration in single location - Fixed embedding_dimensions access in main.py startup - Updated property return types to ModelConfig | None - All imports updated to use config.py definitions This resolves the AttributeError where None.embedding_dimensions was being accessed when model lookups failed due to inconsistent MODEL_CONFIGS.
1 parent 5696814 commit b943664

File tree

3 files changed

+173
-207
lines changed

3 files changed

+173
-207
lines changed

‎agent_memory_server/config.py‎

Lines changed: 166 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,49 +1,178 @@
11
import os
2+
from enum import Enum
23
from typing import Any, Literal
34

45
import yaml
56
from dotenv import load_dotenv
7+
from pydantic import BaseModel
68
from pydantic_settings import BaseSettings
79

810

911
load_dotenv()
1012

1113

14+
class ModelProvider(str, Enum):
15+
"""Type of model provider"""
16+
17+
OPENAI = "openai"
18+
ANTHROPIC = "anthropic"
19+
20+
21+
class ModelConfig(BaseModel):
22+
"""Configuration for a model"""
23+
24+
provider: ModelProvider
25+
name: str
26+
max_tokens: int
27+
embedding_dimensions: int = 1536 # Default for OpenAI ada-002
28+
29+
1230
# Model configuration mapping
1331
MODEL_CONFIGS = {
14-
"gpt-4o": {"provider": "openai", "embedding_dimensions": None},
15-
"gpt-4o-mini": {"provider": "openai", "embedding_dimensions": None},
16-
"gpt-4": {"provider": "openai", "embedding_dimensions": None},
17-
"gpt-3.5-turbo": {"provider": "openai", "embedding_dimensions": None},
18-
"text-embedding-3-small": {"provider": "openai", "embedding_dimensions": 1536},
19-
"text-embedding-3-large": {"provider": "openai", "embedding_dimensions": 3072},
20-
"text-embedding-ada-002": {"provider": "openai", "embedding_dimensions": 1536},
21-
"claude-3-opus-20240229": {"provider": "anthropic", "embedding_dimensions": None},
22-
"claude-3-sonnet-20240229": {"provider": "anthropic", "embedding_dimensions": None},
23-
"claude-3-haiku-20240307": {"provider": "anthropic", "embedding_dimensions": None},
24-
"claude-3-5-sonnet-20240620": {
25-
"provider": "anthropic",
26-
"embedding_dimensions": None,
27-
},
28-
"claude-3-5-sonnet-20241022": {
29-
"provider": "anthropic",
30-
"embedding_dimensions": None,
31-
},
32-
"claude-3-5-haiku-20241022": {
33-
"provider": "anthropic",
34-
"embedding_dimensions": None,
35-
},
36-
"claude-3-7-sonnet-20250219": {
37-
"provider": "anthropic",
38-
"embedding_dimensions": None,
39-
},
40-
"claude-3-7-sonnet-latest": {"provider": "anthropic", "embedding_dimensions": None},
41-
"claude-3-5-sonnet-latest": {"provider": "anthropic", "embedding_dimensions": None},
42-
"claude-3-5-haiku-latest": {"provider": "anthropic", "embedding_dimensions": None},
43-
"claude-3-opus-latest": {"provider": "anthropic", "embedding_dimensions": None},
44-
"o1": {"provider": "openai", "embedding_dimensions": None},
45-
"o1-mini": {"provider": "openai", "embedding_dimensions": None},
46-
"o3-mini": {"provider": "openai", "embedding_dimensions": None},
32+
# OpenAI Models
33+
"gpt-3.5-turbo": ModelConfig(
34+
provider=ModelProvider.OPENAI,
35+
name="gpt-3.5-turbo",
36+
max_tokens=4096,
37+
embedding_dimensions=1536,
38+
),
39+
"gpt-3.5-turbo-16k": ModelConfig(
40+
provider=ModelProvider.OPENAI,
41+
name="gpt-3.5-turbo-16k",
42+
max_tokens=16384,
43+
embedding_dimensions=1536,
44+
),
45+
"gpt-4": ModelConfig(
46+
provider=ModelProvider.OPENAI,
47+
name="gpt-4",
48+
max_tokens=8192,
49+
embedding_dimensions=1536,
50+
),
51+
"gpt-4-32k": ModelConfig(
52+
provider=ModelProvider.OPENAI,
53+
name="gpt-4-32k",
54+
max_tokens=32768,
55+
embedding_dimensions=1536,
56+
),
57+
"gpt-4o": ModelConfig(
58+
provider=ModelProvider.OPENAI,
59+
name="gpt-4o",
60+
max_tokens=128000,
61+
embedding_dimensions=1536,
62+
),
63+
"gpt-4o-mini": ModelConfig(
64+
provider=ModelProvider.OPENAI,
65+
name="gpt-4o-mini",
66+
max_tokens=128000,
67+
embedding_dimensions=1536,
68+
),
69+
# Newer reasoning models
70+
"o1": ModelConfig(
71+
provider=ModelProvider.OPENAI,
72+
name="o1",
73+
max_tokens=200000,
74+
embedding_dimensions=1536,
75+
),
76+
"o1-mini": ModelConfig(
77+
provider=ModelProvider.OPENAI,
78+
name="o1-mini",
79+
max_tokens=128000,
80+
embedding_dimensions=1536,
81+
),
82+
"o3-mini": ModelConfig(
83+
provider=ModelProvider.OPENAI,
84+
name="o3-mini",
85+
max_tokens=200000,
86+
embedding_dimensions=1536,
87+
),
88+
# Embedding models
89+
"text-embedding-ada-002": ModelConfig(
90+
provider=ModelProvider.OPENAI,
91+
name="text-embedding-ada-002",
92+
max_tokens=8191,
93+
embedding_dimensions=1536,
94+
),
95+
"text-embedding-3-small": ModelConfig(
96+
provider=ModelProvider.OPENAI,
97+
name="text-embedding-3-small",
98+
max_tokens=8191,
99+
embedding_dimensions=1536,
100+
),
101+
"text-embedding-3-large": ModelConfig(
102+
provider=ModelProvider.OPENAI,
103+
name="text-embedding-3-large",
104+
max_tokens=8191,
105+
embedding_dimensions=3072,
106+
),
107+
# Anthropic Models
108+
"claude-3-opus-20240229": ModelConfig(
109+
provider=ModelProvider.ANTHROPIC,
110+
name="claude-3-opus-20240229",
111+
max_tokens=200000,
112+
embedding_dimensions=1536,
113+
),
114+
"claude-3-sonnet-20240229": ModelConfig(
115+
provider=ModelProvider.ANTHROPIC,
116+
name="claude-3-sonnet-20240229",
117+
max_tokens=200000,
118+
embedding_dimensions=1536,
119+
),
120+
"claude-3-haiku-20240307": ModelConfig(
121+
provider=ModelProvider.ANTHROPIC,
122+
name="claude-3-haiku-20240307",
123+
max_tokens=200000,
124+
embedding_dimensions=1536,
125+
),
126+
"claude-3-5-sonnet-20240620": ModelConfig(
127+
provider=ModelProvider.ANTHROPIC,
128+
name="claude-3-5-sonnet-20240620",
129+
max_tokens=200000,
130+
embedding_dimensions=1536,
131+
),
132+
# Latest Anthropic Models
133+
"claude-3-7-sonnet-20250219": ModelConfig(
134+
provider=ModelProvider.ANTHROPIC,
135+
name="claude-3-7-sonnet-20250219",
136+
max_tokens=200000,
137+
embedding_dimensions=1536,
138+
),
139+
"claude-3-5-sonnet-20241022": ModelConfig(
140+
provider=ModelProvider.ANTHROPIC,
141+
name="claude-3-5-sonnet-20241022",
142+
max_tokens=200000,
143+
embedding_dimensions=1536,
144+
),
145+
"claude-3-5-haiku-20241022": ModelConfig(
146+
provider=ModelProvider.ANTHROPIC,
147+
name="claude-3-5-haiku-20241022",
148+
max_tokens=200000,
149+
embedding_dimensions=1536,
150+
),
151+
# Convenience aliases
152+
"claude-3-7-sonnet-latest": ModelConfig(
153+
provider=ModelProvider.ANTHROPIC,
154+
name="claude-3-7-sonnet-20250219",
155+
max_tokens=200000,
156+
embedding_dimensions=1536,
157+
),
158+
"claude-3-5-sonnet-latest": ModelConfig(
159+
provider=ModelProvider.ANTHROPIC,
160+
name="claude-3-5-sonnet-20241022",
161+
max_tokens=200000,
162+
embedding_dimensions=1536,
163+
),
164+
"claude-3-5-haiku-latest": ModelConfig(
165+
provider=ModelProvider.ANTHROPIC,
166+
name="claude-3-5-haiku-20241022",
167+
max_tokens=200000,
168+
embedding_dimensions=1536,
169+
),
170+
"claude-3-opus-latest": ModelConfig(
171+
provider=ModelProvider.ANTHROPIC,
172+
name="claude-3-opus-20240229",
173+
max_tokens=200000,
174+
embedding_dimensions=1536,
175+
),
47176
}
48177

49178

@@ -167,14 +296,14 @@ class Config:
167296
extra = "ignore" # Ignore extra environment variables
168297

169298
@property
170-
def generation_model_config(self) -> dict[str, Any]:
299+
def generation_model_config(self) -> ModelConfig | None:
171300
"""Get configuration for the generation model."""
172-
return MODEL_CONFIGS.get(self.generation_model, {})
301+
return MODEL_CONFIGS.get(self.generation_model)
173302

174303
@property
175-
def embedding_model_config(self) -> dict[str, Any]:
304+
def embedding_model_config(self) -> ModelConfig | None:
176305
"""Get configuration for the embedding model."""
177-
return MODEL_CONFIGS.get(self.embedding_model, {})
306+
return MODEL_CONFIGS.get(self.embedding_model)
178307

179308
def load_yaml_config(self, config_path: str) -> dict[str, Any]:
180309
"""Load configuration from YAML file."""

0 commit comments

Comments
 (0)