Skip to content

Commit 183377c

Browse files
chitalianJustin Torregreptile-apps[bot]
authored
Justin/request body buffer phase 2.2 (#4781)
* initial scaffolding for requestBodyBuffer * init some docker scaffolding * initial scaffolding for requestBodyBuffer * added signing aws requests * added docker containers and requestbodybuffer abstractions * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * rm * renmaed file * moved to new file * comment * clean up * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * types * update wrangler * request body buffer impl * container mvp * bump timeout * remove unecessary auth * types * added formatting * types * Update worker/RequestBodyBufferContainer/README.md Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update worker/src/RequestBodyBuffer/RequestBodyBuffer_Remote.ts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * cleanup * revert * renamed everything to unsafe * fix build * fix tests --------- Co-authored-by: Justin Torre <justin@helicone.ai> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
1 parent fbc3c9b commit 183377c

File tree

11 files changed

+74
-64
lines changed

11 files changed

+74
-64
lines changed

‎worker/src/lib/RequestWrapper.ts‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -309,7 +309,7 @@ export class RequestWrapper {
309309
}
310310

311311
// TODO deprecate this function
312-
async getRawText(): Promise<string> {
312+
async unsafeGetRawText(): Promise<string> {
313313
return this.requestBodyBuffer.unsafeGetRawText();
314314
}
315315

@@ -332,8 +332,8 @@ export class RequestWrapper {
332332
return !!hostParts.includes("eu") || !!auth?.includes("helicone-eu");
333333
}
334334

335-
async getText(): Promise<string> {
336-
let text = await this.getRawText();
335+
async unsafeGetText(): Promise<string> {
336+
let text = await this.unsafeGetRawText();
337337

338338
if (this.bodyKeyOverride) {
339339
try {
@@ -354,11 +354,11 @@ export class RequestWrapper {
354354
return text;
355355
}
356356

357-
async getJson<T>(): Promise<T> {
357+
async unsafeGetJson<T>(): Promise<T> {
358358
try {
359-
return JSON.parse(await this.getText());
359+
return JSON.parse(await this.unsafeGetText());
360360
} catch (e) {
361-
console.error("RequestWrapper.getJson", e, await this.getText());
361+
console.error("RequestWrapper.getJson", e, await this.unsafeGetText());
362362
return {} as T;
363363
}
364364
}
@@ -468,7 +468,7 @@ export class RequestWrapper {
468468
async getUserId(): Promise<string | undefined> {
469469
const userId =
470470
this.heliconeHeaders.userId ||
471-
(await this.getJson<{ user?: string }>()).user;
471+
(await this.unsafeGetJson<{ user?: string }>()).user;
472472
return userId;
473473
}
474474

‎worker/src/lib/ai-gateway/SimpleAIGateway.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ export class SimpleAIGateway {
165165
Result<{ modelStrings: string[]; body: any }, Response>
166166
> {
167167
// Get raw text body once
168-
const rawBody = await this.requestWrapper.getText();
168+
const rawBody = await this.requestWrapper.unsafeGetText();
169169
const parsedBody: any = tryJSONParse(rawBody ?? "{}");
170170

171171
if (!parsedBody || !parsedBody.model) {

‎worker/src/lib/managers/AsyncLogManager.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ export async function logAsync(
3131
ctx: ExecutionContext,
3232
provider: Provider
3333
): Promise<Response> {
34-
const asyncLogModel = await requestWrapper.getJson<AsyncLogModel>();
34+
const asyncLogModel = await requestWrapper.unsafeGetJson<AsyncLogModel>();
3535
// if payload is larger than 10MB, return 400
3636
const MAX_PAYLOAD_SIZE = 10 * 1024 * 1024;
3737
if (JSON.stringify(asyncLogModel).length > MAX_PAYLOAD_SIZE) {

‎worker/src/lib/managers/FeedbackManager.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ interface FeedbackRequestBodyV2 {
1313
}
1414

1515
export async function handleFeedback(request: RequestWrapper) {
16-
const body = await request.getJson<FeedbackRequestBodyV2>();
16+
const body = await request.unsafeGetJson<FeedbackRequestBodyV2>();
1717
const heliconeId = body["helicone-id"];
1818
const rating = body["rating"];
1919

‎worker/src/lib/managers/PropertiesManager.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ export async function handleLoggingEndpoint(
1010
request: RequestWrapper,
1111
env: Env
1212
): Promise<Response> {
13-
const body = await request.getJson<LoggingRequestBody>();
13+
const body = await request.unsafeGetJson<LoggingRequestBody>();
1414
const heliconeId = body["helicone-id"];
1515
const propTag = "helicone-property-";
1616
const heliconeHeaders = Object.fromEntries(

‎worker/src/lib/models/HeliconeProxyRequest.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -196,15 +196,15 @@ export class HeliconeProxyRequestMapper {
196196

197197
if (this.request.heliconeHeaders.featureFlags.streamUsage) {
198198
// eslint-disable-next-line @typescript-eslint/no-explicit-any
199-
const jsonBody = (await this.request.getJson()) as any;
199+
const jsonBody = (await this.request.unsafeGetJson()) as any;
200200
if (!jsonBody["stream_options"]) {
201201
jsonBody["stream_options"] = {};
202202
}
203203
jsonBody["stream_options"]["include_usage"] = true;
204204
return JSON.stringify(jsonBody);
205205
}
206206

207-
return await this.request.getText();
207+
return await this.request.unsafeGetText();
208208
}
209209

210210
private validateApiConfiguration(api_base: string | undefined): boolean {
@@ -264,7 +264,7 @@ export class HeliconeProxyRequestMapper {
264264

265265
async requestJson(): Promise<HeliconeProxyRequest["requestJson"]> {
266266
return this.request.getMethod() === "POST"
267-
? await this.request.getJson()
267+
? await this.request.unsafeGetJson()
268268
: {};
269269
}
270270
}

‎worker/src/lib/util/cache/cacheFunctions.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ export async function kvKeyFromRequest(
5151
request.requestWrapper.heliconeHeaders.cacheHeaders.cacheIgnoreKeys ?? [];
5252

5353
const body = tryGetBodyAndRemoveKeys(
54-
await request.requestWrapper.getText(),
54+
await request.requestWrapper.unsafeGetText(),
5555
ignoreKeys
5656
);
5757

‎worker/src/routers/api/apiRouter.ts‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ function getAPIRouterV1(
3737
return new Response("not allowed", { status: 403 });
3838
}
3939

40-
const data = await requestWrapper.getJson<{
40+
const data = await requestWrapper.unsafeGetJson<{
4141
apiKeyHash: string;
4242
orgId: string;
4343
softDelete?: boolean;
@@ -91,7 +91,7 @@ function getAPIRouterV1(
9191
return new Response("not allowed", { status: 403 });
9292
}
9393

94-
const data = await requestWrapper.getJson<
94+
const data = await requestWrapper.unsafeGetJson<
9595
{
9696
provider: ModelProviderName;
9797
decryptedProviderKey: string;
@@ -150,7 +150,7 @@ function getAPIRouterV1(
150150
if (authParams.error !== null) {
151151
return client.response.unauthorized();
152152
}
153-
const job = await requestWrapper.getJson<Job>();
153+
const job = await requestWrapper.unsafeGetJson<Job>();
154154

155155
if (!job) {
156156
return client.response.newError("Invalid run", 400);
@@ -208,7 +208,7 @@ function getAPIRouterV1(
208208
}
209209

210210
const status =
211-
(await requestWrapper.getJson<{ status: string }>()).status ?? "";
211+
(await requestWrapper.unsafeGetJson<{ status: string }>()).status ?? "";
212212

213213
if (!isValidStatus(status)) {
214214
return client.response.newError("Invalid status", 400);
@@ -237,7 +237,7 @@ function getAPIRouterV1(
237237
return client.response.unauthorized();
238238
}
239239

240-
const node = await requestWrapper.getJson<HeliconeNode>();
240+
const node = await requestWrapper.unsafeGetJson<HeliconeNode>();
241241
if (!node) {
242242
return client.response.newError("Invalid task", 400);
243243
}
@@ -297,7 +297,7 @@ function getAPIRouterV1(
297297
}
298298

299299
const status =
300-
(await requestWrapper.getJson<{ status: string }>()).status ?? "";
300+
(await requestWrapper.unsafeGetJson<{ status: string }>()).status ?? "";
301301

302302
if (!isValidStatus(status)) {
303303
return client.response.newError("Invalid status", 400);
@@ -372,7 +372,7 @@ function getAPIRouterV1(
372372
value: string;
373373
}
374374

375-
const newProperty = await requestWrapper.getJson<Body>();
375+
const newProperty = await requestWrapper.unsafeGetJson<Body>();
376376

377377
const auth = await requestWrapper.auth();
378378

@@ -433,7 +433,7 @@ function getAPIRouterV1(
433433
}
434434

435435
const requestData =
436-
await requestWrapper.getJson<
436+
await requestWrapper.unsafeGetJson<
437437
Database["public"]["Tables"]["alert"]["Insert"]
438438
>();
439439

‎worker/src/routers/generateRouter.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,8 @@ const generateHandler = async (
136136
}
137137

138138
// 2. BUILD GENERATE PARAMETERS FROM REQUEST BODY AND VALIDATE
139-
const rawBody = await requestWrapper.getJson<Record<string, unknown>>();
139+
const rawBody =
140+
await requestWrapper.unsafeGetJson<Record<string, unknown>>();
140141
const paramsResult = validateGenerateParams(rawBody);
141142
if (!paramsResult.success) {
142143
return createErrorResponse(

‎worker/test/cache/cacheFunctions.spec.ts‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ describe("cacheFunctions", () => {
2020
mockRequest = {
2121
url: "https://api.openai.com/v1/chat/completions",
2222
requestWrapper: {
23-
getText: vi.fn() as Mock<[], Promise<string>>,
23+
unsafeGetText: vi.fn() as Mock<[], Promise<string>>,
2424
getHeaders: vi.fn(() => headers) as Mock<[], Headers>,
2525
heliconeHeaders: {
2626
cacheHeaders: {
@@ -44,13 +44,13 @@ describe("cacheFunctions", () => {
4444
timestamp: "2024-01-01T00:00:00Z",
4545
});
4646

47-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(requestBody);
47+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(requestBody);
4848

4949
const result = await kvKeyFromRequest(mockRequest, 0, null);
5050

5151
// The hash should include the full request body
5252
expect(result).toContain("hashed_");
53-
expect(mockRequest.requestWrapper.getText).toHaveBeenCalled();
53+
expect(mockRequest.requestWrapper.unsafeGetText).toHaveBeenCalled();
5454
expect(result).toContain(requestBody);
5555
});
5656

@@ -62,7 +62,7 @@ describe("cacheFunctions", () => {
6262
timestamp: "2024-01-01T00:00:00Z",
6363
};
6464

65-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(JSON.stringify(requestBody));
65+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(JSON.stringify(requestBody));
6666
mockRequest.requestWrapper.heliconeHeaders.cacheHeaders.cacheIgnoreKeys = ["request_id", "timestamp"];
6767

6868
const result = await kvKeyFromRequest(mockRequest, 0, null);
@@ -82,7 +82,7 @@ describe("cacheFunctions", () => {
8282
it("should handle non-JSON body gracefully", async () => {
8383
const textBody = "This is plain text, not JSON";
8484

85-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(textBody);
85+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(textBody);
8686
mockRequest.requestWrapper.heliconeHeaders.cacheHeaders.cacheIgnoreKeys = ["some_key"];
8787

8888
const result = await kvKeyFromRequest(mockRequest, 0, null);
@@ -99,7 +99,7 @@ describe("cacheFunctions", () => {
9999
request_id: "unique-123",
100100
});
101101

102-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(requestBody);
102+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(requestBody);
103103
mockRequest.requestWrapper.heliconeHeaders.cacheHeaders.cacheIgnoreKeys = [];
104104

105105
const result = await kvKeyFromRequest(mockRequest, 0, null);
@@ -120,7 +120,7 @@ describe("cacheFunctions", () => {
120120
request_id: "top-level-123",
121121
};
122122

123-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(JSON.stringify(requestBody));
123+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(JSON.stringify(requestBody));
124124
mockRequest.requestWrapper.heliconeHeaders.cacheHeaders.cacheIgnoreKeys = ["request_id"];
125125

126126
const result = await kvKeyFromRequest(mockRequest, 0, null);
@@ -156,10 +156,10 @@ describe("cacheFunctions", () => {
156156

157157
mockRequest.requestWrapper.heliconeHeaders.cacheHeaders.cacheIgnoreKeys = ["request_id"];
158158

159-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(requestBody1);
159+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(requestBody1);
160160
const result1 = await kvKeyFromRequest(mockRequest, 0, null);
161161

162-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(requestBody2);
162+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(requestBody2);
163163
const result2 = await kvKeyFromRequest(mockRequest, 0, null);
164164

165165
// Keys should be different because model is different (not ignored)
@@ -183,10 +183,10 @@ describe("cacheFunctions", () => {
183183

184184
mockRequest.requestWrapper.heliconeHeaders.cacheHeaders.cacheIgnoreKeys = ["request_id", "timestamp"];
185185

186-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(requestBody1);
186+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(requestBody1);
187187
const result1 = await kvKeyFromRequest(mockRequest, 0, null);
188188

189-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(requestBody2);
189+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(requestBody2);
190190
const result2 = await kvKeyFromRequest(mockRequest, 0, null);
191191

192192
// Keys should be the same because only ignored fields are different
@@ -199,7 +199,7 @@ describe("cacheFunctions", () => {
199199
messages: [{ role: "user", content: "Hello" }],
200200
});
201201

202-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(requestBody);
202+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(requestBody);
203203

204204
const resultWithSeed = await kvKeyFromRequest(mockRequest, 0, "test-seed");
205205
const resultWithoutSeed = await kvKeyFromRequest(mockRequest, 0, null);
@@ -215,7 +215,7 @@ describe("cacheFunctions", () => {
215215
messages: [{ role: "user", content: "Hello" }],
216216
});
217217

218-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(requestBody);
218+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(requestBody);
219219

220220
const resultWithIndex0 = await kvKeyFromRequest(mockRequest, 0, null);
221221
const resultWithIndex1 = await kvKeyFromRequest(mockRequest, 1, null);
@@ -243,7 +243,7 @@ describe("cacheFunctions", () => {
243243
"Helicone-Cache-Control": "max-age=3600",
244244
});
245245

246-
(mockRequest.requestWrapper.getText as Mock).mockResolvedValue(requestBody);
246+
(mockRequest.requestWrapper.unsafeGetText as Mock).mockResolvedValue(requestBody);
247247
mockRequest.requestWrapper.headers = googleHeaders;
248248
mockRequest.requestWrapper.getHeaders = vi.fn(() => googleHeaders);
249249

@@ -255,4 +255,4 @@ describe("cacheFunctions", () => {
255255
expect(result).toContain("helicone-cache-control");
256256
});
257257
});
258-
});
258+
});

0 commit comments

Comments
 (0)