Upsert fix for kafka

valhalla/jawn/src/lib/db/database.types.ts

@@ -1936,6 +1936,7 @@ export type Database = {
           created_at: string
           delay_ms: number | null
           feedback: Json | null
+          helicone_org_id: string | null
           id: string
           model: string | null
           prompt_tokens: number | null
@@ -1949,6 +1950,7 @@ export type Database = {
           created_at?: string
           delay_ms?: number | null
           feedback?: Json | null
+          helicone_org_id?: string | null
           id?: string
           model?: string | null
           prompt_tokens?: number | null
@@ -1962,6 +1964,7 @@ export type Database = {
           created_at?: string
           delay_ms?: number | null
           feedback?: Json | null
+          helicone_org_id?: string | null
           id?: string
           model?: string | null
           prompt_tokens?: number | null
@@ -2515,101 +2518,6 @@ export type Database = {
           },
         ]
       }
-      s3_multipart_uploads: {
-        Row: {
-          bucket_id: string
-          created_at: string
-          id: string
-          in_progress_size: number
-          key: string
-          owner_id: string | null
-          upload_signature: string
-          version: string
-        }
-        Insert: {
-          bucket_id: string
-          created_at?: string
-          id: string
-          in_progress_size?: number
-          key: string
-          owner_id?: string | null
-          upload_signature: string
-          version: string
-        }
-        Update: {
-          bucket_id?: string
-          created_at?: string
-          id?: string
-          in_progress_size?: number
-          key?: string
-          owner_id?: string | null
-          upload_signature?: string
-          version?: string
-        }
-        Relationships: [
-          {
-            foreignKeyName: "s3_multipart_uploads_bucket_id_fkey"
-            columns: ["bucket_id"]
-            isOneToOne: false
-            referencedRelation: "buckets"
-            referencedColumns: ["id"]
-          },
-        ]
-      }
-      s3_multipart_uploads_parts: {
-        Row: {
-          bucket_id: string
-          created_at: string
-          etag: string
-          id: string
-          key: string
-          owner_id: string | null
-          part_number: number
-          size: number
-          upload_id: string
-          version: string
-        }
-        Insert: {
-          bucket_id: string
-          created_at?: string
-          etag: string
-          id?: string
-          key: string
-          owner_id?: string | null
-          part_number: number
-          size?: number
-          upload_id: string
-          version: string
-        }
-        Update: {
-          bucket_id?: string
-          created_at?: string
-          etag?: string
-          id?: string
-          key?: string
-          owner_id?: string | null
-          part_number?: number
-          size?: number
-          upload_id?: string
-          version?: string
-        }
-        Relationships: [
-          {
-            foreignKeyName: "s3_multipart_uploads_parts_bucket_id_fkey"
-            columns: ["bucket_id"]
-            isOneToOne: false
-            referencedRelation: "buckets"
-            referencedColumns: ["id"]
-          },
-          {
-            foreignKeyName: "s3_multipart_uploads_parts_upload_id_fkey"
-            columns: ["upload_id"]
-            isOneToOne: false
-            referencedRelation: "s3_multipart_uploads"
-            referencedColumns: ["id"]
-          },
-        ]
-      }
     }
     Views: {
       [_ in never]: never
@@ -2649,37 +2557,6 @@ export type Database = {
           bucket_id: string
         }[]
       }
-      list_multipart_uploads_with_delimiter: {
-        Args: {
-          bucket_id: string
-          prefix_param: string
-          delimiter_param: string
-          max_keys?: number
-          next_key_token?: string
-          next_upload_token?: string
-        }
-        Returns: {
-          key: string
-          id: string
-          created_at: string
-        }[]
-      }
-      list_objects_with_delimiter: {
-        Args: {
-          bucket_id: string
-          prefix_param: string
-          delimiter_param: string
-          max_keys?: number
-          start_after?: string
-          next_token?: string
-        }
-        Returns: {
-          name: string
-          id: string
-          metadata: Json
-          updated_at: string
-        }[]
-      }
       search: {
         Args: {
           prefix: string

valhalla/jawn/src/lib/handlers/LoggingHandler.ts

@@ -366,10 +366,12 @@ export class LoggingHandler extends AbstractLogHandler {
   ): Database["public"]["Tables"]["response"]["Insert"] {
     const response = context.message.log.response;
     const processedResponse = context.processedLog.response;
+    const orgParams = context.orgParams;
 
     const responseInsert: Database["public"]["Tables"]["response"]["Insert"] = {
       id: response.id,
       request: context.message.log.request.id,
+      helicone_org_id: orgParams?.id ?? null,
       body: "{}",
       status: response.status,
       model: processedResponse.model,

valhalla/jawn/src/lib/stores/LogStore.ts

@@ -47,7 +47,7 @@ const requestColumns = new pgp.helpers.ColumnSet(
   { table: "request" }
 );
 const onConflictRequest =
-  " ON CONFLICT (id) DO UPDATE SET " +
+  " ON CONFLICT (id, helicone_org_id) DO UPDATE SET " +
   requestColumns.assignColumns({ from: "EXCLUDED", skip: "id" });
 
 const responseColumns = new pgp.helpers.ColumnSet(
@@ -56,6 +56,7 @@ const responseColumns = new pgp.helpers.ColumnSet(
     "body",
     "request",
     "created_at",
+    { name: "helicone_org_id", def: null },
     { name: "model", def: null },
     { name: "completion_tokens", def: null },
     { name: "delay_ms", def: null },
@@ -68,7 +69,7 @@ const responseColumns = new pgp.helpers.ColumnSet(
 );
 
 const onConflictResponse =
-  " ON CONFLICT (id) DO UPDATE SET " +
+  " ON CONFLICT (request, helicone_org_id) DO UPDATE SET " +
   responseColumns.assignColumns({ from: "EXCLUDED", skip: "id" });
 
 const assetColumns = new pgp.helpers.ColumnSet(

worker/package.json

@@ -12,6 +12,7 @@
     "@types/jest": "^29.5.1",
     "@types/node-gzip": "^1.1.3",
     "@types/supertest": "^2.0.12",
+    "@types/uuid": "^9.0.8",
     "@typescript-eslint/eslint-plugin": "^5.59.0",
     "@typescript-eslint/parser": "^5.59.0",
     "axios": "^1.4.0",
@@ -47,6 +48,7 @@
     "aws4fetch": "^1.0.18",
     "dateformat": "^5.0.3",
     "events": "^3.3.0",
-    "itty-router": "^4.0.13"
+    "itty-router": "^4.0.13",
+    "uuid": "^9.0.1"
   }
 }

worker/src/lib/models/HeliconeHeaders.ts

@@ -1,4 +1,5 @@
 import { Headers } from "@cloudflare/workers-types";
+import { validate as uuidValidate } from "uuid";
 
 type Nullable<T> = T | null;
 
@@ -202,6 +203,7 @@ export class HeliconeHeaders implements IHeliconeHeaders {
   }
 
   private getHeliconeHeaders(): IHeliconeHeaders {
+    const heliconeRequestId = this.headers.get("Helicone-Request-Id");
     return {
       heliconeAuth: this.headers.get("helicone-auth") ?? null,
       heliconeAuthV2: this.getHeliconeAuthV2(),
@@ -211,7 +213,9 @@ export class HeliconeHeaders implements IHeliconeHeaders {
       targetBaseUrl: this.headers.get("Helicone-Target-URL") ?? null,
       retryHeaders: this.getRetryHeaders(),
       promptFormat: this.headers.get("Helicone-Prompt-Format") ?? null,
-      requestId: this.headers.get("Helicone-Request-Id") ?? crypto.randomUUID(),
+      requestId: uuidValidate(heliconeRequestId ?? "")
+        ? heliconeRequestId
+        : crypto.randomUUID(),
       promptHeaders: {
         promptId: this.headers.get("Helicone-Prompt-Id") ?? null,
         promptMode: this.headers.get("Helicone-Prompt-Mode") ?? null,