From 80cd6a11bcb2bd2148b1992d00bb756c44daa11c Mon Sep 17 00:00:00 2001
From: Travis Fischer <fisch0920@gmail.com>
Date: Wed, 31 May 2023 23:20:43 -0700
Subject: [PATCH] feat: add some more unit tests for structured llm outputs

---
 test/_preload.ts    |  1 -
 test/_utils.ts      |  2 +-
 test/openai.test.ts | 57 ++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 57 insertions(+), 3 deletions(-)
 delete mode 100644 test/_preload.ts

diff --git a/test/_preload.ts b/test/_preload.ts
deleted file mode 100644
index 9a1976a..0000000
--- a/test/_preload.ts
+++ /dev/null
@@ -1 +0,0 @@
-import 'dotenv/config'
diff --git a/test/_utils.ts b/test/_utils.ts
index deb0cb5..d47d345 100644
--- a/test/_utils.ts
+++ b/test/_utils.ts
@@ -20,7 +20,7 @@ const redis = new Redis(process.env.REDIS_URL_TEST!)
 const keyvRedis = new KeyvRedis(redis)
 const keyv = new Keyv({ store: keyvRedis, namespace: 'agentic-test' })
 
-// TODO: this is pretty hacky
+// TODO: this is a lil hacky
 const keyvHas = (keyv.has as any).bind(keyv)
 keyv.has = async (key, ...rest) => {
   if (refreshTestCache) {
diff --git a/test/openai.test.ts b/test/openai.test.ts
index c439ae6..5a782c9 100644
--- a/test/openai.test.ts
+++ b/test/openai.test.ts
@@ -1,11 +1,43 @@
 import test from 'ava'
-// import 'dotenv/config'
 import { expectTypeOf } from 'expect-type'
 import { z } from 'zod'
 
 import { OpenAIChatModelBuilder } from '../src/openai'
 import { createOpenAITestClient } from './_utils'
 
+test('OpenAIChatModel ⇒ string output', async (t) => {
+  t.timeout(2 * 60 * 1000)
+  const client = createOpenAITestClient()
+
+  const builder = new OpenAIChatModelBuilder(client, {
+    modelParams: {
+      temperature: 0,
+      max_tokens: 30
+    },
+    messages: [
+      {
+        role: 'system',
+        content:
+          'You are a human, not a large language model. You are not an AI. Respond like a human would. Your name is Leeroy Jenkins.'
+      },
+      {
+        role: 'user',
+        content: 'what is your name?'
+      }
+    ]
+  })
+
+  const result = await builder.call()
+  t.truthy(typeof result === 'string')
+
+  expectTypeOf(result).toMatchTypeOf<string>()
+
+  const result2 = await builder.output(z.string()).call()
+  t.truthy(typeof result2 === 'string')
+
+  expectTypeOf(result2).toMatchTypeOf<string>()
+})
+
 test('OpenAIChatModel ⇒ json output', async (t) => {
   t.timeout(2 * 60 * 1000)
   const client = createOpenAITestClient()
@@ -30,3 +62,26 @@ test('OpenAIChatModel ⇒ json output', async (t) => {
 
   expectTypeOf(result).toMatchTypeOf<{ foo: string; bar: number }>()
 })
+
+test.only('OpenAIChatModel ⇒ boolean output', async (t) => {
+  t.timeout(2 * 60 * 1000)
+  const client = createOpenAITestClient()
+
+  const builder = new OpenAIChatModelBuilder(client, {
+    modelParams: {
+      temperature: 0,
+      max_tokens: 30
+    },
+    messages: [
+      {
+        role: 'user',
+        content: 'are you alive?'
+      }
+    ]
+  }).output(z.boolean())
+
+  const result = await builder.call()
+  t.truthy(typeof result === 'boolean')
+
+  expectTypeOf(result).toMatchTypeOf<boolean>()
+})