From edb5a16465d75faf0ba1205e2aa2bd0912c66eb6 Mon Sep 17 00:00:00 2001
From: Simon Farshid <simon.farshid@outlook.com>
Date: Wed, 21 Aug 2024 18:52:50 -0700
Subject: [PATCH] feat: DangerousInBrowserRuntime (#698)

---
 .changeset/four-bears-explain.md              |  5 +++
 packages/react/src/edge.ts                    |  5 ++-
 .../DangerousInBrowserAdapter.ts              | 39 +++++++++++++++++++
 .../runtimes/dangerous-in-browser/index.ts    |  4 ++
 .../useDangerousInBrowserRuntime.ts           | 17 ++++++++
 .../src/runtimes/edge/createEdgeRuntimeAPI.ts |  4 +-
 packages/react/src/runtimes/index.ts          |  1 +
 7 files changed, 72 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/four-bears-explain.md
 create mode 100644 packages/react/src/runtimes/dangerous-in-browser/DangerousInBrowserAdapter.ts
 create mode 100644 packages/react/src/runtimes/dangerous-in-browser/index.ts
 create mode 100644 packages/react/src/runtimes/dangerous-in-browser/useDangerousInBrowserRuntime.ts

diff --git a/.changeset/four-bears-explain.md b/.changeset/four-bears-explain.md
new file mode 100644
index 0000000000..075f938eae
--- /dev/null
+++ b/.changeset/four-bears-explain.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react": patch
+---
+
+feat: DangerousInBrowserRuntime
diff --git a/packages/react/src/edge.ts b/packages/react/src/edge.ts
index 40d568156b..f5ba2f38c7 100644
--- a/packages/react/src/edge.ts
+++ b/packages/react/src/edge.ts
@@ -1 +1,4 @@
-export { createEdgeRuntimeAPI, getEdgeRuntimeResponse } from "./runtimes/edge/createEdgeRuntimeAPI";
+export {
+  createEdgeRuntimeAPI,
+  getEdgeRuntimeResponse,
+} from "./runtimes/edge/createEdgeRuntimeAPI";
diff --git a/packages/react/src/runtimes/dangerous-in-browser/DangerousInBrowserAdapter.ts b/packages/react/src/runtimes/dangerous-in-browser/DangerousInBrowserAdapter.ts
new file mode 100644
index 0000000000..d63e67044a
--- /dev/null
+++ b/packages/react/src/runtimes/dangerous-in-browser/DangerousInBrowserAdapter.ts
@@ -0,0 +1,39 @@
+import { ChatModelAdapter, ChatModelRunOptions } from "../local";
+import { toCoreMessages } from "../edge/converters/toCoreMessages";
+import { toLanguageModelTools } from "../edge/converters/toLanguageModelTools";
+import { EdgeRuntimeRequestOptions } from "../edge/EdgeRuntimeRequestOptions";
+import { runResultStream } from "../edge/streams/runResultStream";
+import { toolResultStream } from "../edge/streams/toolResultStream";
+import { asAsyncIterable } from "../edge/EdgeChatAdapter";
+import {
+  CreateEdgeRuntimeAPIOptions,
+  getEdgeRuntimeStream,
+} from "../edge/createEdgeRuntimeAPI";
+
+export type DangerousInBrowserAdapterOptions = CreateEdgeRuntimeAPIOptions;
+
+export class DangerousInBrowserAdapter implements ChatModelAdapter {
+  constructor(private options: DangerousInBrowserAdapterOptions) {}
+
+  async *run({ messages, abortSignal, config }: ChatModelRunOptions) {
+    const res = await getEdgeRuntimeStream({
+      options: this.options,
+      abortSignal,
+      requestData: {
+        system: config.system,
+        messages: toCoreMessages(messages),
+        tools: config.tools ? toLanguageModelTools(config.tools) : [],
+        ...config.callSettings,
+        ...config.config,
+      } satisfies EdgeRuntimeRequestOptions,
+    });
+
+    const stream = res
+      .pipeThrough(toolResultStream(config.tools, abortSignal))
+      .pipeThrough(runResultStream());
+
+    for await (const update of asAsyncIterable(stream)) {
+      yield update;
+    }
+  }
+}
diff --git a/packages/react/src/runtimes/dangerous-in-browser/index.ts b/packages/react/src/runtimes/dangerous-in-browser/index.ts
new file mode 100644
index 0000000000..461db9d165
--- /dev/null
+++ b/packages/react/src/runtimes/dangerous-in-browser/index.ts
@@ -0,0 +1,4 @@
+export {
+  useDangerousInBrowserRuntime,
+  type DangerousInBrowserRuntimeOptions,
+} from "./useDangerousInBrowserRuntime";
diff --git a/packages/react/src/runtimes/dangerous-in-browser/useDangerousInBrowserRuntime.ts b/packages/react/src/runtimes/dangerous-in-browser/useDangerousInBrowserRuntime.ts
new file mode 100644
index 0000000000..e06949b206
--- /dev/null
+++ b/packages/react/src/runtimes/dangerous-in-browser/useDangerousInBrowserRuntime.ts
@@ -0,0 +1,17 @@
+import { LocalRuntimeOptions, useLocalRuntime } from "..";
+import { useState } from "react";
+import {
+  DangerousInBrowserAdapter,
+  DangerousInBrowserAdapterOptions,
+} from "./DangerousInBrowserAdapter";
+
+export type DangerousInBrowserRuntimeOptions =
+  DangerousInBrowserAdapterOptions & LocalRuntimeOptions;
+
+export const useDangerousInBrowserRuntime = ({
+  initialMessages,
+  ...options
+}: DangerousInBrowserRuntimeOptions) => {
+  const [adapter] = useState(() => new DangerousInBrowserAdapter(options));
+  return useLocalRuntime(adapter, { initialMessages });
+};
diff --git a/packages/react/src/runtimes/edge/createEdgeRuntimeAPI.ts b/packages/react/src/runtimes/edge/createEdgeRuntimeAPI.ts
index 0ac682a295..04563295fc 100644
--- a/packages/react/src/runtimes/edge/createEdgeRuntimeAPI.ts
+++ b/packages/react/src/runtimes/edge/createEdgeRuntimeAPI.ts
@@ -41,7 +41,7 @@ type LanguageModelCreator = (
   config: LanguageModelConfig,
 ) => Promise<LanguageModelV1> | LanguageModelV1;
 
-type CreateEdgeRuntimeAPIOptions = LanguageModelV1CallSettings & {
+export type CreateEdgeRuntimeAPIOptions = LanguageModelV1CallSettings & {
   model: LanguageModelV1 | LanguageModelCreator;
   system?: string;
   tools?: Record<string, Tool<any, any>>;
@@ -63,7 +63,7 @@ type GetEdgeRuntimeStreamOptions = {
   options: CreateEdgeRuntimeAPIOptions;
 };
 
-const getEdgeRuntimeStream = async ({
+export const getEdgeRuntimeStream = async ({
   abortSignal,
   requestData: unsafeRequest,
   options: {
diff --git a/packages/react/src/runtimes/index.ts b/packages/react/src/runtimes/index.ts
index 420a951784..a5746e8ed1 100644
--- a/packages/react/src/runtimes/index.ts
+++ b/packages/react/src/runtimes/index.ts
@@ -2,3 +2,4 @@ export * from "./core";
 export * from "./local";
 export * from "./edge";
 export * from "./external-store";
+export * from "./dangerous-in-browser";