feat: semantic cache (#1643)

* setup semantic cache in unkey repo

* gitignore

* add action deploy

* update action

* update action

* update actions

* try alternative wrangler deploy

* install

* update workflow

* rerun

* add fe

* check in fe

* log

* update deploy path

* only use latest message

* regen lockfile

* use @unkey/cache

* lockfile

* update dependencies

* lockfile: use versioned cache

* remove fe

* lockfile

* lockfile

* update workflow

* update workflow

* restore pnpm workspace

* send data to tinybird correctly

* add scripts

* fix types

* move app to /apps

* remove app from packages

* Add semantic cache to desktop sidebar

* Add semantic cache to mobile sidebar

* Add header to semantic caching page

* Add basic logging

* merge main

* formatting fixes

* remove action

* format code

* remove from sidebars

* update types

* clean up deps

* parameterize

* fix types

* format

* Regenerate lockfile

* formatting

* Refactor into multiple files

* remove tsup

* remove tsc build step

* infer analyticsevent type

* respond to feedback

* lockfile

* add tinybird data

* fix(worker.ts): remove unused async middleware function
chore(package.json): update license to AGPL-3.0
chore(package.json): remove unnecessary files field
chore(package.json): update dev script in semantic-cache package
chore(package.json): update dev script in api package
chore(package.json): update dev script in hono package
chore(package.json): update dev script in nextjs package
chore(package.json): update dev script in nuxt package
chore(package.json): update dev script in ratelimit package
chore(package.json): update dev script in rbac package

---------

Co-authored-by: chronark <[email protected]>

apps/api/src/worker.ts

@@ -35,34 +35,6 @@ app.use("*", init());
 app.use("*", cors());
 app.use("*", metrics());
 
-app.use("*", async (c, next) => {
-  try {
-    if (c.env.TINYBIRD_PROXY_URL) {
-      const start = performance.now();
-      const p = fetch(new URL("/v0/incr", c.env.TINYBIRD_PROXY_URL), {
-        method: "POST",
-      }).then(() => {
-        const { metrics } = c.get("services");
-
-        metrics.emit({
-          metric: "metric.koyeb.latency",
-          // @ts-expect-error
-          continent: c.req.raw?.cf?.continent,
-          // @ts-expect-error
-          colo: c.req.raw?.cf?.colo,
-          latency: performance.now() - start,
-        });
-      });
-
-      c.executionCtx.waitUntil(p);
-    }
-  } catch (e) {
-    console.error(e);
-  }
-
-  return next();
-});
-
 /**
  * Registering all route handlers
  */

apps/dashboard/app/(app)/desktop-sidebar.tsx

@@ -8,6 +8,7 @@ import {
   BookOpen,
   Code,
   Crown,
+  DatabaseZap,
   GlobeLock,
   Loader2,
   type LucideIcon,
@@ -113,6 +114,13 @@ export const DesktopSidebar: React.FC<Props> = ({ workspace, className }) => {
       tag: <Tag label="internal" />,
       hidden: !workspace.features.successPage,
     },
+    {
+      icon: DatabaseZap,
+      href: "/semantic-cache",
+      label: "Semantic Cache",
+      active: segments.at(0) === "semantic-cache",
+      hidden: true,
+    },
   ].filter((n) => !n.hidden);
 
   const firstOfNextMonth = new Date();

apps/dashboard/app/(app)/mobile-sidebar.tsx

@@ -3,7 +3,7 @@ import { Button } from "@/components/ui/button";
 import { Sheet, SheetClose, SheetContent, SheetHeader, SheetTrigger } from "@/components/ui/sheet";
 import { cn } from "@/lib/utils";
 import { SignOutButton } from "@clerk/nextjs";
-import { BookOpen, FileJson, LogOut, Menu, Settings } from "lucide-react";
+import { BookOpen, DatabaseZap, FileJson, LogOut, Menu, Settings } from "lucide-react";
 import Link from "next/link";
 import { useRouter } from "next/navigation";
 import { WorkspaceSwitcher } from "./team-switcher";

apps/dashboard/app/(app)/semantic-cache/page.tsx

@@ -0,0 +1,55 @@
+import { PageHeader } from "@/components/dashboard/page-header";
+import { Separator } from "@/components/ui/separator";
+
+import {
+  Table,
+  TableBody,
+  TableCaption,
+  TableCell,
+  TableHead,
+  TableHeader,
+  TableRow,
+} from "@/components/ui/table";
+
+import { getAllSemanticCacheLogs } from "@/lib/tinybird";
+
+export default async function SemanticCachePage() {
+  const { data } = await getAllSemanticCacheLogs({ limit: 10 });
+  return (
+    <div>
+      <PageHeader
+        title="Semantic Cache"
+        description="Faster, cheaper LLM API calls through semantic caching"
+      />
+      <Separator className="my-6" />
+      <h1 className="font-medium">Logs</h1>
+      <Table className="mt-4">
+        <TableCaption>View real-time logs from the semantic cache.</TableCaption>
+        <TableHeader>
+          <TableRow>
+            <TableHead>Time</TableHead>
+            <TableHead>Model</TableHead>
+            <TableHead>Cache status</TableHead>
+            <TableHead>Query</TableHead>
+            <TableHead>Response</TableHead>
+            <TableHead>Request ID</TableHead>
+            <TableHead>Request timing</TableHead>
+          </TableRow>
+        </TableHeader>
+        <TableBody>
+          {data.map((data) => (
+            <TableRow key={data.requestId}>
+              <TableCell className="font-medium">{data.timestamp}</TableCell>
+              <TableCell>{data.model}</TableCell>
+              <TableCell>{data.cache}</TableCell>
+              <TableCell>{data.query}</TableCell>
+              <TableCell>{data.response}</TableCell>
+              <TableCell>{data.requestId}</TableCell>
+              <TableCell>{data.timing}</TableCell>
+            </TableRow>
+          ))}
+        </TableBody>
+      </Table>
+    </div>
+  );
+}

apps/dashboard/lib/tinybird.ts

@@ -721,6 +721,28 @@ export const getRatelimitEvents = tb.buildPipe({
   },
 });
 
+export const getAllSemanticCacheLogs = tb.buildPipe({
+  pipe: "get_all_semantic_cache_logs__v1",
+  parameters: z.object({
+    limit: z.number().optional(),
+  }),
+  data: z.object({
+    timestamp: z.string(),
+    model: z.string(),
+    stream: z.number(),
+    query: z.string(),
+    vector: z.array(z.number()),
+    response: z.string(),
+    cache: z.number(),
+    timing: z.number(),
+    tokens: z.number(),
+    requestId: z.string(),
+  }),
+  opts: {
+    cache: "no-store",
+  },
+});
+
 // public get getVerificationsByOwnerId() {
 //   return this.client.buildPipe({
 //     pipe: "get_verifictions_by_keySpaceId__v1",

apps/semantic-cache/.gitignore

@@ -0,0 +1,5 @@
+dist
+node_modules
+.env
+scripts
+.dev.vars

apps/semantic-cache/lib/cache.ts

@@ -0,0 +1,27 @@
+import { CloudflareStore, type MemoryStore, Namespace, createCache } from "@unkey/cache";
+import type { Context } from "hono";
+
+import type { LLMResponse } from "../types";
+
+export async function initCache(c: Context, memory: MemoryStore<string, any>) {
+  const context = c.executionCtx;
+  const fresh = 6_000_000;
+  const stale = 300_000_000;
+
+  const cache = createCache({
+    response: new Namespace<LLMResponse>(context, {
+      stores: [
+        memory,
+        new CloudflareStore({
+          cloudflareApiKey: c.env.CLOUDFLARE_API_KEY,
+          zoneId: c.env.CLOUDFLARE_ZONE_ID,
+          domain: "cache.unkey.dev",
+        }),
+      ],
+      fresh,
+      stale,
+    }),
+  });
+
+  return cache;
+}

apps/semantic-cache/lib/streaming.ts

@@ -0,0 +1,48 @@
+export class ManagedStream {
+  stream: ReadableStream;
+  reader: ReadableStreamDefaultReader<Uint8Array>;
+  isDone: boolean;
+  data: string;
+  isComplete: boolean;
+
+  constructor(stream: ReadableStream) {
+    this.stream = stream;
+    this.reader = this.stream.getReader();
+    this.isDone = false;
+    this.data = "";
+    this.isComplete = false;
+  }
+
+  async readToEnd() {
+    try {
+      while (true) {
+        const { done, value } = await this.reader.read();
+        if (done) {
+          this.isDone = true;
+          break;
+        }
+        this.data += new TextDecoder().decode(value);
+      }
+    } catch (error) {
+      console.error("Stream error:", error);
+      this.isDone = false;
+    } finally {
+      this.reader.releaseLock();
+    }
+    return this.isDone;
+  }
+
+  checkComplete() {
+    if (this.data.includes("[DONE]")) {
+      this.isComplete = true;
+    }
+  }
+
+  getReader() {
+    return this.reader;
+  }
+
+  getData() {
+    return this.data;
+  }
+}

apps/semantic-cache/package.json

@@ -0,0 +1,35 @@
+{
+  "name": "@unkey/semantic-cache",
+  "version": "1.0.0",
+  "license": "AGPL-3.0",
+  "private": true,
+  "keywords": ["unkey", "semantic", "cache", "ai"],
+  "bugs": {
+    "url": "https://github.com/unkeyed/unkey/issues"
+  },
+  "homepage": "https://github.com/unkeyed/unkey#readme",
+  "author": "Dominic Eccleston [email protected]",
+  "scripts": {
+    "dev": "wrangler dev src/worker.ts"
+  },
+  "devDependencies": {
+    "@cloudflare/workers-types": "^4.20240403.0",
+    "@types/node": "^20.12.7",
+    "dotenv": "^16.4.5",
+    "typescript": "^5.3.3",
+    "wrangler": "^3.47.0"
+  },
+  "dependencies": {
+    "@chronark/zod-bird": "^0.3.9",
+    "@cloudflare/ai": "^1.1.0",
+    "@unkey/cache": "workspace:^",
+    "@unkey/logs": "workspace:^",
+    "@unkey/metrics": "workspace:^",
+    "ai": "^3.0.23",
+    "hono": "^4.2.7",
+    "nanoid": "^5.0.7",
+    "openai": "^4.35.0",
+    "superjson": "^2.2.1",
+    "zod": "^3.23.5"
+  }
+}

apps/semantic-cache/src/analytics.ts

@@ -0,0 +1,32 @@
+import { Tinybird } from "@chronark/zod-bird";
+import { z } from "zod";
+
+export const eventSchema = z.object({
+  timestamp: z.string(),
+  model: z.string(),
+  stream: z.boolean(),
+  query: z.string(),
+  vector: z.array(z.number()),
+  response: z.string(),
+  cache: z.boolean(),
+  timing: z.number(),
+  tokens: z.number(),
+  requestId: z.string(),
+});
+
+export class Analytics {
+  public readonly client: Tinybird;
+
+  constructor(opts: {
+    tinybirdToken: string;
+  }) {
+    this.client = new Tinybird({ token: opts.tinybirdToken });
+  }
+
+  public get ingestLogs() {
+    return this.client.buildIngestEndpoint({
+      datasource: "semantic_cache__v3",
+      event: eventSchema,
+    });
+  }
+}