From c5e5fe17bc98aade40cc0e2e4805b1bd2f5d3ec6 Mon Sep 17 00:00:00 2001
From: colegottdank <colegottdank@gmail.com>
Date: Thu, 19 Sep 2024 19:23:18 -0700
Subject: [PATCH] fix async prompt logging (#2673)

---
 valhalla/jawn/package.json                 |  1 +
 worker/src/lib/dbLogger/DBLoggable.ts      | 16 +++++++++++++++-
 worker/src/lib/managers/AsyncLogManager.ts | 13 +++++++++++++
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/valhalla/jawn/package.json b/valhalla/jawn/package.json
index 8deafad516..0da32ceec0 100644
--- a/valhalla/jawn/package.json
+++ b/valhalla/jawn/package.json
@@ -43,6 +43,7 @@
     "pg-promise": "^11.6.0",
     "posthog-node": "^4.0.1",
     "rate-limit-redis": "^4.2.0",
+    "stripe": "^16.12.0",
     "swagger-ui-express": "^5.0.0",
     "tiktoken": "^1.0.15",
     "typescript": "^5.2.2",
diff --git a/worker/src/lib/dbLogger/DBLoggable.ts b/worker/src/lib/dbLogger/DBLoggable.ts
index 1c0ac874aa..0733045a5d 100644
--- a/worker/src/lib/dbLogger/DBLoggable.ts
+++ b/worker/src/lib/dbLogger/DBLoggable.ts
@@ -116,6 +116,7 @@ interface DBLoggableRequestFromAsyncLogModelProps {
   providerRequestHeaders: HeliconeHeaders;
   providerResponseHeaders: Headers;
   provider: Provider;
+  heliconeTemplate?: TemplateWithInputs;
 }
 
 function getResponseBodyFromJSON(json: Record<string, Json>): {
@@ -145,12 +146,24 @@ export async function dbLoggableRequestFromAsyncLogModel(
     providerRequestHeaders,
     providerResponseHeaders,
     provider,
+    heliconeTemplate,
   } = props;
 
   return new DBLoggable({
     request: {
       requestId: providerRequestHeaders.requestId ?? crypto.randomUUID(),
-      promptSettings: requestWrapper.promptSettings,
+      promptSettings: providerRequestHeaders.promptHeaders?.promptId
+        ? {
+            promptId: providerRequestHeaders.promptHeaders.promptId,
+            promptVersion:
+              providerRequestHeaders.promptHeaders.promptVersion ?? "",
+            promptMode: "production",
+          }
+        : {
+            promptId: undefined,
+            promptVersion: "",
+            promptMode: "deactivated",
+          },
       userId: providerRequestHeaders.userId ?? undefined,
       startTime: new Date(
         asyncLogModel.timing.startTime.seconds * 1000 +
@@ -169,6 +182,7 @@ export async function dbLoggableRequestFromAsyncLogModel(
       flaggedForModeration: null,
       request_ip: null,
       country_code: (requestWrapper.cf?.country as string) ?? null,
+      heliconeTemplate: heliconeTemplate ?? undefined,
     },
     response: {
       responseId: crypto.randomUUID(),
diff --git a/worker/src/lib/managers/AsyncLogManager.ts b/worker/src/lib/managers/AsyncLogManager.ts
index 73b2a221a1..105acf9ef7 100644
--- a/worker/src/lib/managers/AsyncLogManager.ts
+++ b/worker/src/lib/managers/AsyncLogManager.ts
@@ -12,6 +12,8 @@ import { DBQueryTimer } from "../util/loggers/DBQueryTimer";
 import { S3Client } from "../clients/S3Client";
 import { RequestResponseManager } from "./RequestResponseManager";
 import { KafkaProducer } from "../clients/KafkaProducer";
+import { parseJSXObject } from "@helicone/prompts";
+import { TemplateWithInputs } from "@helicone/prompts/dist/objectParser";
 
 function mergeHeaders(x: Headers, y: Headers) {
   const merged = new Headers();
@@ -51,6 +53,16 @@ export async function logAsync(
     mergeHeaders(requestHeaders, requestWrapper.headers)
   );
 
+  let templateWithInputs: TemplateWithInputs | undefined;
+  try {
+    if (requestWrapper.heliconeHeaders.promptHeaders.promptId) {
+      const parseResult = parseJSXObject(asyncLogModel.providerRequest.json);
+      templateWithInputs = parseResult?.templateWithInputs;
+    }
+  } catch (e) {
+    console.error("Error parsing prompt:", e);
+  }
+
   const loggable = await dbLoggableRequestFromAsyncLogModel({
     requestWrapper,
     env,
@@ -58,6 +70,7 @@ export async function logAsync(
     providerRequestHeaders: heliconeHeaders,
     providerResponseHeaders: responseHeaders,
     provider: provider,
+    heliconeTemplate: templateWithInputs,
   });
 
   const { data: auth, error: authError } = await requestWrapper.auth();