Merge pull request #37 from upstash/add-helicone

feat: add helicone
upstash · Aug 9, 2024 · 1f50c91 · 1f50c91
2 parents 36852ae + edaa3a1
commit 1f50c91
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -360,6 +360,39 @@ await ragChat.history.addMessage({
 });
 ```
 
+### Add Observability via Helicone
+
+Helicone is a powerful observability platform that provides valuable insights into your LLM usage. Integrating Helicone with RAGChat is straightforward.
+
+To enable Helicone observability in RAGChat, you simply need to pass your Helicone API key when initializing your model. Here's how to do it for both custom models and OpenAI:
+
+#### For Custom Models (e.g., Meta-Llama)
+
+```ts
+import { RAGChat, custom } from "ragchat";
+
+const ragChat = new RAGChat({
+  model: custom("meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo", {
+    apiKey: "xxx",
+    baseUrl: "https://api.together.xyz",
+    analytics: { name: "helicone", token: process.env.HELICONE_API_KEY! },
+  }),
+});
+```
+
+#### For OpenAI Models
+
+```ts
+import { RAGChat, openai } from "ragchat";
+
+const ragChat = new RAGChat({
+  model: openai("gpt-3.5-turbo", {
+    apiKey: process.env.OPENAI_API_KEY!,
+    analytics: { name: "helicone", token: process.env.HELICONE_API_KEY! },
+  }),
+});
+```
+
 ## Example usage
 
 ### Nextjs route handlers

diff --git a/src/models.ts b/src/models.ts
@@ -42,7 +42,36 @@ export type LLMClientConfig = {
   baseUrl: string;
 };
 
-type ModelOptions = Omit<LLMClientConfig, "model">;
+type ModelOptions = Omit<LLMClientConfig, "model"> & {
+  analytics?: { name: "helicone"; token: string };
+};
+
+const analyticsBaseUrlMap = (
+  analyticsName: "helicone",
+  analyticsToken: string,
+  providerApiKey: string,
+  providerBaseUrl?: string
+) => {
+  return {
+    helicone: {
+      custom: {
+        baseURL: "https://gateway.helicone.ai",
+        defaultHeaders: {
+          "Helicone-Auth": `Bearer ${analyticsToken}`,
+          "Helicone-Target-Url": providerBaseUrl,
+          Authorization: `Bearer ${providerApiKey}`,
+        },
+      },
+      openai: {
+        basePath: "https://oai.helicone.ai/v1",
+        defaultHeaders: {
+          "Helicone-Auth": `Bearer ${analyticsToken}`,
+          Authorization: `Bearer ${providerApiKey}`,
+        },
+      },
+    },
+  }[analyticsName];
+};
 
 export const upstash = (model: UpstashChatModel, options?: Omit<ModelOptions, "baseUrl">) => {
   const apiKey = options?.apiKey ?? process.env.QSTASH_TOKEN ?? "";
@@ -70,18 +99,35 @@ export const custom = (model: string, options?: ModelOptions) => {
   return new ChatOpenAI({
     modelName: model,
     ...options,
-    configuration: {
-      apiKey: options.apiKey,
-      baseURL: options.baseUrl,
-    },
+    ...(options.analytics
+      ? {
+          configuration: analyticsBaseUrlMap(
+            options.analytics.name,
+            options.analytics.token,
+            options.apiKey,
+            options.baseUrl
+          ).custom,
+        }
+      : {
+          configuration: {
+            apiKey: options.apiKey,
+            baseURL: options.baseUrl,
+          },
+        }),
   });
 };
 
 export const openai = (model: OpenAIChatModel, options?: Omit<ModelOptions, "baseUrl">) => {
+  const apiKey = process.env.OPENAI_API_KEY ?? options?.apiKey ?? "";
+  const { analytics, ...optionsWithout } = options ?? {};
+
   return new ChatOpenAI({
     modelName: model,
     temperature: 0,
-    apiKey: process.env.OPENAI_API_KEY ?? options?.apiKey ?? "",
-    ...options,
+    ...optionsWithout,
+    apiKey,
+    ...(analytics
+      ? { configuration: analyticsBaseUrlMap(analytics.name, analytics.token, apiKey).openai }
+      : {}),
   });
 };