feat(aiproxy): api-key header (#32)

Handle error from OpenAI without locking the conversation
Support Api-Key header that Azure OpenAI use ( use SPIN_VARIABLE_OPENAI_API_KEY_METHOD environment variable to decide which method to use for passing the api key )
Allow refunding all tokens if there is no usage
Lock emscripten version. After emscripten 3.1.74 the standard libraries are built with non-mvp features. The NEAR wasm spec is at mvp, and so the build was broken.
Enhance test suites to support the different Api-Key passing methods

.devcontainer/install-dependencies.sh

@@ -34,6 +34,7 @@ echo 'export PATH="$(pwd)/wabt-1.0.35/bin:$PATH"' >> ~/.bashrc
 # Install Emscripten
 git clone https://github.com/emscripten-core/emsdk.git
 cd emsdk
+git checkout 3.1.74 
 ./emsdk install latest
 ./emsdk activate latest
 cd ..

examples/aiproxy/.gitignore

@@ -2,3 +2,4 @@ target
 .spin
 dist
 web4.js
+.env

examples/aiproxy/.test.env

@@ -1,5 +1,6 @@
 SPIN_VARIABLE_OPENAI_API_KEY=abcd
+SPIN_VARIABLE_OPENAI_API_KEY_METHOD=authorization  # or "api-key"
 SPIN_VARIABLE_REFUND_SIGNING_KEY=48QM3KLHFY22hDNnDx6zvgakY1dy66Jsv4dtTT6mt131DtjvPrQn7zyr3CVb1ZKPuVLbmvjQSK9o5vuEvMyiLR5Y
 SPIN_VARIABLE_FT_CONTRACT_ID=aitoken.test.near
 SPIN_VARIABLE_OPENAI_COMPLETIONS_ENDPOINT=http://127.0.0.1:3001/v1/chat/completions
-SPIN_VARIABLE_RPC_URL=http://localhost:14500
+SPIN_VARIABLE_RPC_URL=http://localhost:14500

examples/aiproxy/README.md

@@ -4,13 +4,14 @@ This folder contains a [Spin](https://www.fermyon.com/spin) application, based o
 
 There is a simple example of a web client in the [web](./web/) folder.
 
-The application will keep track of of token usage per conversation in the built-in key-value storage of Spin. The initial balance for a conversation is retrieved from the Fungible Token smart contract.
+The application will keep track of token usage per conversation in the built-in key-value storage of Spin. The initial balance for a conversation is retrieved from the Fungible Token smart contract.
 
 To launch the application, make sure to have the Spin SDK installed.
 
 You also need to set some environment variables:
 
 - `SPIN_VARIABLE_OPENAI_API_KEY` your OpenAI API key.
+- `SPIN_VARIABLE_OPENAI_API_KEY_METHOD` specifies the method to provide the API key. Use `authorization` for OpenAI (default) and `api-key` for Azure OpenAI.
 - `SPIN_VARIABLE_REFUND_SIGNING_KEY` an ed21159 secret key that will be used to sign refund requests. You can run the [create-refund-signing-keypair.js](./create-refund-signing-keypair.js) script to create the keypair. Run it using the command `$(node create-refund-signing-keypair.js)` and it will set the environment variable for you.
 - `SPIN_VARIABLE_FT_CONTRACT_ID` the NEAR contract account id. e.g `aitoken.test.near`
 - `SPIN_VARIABLE_OPENAI_COMPLETIONS_ENDPOINT` OpenAI API completions endpoint. E.g. https://api.openai.com/v1/chat/completions

examples/aiproxy/openai-proxy/src/lib.rs

@@ -247,6 +247,19 @@ async fn handle_request(request: Request, response_out: ResponseOutparam) {
 
             match proxy_openai(messages).await {
                 Ok(incoming_response) => {
+                    if incoming_response.status() != 200 {
+                        conversation_balance.locked_for_ongoing_request = false;
+                        conversation_balance_store
+                            .set_json(conversation_id, &conversation_balance)
+                            .unwrap();
+                        let response_data = incoming_response.into_body().await.unwrap();
+                        let response_string = String::from_utf8(response_data).unwrap();
+                        eprintln!(
+                            "error in response from OpenAI endpoint: {:?}",
+                            response_string
+                        );
+                        return server_error(response_out);
+                    }
                     let mut incoming_response_body = incoming_response.take_body_stream();
                     let outgoing_response = OutgoingResponse::new(headers);
                     let mut outgoing_response_body = outgoing_response.take_body();
@@ -458,18 +471,24 @@ async fn proxy_openai(messages: Value) -> anyhow::Result<IncomingResponse> {
     });
 
     let openai_completions_endpoint = variables::get("openai_completions_endpoint")?;
-    let outgoing_request = Request::builder()
+    let api_key = variables::get("openai_api_key").unwrap();
+    let api_key_method =
+        variables::get("openai_api_key_method").unwrap_or_else(|_| "authorization".to_string());
+
+    let mut openai_request_builder = Request::builder();
+    openai_request_builder
         .method(Method::Post)
         .uri(openai_completions_endpoint)
-        .header(
-            "Authorization",
-            format!("Bearer {}", variables::get("openai_api_key").unwrap()),
-        )
-        .header("Content-Type", "application/json")
-        .body(request_body.to_string())
-        .build();
-
-    let response = match http::send::<_, IncomingResponse>(outgoing_request).await {
+        .header("Content-Type", "application/json");
+
+    let openai_request = match api_key_method.as_str() {
+        "api-key" => openai_request_builder.header("Api-Key", api_key),
+        _ => openai_request_builder.header("Authorization", format!("Bearer {}", api_key)),
+    }
+    .body(request_body.to_string())
+    .build();
+
+    let response = match http::send::<_, IncomingResponse>(openai_request).await {
         Ok(resp) => resp,
         Err(e) => {
             eprintln!("Error sending request to OpenAI: {e}");

examples/aiproxy/playwright-tests/aiproxy.spec.js

@@ -1,7 +1,64 @@
-
 import { test, expect } from '@playwright/test';
+import { spawn } from 'child_process';
+import { fileURLToPath } from 'url';
+import path from 'path';
+import http from 'http';
+
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = path.dirname(__filename);
+const mockServerPath = path.resolve(__dirname, 'openaimockserver.js');
+
+let mockServerProcess;
+
+async function startMockServer(apiKeyMethod, apikey = 'abcd') {
+  if (mockServerProcess) {
+    await new Promise((resolve) => {
+      mockServerProcess.on('close', resolve);
+      mockServerProcess.kill();
+    });
+  }
+
+  mockServerProcess = spawn('node', [mockServerPath], {
+    env: {
+      ...process.env,
+      SPIN_VARIABLE_OPENAI_API_KEY: apikey,
+      SPIN_VARIABLE_OPENAI_API_KEY_METHOD: apiKeyMethod,
+    },
+  });
+
+  mockServerProcess.stdout.on('data', (data) => {
+    console.log(`stdout: ${data}`);
+  });
+
+  mockServerProcess.stderr.on('data', (data) => {
+    console.error(`stderr: ${data}`);
+  });
+
+  // Wait for the server to start and respond on port 3001
+  await new Promise((resolve, reject) => {
+    const interval = setInterval(() => {
+      http.get('http://127.0.0.1:3001', (res) => {
+        if (res.statusCode === 200) {
+          clearInterval(interval);
+          resolve();
+        }
+      }).on('error', () => {
+        // Ignore errors, keep trying
+      });
+    }, 500);
+  });
+}
+
+test.afterEach(async () => {
+  if (mockServerProcess) {
+    await new Promise((resolve) => {
+      mockServerProcess.on('close', resolve);
+      mockServerProcess.kill();
+    });
+  }
+});
 
-test('ask question', async ({ page }) => {
+async function testConversation({page, expectedRefundAmount = "127999973", expectedOpenAIResponse = "Hello! How can I assist you today?"}) {
   const { functionAccessKeyPair, publicKey, accountId, contractId } = await fetch('http://localhost:14501').then(r => r.json());
 
   await page.goto('/');
@@ -20,10 +77,28 @@ test('ask question', async ({ page }) => {
   questionArea.fill("Hello!");
   await page.waitForTimeout(500);
   await page.getByRole('button', { name: 'Ask AI' }).click();
-  await expect(await page.getByText("Hello! How can I assist you today?")).toBeVisible();
+  await expect(await page.getByText(expectedOpenAIResponse)).toBeVisible();
 
   await page.waitForTimeout(500);
   await page.locator("#refundButton").click();
 
-  await expect(await page.locator("#refund_message")).toContainText(`EVENT_JSON:{"standard":"nep141","version":"1.0.0","event":"ft_transfer","data":[{"old_owner_id":"${contractId}","new_owner_id":"${accountId}","amount":"127999973"}]}\nrefunded 127999973 to ${accountId}`);
+  await expect(await page.locator("#refund_message")).toContainText(`EVENT_JSON:{"standard":"nep141","version":"1.0.0","event":"ft_transfer","data":[{"old_owner_id":"${contractId}","new_owner_id":"${accountId}","amount":"${expectedRefundAmount}"}]}\nrefunded ${expectedRefundAmount} to ${accountId}`, {timeout: 10_000});
+}
+
+test('start conversation, ask question and refund (using OpenAI authorization header)', async ({ page }) => {
+  await startMockServer('authorization');
+  await testConversation({page});
+});
+
+
+
+test('start conversation, ask question and refund (using Azure OpenAI Api-Key header)', async ({ page }) => {
+  await startMockServer('api-key');
+  await testConversation({page});
+});
+
+test('start conversation, ask question, where openai API fails, and refund (using wrong OpenAI API key)', async ({ page }) => {
+  await startMockServer('api-key', "1234ffff");
+
+  await testConversation({page, expectedRefundAmount: "128000000", expectedOpenAIResponse: "Failed to fetch from proxy: Internal Server Error"});
 });

examples/aiproxy/playwright-tests/near_rpc.js

@@ -37,7 +37,7 @@ await aiuser.call(aiTokenAccount.accountId, 'storage_deposit', {
 
 await aiTokenAccount.call(aiTokenAccount.accountId, 'ft_transfer', {
     receiver_id: aiuser.accountId,
-    amount: 128_000_000n.toString(),
+    amount: (100n*128_000_000n).toString(),
 }, {
     attachedDeposit: 1n.toString()
 });

examples/aiproxy/playwright-tests/openaimockserver.js

@@ -2,6 +2,8 @@ import { createServer } from "http";
 import { Readable } from "stream";
 
 const PORT = 3001;
+const API_KEY = process.env.SPIN_VARIABLE_OPENAI_API_KEY;
+const API_KEY_METHOD = process.env.SPIN_VARIABLE_API_KEY_METHOD || "authorization";
 
 const server = createServer((req, res) => {
   if (req.method === "POST" && req.url.startsWith("/v1/chat/completions")) {
@@ -12,6 +14,21 @@ const server = createServer((req, res) => {
     });
 
     req.on("end", () => {
+      let apiKeyValid = false;
+
+      if (API_KEY_METHOD === "api-key") {
+        apiKeyValid = req.headers["api-key"] === API_KEY;
+      } else {
+        const authHeader = req.headers["authorization"];
+        apiKeyValid = authHeader && authHeader === `Bearer ${API_KEY}`;
+      }
+
+      if (!apiKeyValid) {
+        res.writeHead(401, { "Content-Type": "application/json" });
+        res.end(JSON.stringify({ error: "Unauthorized" }));
+        return;
+      }
+
       const responseChunks = [
         JSON.stringify({
           choices: [

examples/aiproxy/playwright.config.js

@@ -1,4 +1,18 @@
 import { defineConfig, devices } from '@playwright/test';
+import fs from 'fs';
+import path from 'path';
+import dotenv from 'dotenv';
+import { fileURLToPath } from 'url';
+
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = path.dirname(__filename);
+
+const envFilePath = path.resolve(__dirname, '.test.env');
+const envConfig = dotenv.parse(fs.readFileSync(envFilePath));
+
+for (const k in envConfig) {
+  process.env[k] = envConfig[k];
+}
 
 /**
  * @see https://playwright.dev/docs/test-configuration
@@ -64,25 +78,22 @@ export default defineConfig({
   ],
 
   /* Run your local dev server before starting the tests */
-  webServer: [{
-    command: 'static-web-server -p 8080 -d web',
-    url: 'http://127.0.0.1:8080',
-    reuseExistingServer: !process.env.CI,
-  },
-  {
-    command: "export $(grep -v '^#' .test.env | xargs) && spin build && spin up",
-    url: 'http://127.0.0.1:3000',
-    reuseExistingServer: !process.env.CI,
-  },
-  {
-    command: "export $(grep -v '^#' .test.env | xargs) && node playwright-tests/openaimockserver.js",
-    url: 'http://127.0.0.1:3001',
-    reuseExistingServer: !process.env.CI,
-  },
-  {
-    command: "export $(grep -v '^#' .test.env | xargs) && node playwright-tests/near_rpc.js",
-    url: 'http://127.0.0.1:14501',
-    reuseExistingServer: !process.env.CI,
-  }],
+  webServer: [
+    {
+      command: "spin build && spin up",
+      url: 'http://127.0.0.1:3000',
+      reuseExistingServer: !process.env.CI,
+    },
+    {
+      command: "node playwright-tests/near_rpc.js",
+      url: 'http://127.0.0.1:14501',
+      reuseExistingServer: !process.env.CI,
+    },
+    {
+      command: "npx http-server ./web -p 8080",
+      url: 'http://127.0.0.1:8080',
+      reuseExistingServer: !process.env.CI,
+    }
+  ],
 });
 

examples/aiproxy/spin.toml

@@ -21,13 +21,15 @@ workdir = "openai-proxy"
 watch = ["src/**/*.rs", "Cargo.toml"]
 
 [variables]
+openai_api_key_method = { default = "authorization" }
 openai_api_key = { required = true }
 refund_signing_key = { required = true }
 openai_completions_endpoint = { required = true }
 ft_contract_id = { required = true }
 rpc_url = {required = true }
 
 [component.openai-proxy.variables]
+openai_api_key_method = "{{ openai_api_key_method }}"
 openai_api_key = "{{ openai_api_key }}"
 refund_signing_key = "{{ refund_signing_key }}"
 openai_completions_endpoint =  "{{ openai_completions_endpoint }}"