Valhalla ingest (Helicone#1092)

docker-compose-local.yml

@@ -37,12 +37,31 @@ services:
     restart: unless-stopped
     ports:
       - 25432:5432/tcp
-    volumes:
-      - ./volumes/postgres:/var/lib/postgresql/data:z
     environment:
       POSTGRES_USER: postgres
       POSTGRES_DB: helicone
       POSTGRES_PASSWORD: password
+    healthcheck:
+      test: pg_isready -U postgres
+      interval: 5s
+      timeout: 5s
+      retries: 10
+
+  helicone-be-db-migration-runner:
+    container_name: helicone-be-db-migration-runner
+    image: flyway/flyway
+    depends_on:
+      helicone-be-db:
+        condition: service_healthy
+    volumes:
+      - ./valhalla/migrations:/flyway/sql:z
+    environment:
+      - FLYWAY_URL=jdbc:postgresql://helicone-be-db:5432/helicone
+      - FLYWAY_USER=postgres
+      - FLYWAY_PASSWORD=password
+      - FLYWAY_SCHEMAS=public
+      - FLYWAY_LOCATIONS=filesystem:/flyway/sql
+    command: migrate
 
   jawn:
     image: jawn-local

run-local.py

@@ -14,7 +14,7 @@ class Command(str, Enum):
     DOCKER_COMPOSE_BUILD_WORKER = "docker-compose -f docker-compose-local.yml build worker-openai-proxy"
     DOCKER_COMPOSE_LOGS = "docker-compose -f docker-compose-local.yml logs -f"
     DOCKER_COMPOSE_VALHALLA = "docker-compose -f docker-compose-local.yml up -d jawn helicone-be-db"
-    DOCKER_COMPOSE_VALHALLA_DB = "docker-compose -f docker-compose-local.yml up -d helicone-be-db"
+    DOCKER_COMPOSE_VALHALLA_DB = "docker-compose -f docker-compose-local.yml up -d helicone-be-db helicone-be-db-migration-runner"
     DOCKER_COMPOSE_WORKER = "docker-compose -f docker-compose-local.yml up -d worker-openai-proxy worker-anthropic-proxy worker-helicone-api"
     DOCKER_COMPOSE_FRONTEND = "docker-compose -f docker-compose-local.yml up -d web"
     DOCKER_COMPOSE_CLICKHOUSE = "docker-compose -f docker-compose-local.yml up -d clickhouse-migration-runner"

valhalla/migrations/V1__inital.sql

@@ -8,12 +8,8 @@ create table
     url_href text not null,
     user_id text null,
     properties jsonb null,
-    formatted_prompt_id uuid null,
-    prompt_values jsonb null,
-    helicone_api_key_id bigint null,
     helicone_org_id uuid null,
     provider text not null default 'OPENAI'::text,
-    helicone_proxy_key_id uuid null,
     constraint request_pkey primary key (id)
   ) tablespace pg_default;
 
@@ -26,7 +22,7 @@ create table
     body jsonb not null,
     request uuid not null,
     delay_ms integer null,
-    status smallint null,
+    http_status smallint null,
     completion_tokens integer null,
     model text null,
     prompt_tokens integer null,

valhalla/nodemon.jawn.json

@@ -1,7 +1,16 @@
 {
-  "watch": ["packages/jawn/src/**/*", "packages/helicone-shared-ts/src/**/*"],
+  "watch": [
+    "packages/jawn/src/**/*",
+    "packages/helicone-shared-ts/src/**/*",
+    "packages/jawn/src/schema/openapi.yml"
+  ],
+
   "ext": "ts,tsx,js,json",
-  "exec": "yarn workspace helicone-shared-ts build && ts-node packages/jawn/src/index.ts",
-  "ignore": ["node_modules/*", "packages/*/node_modules/*"],
+  "exec": "openapi-typescript packages/jawn/src/schema/openapi.yml --output packages/jawn/src/schema/types.d.ts && yarn workspace helicone-shared-ts build && ts-node packages/jawn/src/index.ts ",
+  "ignore": [
+    "node_modules/*",
+    "packages/*/node_modules/*",
+    "packages/jawn/src/schema/types.d.ts"
+  ],
   "delay": "1000"
 }

valhalla/package.json

@@ -4,7 +4,9 @@
   "license": "MIT",
   "scripts": {
     "dev:jawn": "nodemon --config nodemon.jawn.json",
-    "dev:slay": "nodemon --config nodemon.slay.json"
+    "test:jawn": "npx jest --config packages/jawn/jest.config.js --detectOpenHandles",
+    "dev:slay": "nodemon --config nodemon.slay.json",
+    "test": "yarn test:jawn"
   },
   "workspaces": [
     "packages/*"

valhalla/packages/helicone-shared-ts/src/db/valhalla.database.types.ts

@@ -0,0 +1,36 @@
+export interface ValhallaRequest {
+  id: string; // uuid
+  createdAt: Date; // timestamp with time zone
+  body: any; // jsonb
+  urlHref: string; // text
+  userId: string | null; // text, nullable
+  properties: any | null; // jsonb, nullable
+  heliconeApiKeyID: bigint | null; // bigint, nullable
+  heliconeOrgID: string | null; // uuid, nullable
+  provider: string; // text with default value 'OPENAI'
+  heliconeProxyKeyID: string | null; // uuid, nullable
+}
+
+export interface ValhallaResponse {
+  id: string; // uuid
+  createdAt: Date; // timestamp with time zone
+  body: any; // jsonb
+  request: string; // uuid (foreign key to Request)
+  delayMs: number | null; // integer, nullable
+  http_status: number | null; // smallint, nullable
+  completionTokens: number | null; // integer, nullable
+  model: string | null; // text, nullable
+  promptTokens: number | null; // integer, nullable
+}
+
+export interface ValhallaCacheHits {
+  createdAt: Date; // timestamp with time zone
+  requestID: string; // uuid (foreign key to Request)
+}
+
+export interface ValhallaFeedback {
+  responseID: string; // uuid (foreign key to Response)
+  rating: boolean; // boolean
+  createdAt: Date; // timestamp with time zone
+  id: string; // uuid
+}

valhalla/packages/helicone-shared-ts/src/db/valhalla.ts

@@ -6,10 +6,18 @@ import {
   err,
   ok,
 } from "../modules/result";
+import { ValhallaRequest, ValhallaResponse } from "./valhalla.database.types";
 
 export interface IValhallaDB {
-  query(query: string): PromiseGenericResult<QueryResult<any>>;
+  query(query: string, values: any[]): PromiseGenericResult<QueryResult<any>>;
   now(): PromiseGenericResult<QueryResult<any>>;
+  insertRequest(
+    request: ValhallaRequest
+  ): PromiseGenericResult<QueryResult<any>>;
+  insertResponse(
+    response: ValhallaResponse
+  ): PromiseGenericResult<QueryResult<any>>;
+  close(): Promise<void>;
 }
 
 class ValhallaDB implements IValhallaDB {
@@ -68,18 +76,86 @@ class ValhallaDB implements IValhallaDB {
     }
   }
 
-  async query(query: string): PromiseGenericResult<QueryResult<any>> {
+  async close() {
+    await this.client.end();
+  }
+
+  async query(
+    query: string,
+    values: any[] = []
+  ): PromiseGenericResult<QueryResult<any>> {
     try {
       this.connect();
-      return ok(await this.client.query(query));
+      return ok(await this.client.query(query, values));
     } catch (thrownErr) {
+      console.error("Error in query", query, thrownErr);
       return err(JSON.stringify(thrownErr));
     }
   }
 
   async now() {
     return this.query("SELECT NOW() as now");
   }
+
+  async insertRequest(
+    request: ValhallaRequest
+  ): PromiseGenericResult<QueryResult<any>> {
+    const query = `
+      INSERT INTO request (
+        id,
+        created_at,
+        url_href,
+        user_id,
+        properties,
+        helicone_org_id,
+        provider,
+        body
+      )
+      VALUES (
+        $1, $2, $3, $4, $5, $6, $7, $8
+      );
+    `;
+    return this.query(query, [
+      request.id,
+      request.createdAt.toISOString(),
+      request.urlHref,
+      request.userId,
+      JSON.stringify(request.properties),
+      request.heliconeOrgID,
+      request.provider,
+      JSON.stringify(request.body),
+    ]);
+  }
+  async insertResponse(
+    response: ValhallaResponse
+  ): PromiseGenericResult<QueryResult<any>> {
+    const query = `
+    INSERT INTO response (
+      id,
+      created_at,
+      body,
+      request,
+      delay_ms,
+      http_status,
+      completion_tokens,
+      model,
+      prompt_tokens
+    )
+    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
+    RETURNING *
+  `;
+    return this.query(query, [
+      response.id,
+      response.createdAt.toISOString(),
+      JSON.stringify(response.body),
+      response.request,
+      response.delayMs,
+      response.http_status,
+      response.completionTokens,
+      response.model,
+      response.promptTokens,
+    ]);
+  }
 }
 
 class StaticValhallaPool {

valhalla/packages/helicone-shared-ts/src/requestWrapper/index.ts

@@ -0,0 +1,20 @@
+import {
+  Request as ExpressRequest,
+  Response as ExpressResponse,
+} from "express";
+
+export class RequestWrapper<T> {
+  constructor(private request: ExpressRequest) {}
+
+  public async getBody(): Promise<T> {
+    return await this.request.body;
+  }
+
+  public authHeader(): string | undefined {
+    return this.request.headers.authorization;
+  }
+
+  public heliconeOrgId(): string | undefined {
+    return (this.request.headers["helicone-org-id"] as string) ?? undefined;
+  }
+}

valhalla/packages/helicone-shared-ts/src/routers/iRouterWrapper.ts

@@ -5,16 +5,17 @@ import {
   Response as ExpressResponse,
 } from "express";
 import { SupabaseConnector } from "../db/supabase";
+import { RequestWrapper } from "../requestWrapper";
 
-export interface IRouterWrapper {
-  req: ExpressRequest;
+export interface IRouterWrapper<T> {
+  request: RequestWrapper<T>;
   res: ExpressResponse;
 }
 
-export interface IRouterWrapperDB extends IRouterWrapper {
+export interface IRouterWrapperDB<T> extends IRouterWrapper<T> {
   db: IValhallaDB;
 }
 
-export interface IRouterWrapperAuth extends IRouterWrapperDB {
+export interface IRouterWrapperAuth<T> extends IRouterWrapperDB<T> {
   supabaseClient: SupabaseConnector;
 }

valhalla/packages/helicone-shared-ts/src/routers/withAuth.ts

@@ -5,13 +5,14 @@ import {
 } from "express";
 import { withDB } from "./withDB";
 import { SupabaseConnector } from "../db/supabase";
+import { RequestWrapper } from "../requestWrapper";
 
-export function withAuth<T extends ExpressRequest, K extends ExpressResponse>(
-  fn: ({ db, req, res, supabaseClient }: IRouterWrapperAuth) => void
+export function withAuth<T>(
+  fn: ({ db, request, res, supabaseClient }: IRouterWrapperAuth<T>) => void
 ) {
-  return withDB(async ({ db, req, res }) => {
+  return withDB<T>(async ({ db, request, res }) => {
     const supabaseClient = new SupabaseConnector();
-    const authorizationString = req.headers.authorization;
+    const authorizationString = request.authHeader();
     if ("string" !== typeof authorizationString) {
       res.status(401).json({
         error: "No authorization header",
@@ -21,7 +22,7 @@ export function withAuth<T extends ExpressRequest, K extends ExpressResponse>(
 
     const isAuthenticated = await supabaseClient.authenticate(
       authorizationString,
-      req.headers["helicone-org-id"] as string
+      request.heliconeOrgId()
     );
     if (isAuthenticated.error) {
       res.status(401).json({
@@ -32,7 +33,7 @@ export function withAuth<T extends ExpressRequest, K extends ExpressResponse>(
 
     fn({
       db,
-      req,
+      request,
       res,
       supabaseClient: supabaseClient,
     });

valhalla/packages/helicone-shared-ts/src/routers/withDB.ts

@@ -1,18 +1,19 @@
 import { createValhallaClient } from "..";
+import { RequestWrapper } from "../requestWrapper";
 import { IRouterWrapperDB } from "./iRouterWrapper";
 import {
   Request as ExpressRequest,
   Response as ExpressResponse,
 } from "express";
 
-export function withDB<T extends ExpressRequest, K extends ExpressResponse>(
-  fn: ({ db, req, res }: IRouterWrapperDB) => void
+export function withDB<T>(
+  fn: ({ db, request, res }: IRouterWrapperDB<T>) => void
 ) {
-  return async (req: T, res: K) => {
+  return async (req: ExpressRequest, res: ExpressResponse) => {
     const valhallaDB = createValhallaClient();
     fn({
       db: valhallaDB,
-      req,
+      request: new RequestWrapper<T>(req),
       res,
     });
   };