Skip to content

Commit 66dba04

Browse files
noahpodgurskinoahpodgurski
authored
feat: search-dev budget (#186)
* feat: search-dev budget * move allowed service types to config * fix tests
1 parent 1b44298 commit 66dba04

6 files changed

Lines changed: 65 additions & 44 deletions

File tree

‎docs/index.html‎

Lines changed: 5 additions & 5 deletions
Large diffs are not rendered by default.

‎src/mlpa/core/config.py‎

Lines changed: 42 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,9 @@ class Env(BaseSettings):
3636
LITELLM_API_BASE: str = "http://localhost:4000"
3737
CHALLENGE_EXPIRY_SECONDS: int = 300 # 5 minutes
3838

39+
# Search allowed service types
40+
SEARCH_ALLOWED_SERVICE_TYPES: set[str] = {"search", "search-dev"}
41+
3942
# User Feature Budget - AI service type
4043
USER_FEATURE_BUDGET_AI_BUDGET_ID: str = "end-user-budget-ai"
4144
USER_FEATURE_BUDGET_AI_MAX_BUDGET: float = 0.1
@@ -64,6 +67,18 @@ class Env(BaseSettings):
6467
USER_FEATURE_BUDGET_MEMORIES_TPM_LIMIT: int = 2000
6568
USER_FEATURE_BUDGET_MEMORIES_BUDGET_DURATION: str = "1d"
6669

70+
USER_FEATURE_BUDGET_SEARCH_BUDGET_ID: str = "end-user-budget-search"
71+
USER_FEATURE_BUDGET_SEARCH_MAX_BUDGET: float = 0.1
72+
USER_FEATURE_BUDGET_SEARCH_RPM_LIMIT: int = 10
73+
USER_FEATURE_BUDGET_SEARCH_TPM_LIMIT: int = 2000
74+
USER_FEATURE_BUDGET_SEARCH_BUDGET_DURATION: str = "1d"
75+
76+
USER_FEATURE_BUDGET_TELEMETRY_BUDGET_ID: str = "end-user-budget-telemetry"
77+
USER_FEATURE_BUDGET_TELEMETRY_MAX_BUDGET: float = 0.1
78+
USER_FEATURE_BUDGET_TELEMETRY_RPM_LIMIT: int = 10
79+
USER_FEATURE_BUDGET_TELEMETRY_TPM_LIMIT: int = 2000
80+
USER_FEATURE_BUDGET_TELEMETRY_BUDGET_DURATION: str = "1d"
81+
6782
# User Feature Budget - ai-dev service type (experimentation, batch predictions)
6883
USER_FEATURE_BUDGET_AI_DEV_BUDGET_ID: str = "end-user-budget-ai-dev"
6984
USER_FEATURE_BUDGET_AI_DEV_MAX_BUDGET: float = 1.0
@@ -84,17 +99,11 @@ class Env(BaseSettings):
8499
USER_FEATURE_BUDGET_MOCHI_DEV_TPM_LIMIT: int = 10000
85100
USER_FEATURE_BUDGET_MOCHI_DEV_BUDGET_DURATION: str = "1d"
86101

87-
USER_FEATURE_BUDGET_SEARCH_BUDGET_ID: str = "end-user-budget-search"
88-
USER_FEATURE_BUDGET_SEARCH_MAX_BUDGET: float = 0.1
89-
USER_FEATURE_BUDGET_SEARCH_RPM_LIMIT: int = 10
90-
USER_FEATURE_BUDGET_SEARCH_TPM_LIMIT: int = 2000
91-
USER_FEATURE_BUDGET_SEARCH_BUDGET_DURATION: str = "1d"
92-
93-
USER_FEATURE_BUDGET_TELEMETRY_BUDGET_ID: str = "end-user-budget-telemetry"
94-
USER_FEATURE_BUDGET_TELEMETRY_MAX_BUDGET: float = 0.1
95-
USER_FEATURE_BUDGET_TELEMETRY_RPM_LIMIT: int = 10
96-
USER_FEATURE_BUDGET_TELEMETRY_TPM_LIMIT: int = 2000
97-
USER_FEATURE_BUDGET_TELEMETRY_BUDGET_DURATION: str = "1d"
102+
USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_ID: str = "end-user-budget-search-dev"
103+
USER_FEATURE_BUDGET_SEARCH_DEV_MAX_BUDGET: float = 1.0
104+
USER_FEATURE_BUDGET_SEARCH_DEV_RPM_LIMIT: int = 200
105+
USER_FEATURE_BUDGET_SEARCH_DEV_TPM_LIMIT: int = 10000
106+
USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_DURATION: str = "1d"
98107

99108
@cached_property
100109
def user_feature_budget(self) -> dict[str, dict]:
@@ -132,6 +141,20 @@ def user_feature_budget(self) -> dict[str, dict]:
132141
"tpm_limit": self.USER_FEATURE_BUDGET_MEMORIES_TPM_LIMIT,
133142
"budget_duration": self.USER_FEATURE_BUDGET_MEMORIES_BUDGET_DURATION,
134143
},
144+
"search": {
145+
"budget_id": self.USER_FEATURE_BUDGET_SEARCH_BUDGET_ID,
146+
"max_budget": self.USER_FEATURE_BUDGET_SEARCH_MAX_BUDGET,
147+
"rpm_limit": self.USER_FEATURE_BUDGET_SEARCH_RPM_LIMIT,
148+
"tpm_limit": self.USER_FEATURE_BUDGET_SEARCH_TPM_LIMIT,
149+
"budget_duration": self.USER_FEATURE_BUDGET_SEARCH_BUDGET_DURATION,
150+
},
151+
"telemetry": {
152+
"budget_id": self.USER_FEATURE_BUDGET_TELEMETRY_BUDGET_ID,
153+
"max_budget": self.USER_FEATURE_BUDGET_TELEMETRY_MAX_BUDGET,
154+
"rpm_limit": self.USER_FEATURE_BUDGET_TELEMETRY_RPM_LIMIT,
155+
"tpm_limit": self.USER_FEATURE_BUDGET_TELEMETRY_TPM_LIMIT,
156+
"budget_duration": self.USER_FEATURE_BUDGET_TELEMETRY_BUDGET_DURATION,
157+
},
135158
"ai-dev": {
136159
"budget_id": self.USER_FEATURE_BUDGET_AI_DEV_BUDGET_ID,
137160
"max_budget": self.USER_FEATURE_BUDGET_AI_DEV_MAX_BUDGET,
@@ -153,19 +176,12 @@ def user_feature_budget(self) -> dict[str, dict]:
153176
"tpm_limit": self.USER_FEATURE_BUDGET_MOCHI_DEV_TPM_LIMIT,
154177
"budget_duration": self.USER_FEATURE_BUDGET_MOCHI_DEV_BUDGET_DURATION,
155178
},
156-
"search": {
157-
"budget_id": self.USER_FEATURE_BUDGET_SEARCH_BUDGET_ID,
158-
"max_budget": self.USER_FEATURE_BUDGET_SEARCH_MAX_BUDGET,
159-
"rpm_limit": self.USER_FEATURE_BUDGET_SEARCH_RPM_LIMIT,
160-
"tpm_limit": self.USER_FEATURE_BUDGET_SEARCH_TPM_LIMIT,
161-
"budget_duration": self.USER_FEATURE_BUDGET_SEARCH_BUDGET_DURATION,
162-
},
163-
"telemetry": {
164-
"budget_id": self.USER_FEATURE_BUDGET_TELEMETRY_BUDGET_ID,
165-
"max_budget": self.USER_FEATURE_BUDGET_TELEMETRY_MAX_BUDGET,
166-
"rpm_limit": self.USER_FEATURE_BUDGET_TELEMETRY_RPM_LIMIT,
167-
"tpm_limit": self.USER_FEATURE_BUDGET_TELEMETRY_TPM_LIMIT,
168-
"budget_duration": self.USER_FEATURE_BUDGET_TELEMETRY_BUDGET_DURATION,
179+
"search-dev": {
180+
"budget_id": self.USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_ID,
181+
"max_budget": self.USER_FEATURE_BUDGET_SEARCH_DEV_MAX_BUDGET,
182+
"rpm_limit": self.USER_FEATURE_BUDGET_SEARCH_DEV_RPM_LIMIT,
183+
"tpm_limit": self.USER_FEATURE_BUDGET_SEARCH_DEV_TPM_LIMIT,
184+
"budget_duration": self.USER_FEATURE_BUDGET_SEARCH_DEV_BUDGET_DURATION,
169185
},
170186
}
171187

@@ -197,6 +213,7 @@ def service_type_purposes(self) -> dict[str, list[str]]:
197213
"s2s": [],
198214
"s2s-android": [],
199215
"search": [],
216+
"search-dev": [],
200217
"telemetry": telemetry_purposes,
201218
}
202219

@@ -214,7 +231,7 @@ def forced_model_service_type_pairs(self) -> dict[str, list[str]]:
214231
Returns a dictionary mapping model names to their valid service types.
215232
"""
216233
# Force certain models to use certain service types
217-
return {"exa": ["search"]}
234+
return {"exa": ["search", "search-dev"]}
218235

219236
def valid_service_type_for_model(self, service_type: str, model: str) -> bool:
220237
"""Check if a service type is valid for a specific model."""

‎src/mlpa/core/pg_services/litellm_pg_service.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,6 @@ async def create_budget(self):
182182
If a budget already exists, it will be overwritten with the new values.
183183
Returns a list of created/updated budget records.
184184
"""
185-
budgets_created = []
186185
user_feature_budgets = env.user_feature_budget
187186

188187
for service_type, budget_config in user_feature_budgets.items():

‎src/mlpa/run.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,7 @@ async def search(
233233
AuthorizedSearchRequest, Depends(authorize_search_request)
234234
],
235235
):
236-
if authorized_search_request.service_type != "search":
236+
if authorized_search_request.service_type not in env.SEARCH_ALLOWED_SERVICE_TYPES:
237237
raise HTTPException(
238238
status_code=400, detail="service-type header must be of type 'search'"
239239
)

‎src/tests/unit/test_authorize.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,5 +86,5 @@ async def test_authorize_chat_request_rejects_invalid_service_type_for_model():
8686
assert exc_info.value.status_code == 400
8787
assert (
8888
exc_info.value.detail
89-
== "Invalid service-type value for model exa. Should be one of ['search']"
89+
== "Invalid service-type value for model exa. Should be one of ['search', 'search-dev']"
9090
)

‎src/tests/unit/test_config.py‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -98,12 +98,13 @@ def test_valid_service_types_all_service_types():
9898
assert "s2s" in service_types
9999
assert "s2s-android" in service_types
100100
assert "memories" in service_types
101+
assert "search" in service_types
102+
assert "telemetry" in service_types
101103
assert "ai-dev" in service_types
102104
assert "memories-dev" in service_types
103105
assert "mochi-dev" in service_types
104-
assert "search" in service_types
105-
assert "telemetry" in service_types
106-
assert len(service_types) == 9
106+
assert "search-dev" in service_types
107+
assert len(service_types) == 10
107108

108109

109110
def test_service_type_purposes():
@@ -112,11 +113,11 @@ def test_service_type_purposes():
112113
purposes = env.service_type_purposes
113114
ai_purposes = ["chat", "title-generation", "convo-starters-sidebar"]
114115
assert purposes["ai"] == ai_purposes
116+
assert purposes["memories"] == ["memory-generation"]
117+
assert purposes["telemetry"] == ["chat"]
115118
assert purposes["ai-dev"] == ai_purposes
116119
assert purposes["mochi-dev"] == ai_purposes
117-
assert purposes["memories"] == ["memory-generation"]
118120
assert purposes["memories-dev"] == ["memory-generation"]
119-
assert purposes["telemetry"] == ["chat"]
120121

121122

122123
def test_service_type_purposes_s2s_empty():
@@ -126,20 +127,22 @@ def test_service_type_purposes_s2s_empty():
126127
assert purposes["s2s"] == []
127128
assert purposes["s2s-android"] == []
128129
assert purposes["search"] == []
130+
assert purposes["search-dev"] == []
129131

130132

131133
def test_service_type_requires_purpose():
132134
"""Test that purpose is required only for service types with configured allowlists."""
133135
env = Env()
134136
assert env.service_type_requires_purpose("ai") is True
135-
assert env.service_type_requires_purpose("ai-dev") is True
136-
assert env.service_type_requires_purpose("mochi-dev") is True
137137
assert env.service_type_requires_purpose("memories") is True
138-
assert env.service_type_requires_purpose("memories-dev") is True
139138
assert env.service_type_requires_purpose("s2s") is False
140139
assert env.service_type_requires_purpose("s2s-android") is False
141140
assert env.service_type_requires_purpose("search") is False
142141
assert env.service_type_requires_purpose("telemetry") is True
142+
assert env.service_type_requires_purpose("ai-dev") is True
143+
assert env.service_type_requires_purpose("mochi-dev") is True
144+
assert env.service_type_requires_purpose("memories-dev") is True
145+
assert env.service_type_requires_purpose("search-dev") is False
143146

144147

145148
def test_valid_purposes_for_service_type():
@@ -170,11 +173,12 @@ def test_user_feature_budget_structure_consistency():
170173
"s2s",
171174
"s2s-android",
172175
"memories",
176+
"search",
177+
"telemetry",
173178
"ai-dev",
174179
"memories-dev",
175180
"mochi-dev",
176-
"search",
177-
"telemetry",
181+
"search-dev",
178182
]:
179183
assert service_type in budgets
180184
service_keys = set(budgets[service_type].keys())
@@ -199,14 +203,15 @@ def test_forced_model_service_type_pairs_defaults():
199203
"""Test that forced model/service-type mappings include search-only models."""
200204
env = Env()
201205

202-
assert env.forced_model_service_type_pairs == {"exa": ["search"]}
206+
assert env.forced_model_service_type_pairs == {"exa": ["search", "search-dev"]}
203207

204208

205209
def test_valid_service_type_for_model_forced_pair():
206210
"""Test that forced model/service-type pairs are enforced."""
207211
env = Env()
208212

209213
assert env.valid_service_type_for_model("search", "exa") is True
214+
assert env.valid_service_type_for_model("search-dev", "exa") is True
210215
assert env.valid_service_type_for_model("ai", "exa") is False
211216

212217

0 commit comments

Comments
 (0)