Add support for token-less FixieCorpus tool (#269)

This change:
- Changes tools to be `Component<any>` (which is compatible with the
existing signature) so they can access context and/or return other
AI.JSX nodes
- Adds an `<ExecuteFunction>` component that encapsulates error handling
so it can be shared between the two `UseTools` implementations
- Adds `FixieAPIContext` to AI.JSX where the Fixie API url/token can be
configured
- Adds `FixieCorpus.createTool` which facilitates `FixieCorpus`-based
`Tool` creation and reads the API configuration from context (rather
than from `process.env`).

packages/ai-jsx/package.json

@@ -4,7 +4,7 @@
   "repository": "fixie-ai/ai-jsx",
   "bugs": "https://github.com/fixie-ai/ai-jsx/issues",
   "homepage": "https://ai-jsx.com",
-  "version": "0.16.0",
+  "version": "0.17.0",
   "volta": {
     "extends": "../../package.json"
   },
@@ -45,6 +45,15 @@
         "default": "./dist/cjs/batteries/sidekick/platform/sidekick.cjs"
       }
     },
+    "./batteries/fixie": {
+      "import": {
+        "types": "./dist/esm/batteries/fixie.d.ts",
+        "default": "./dist/esm/batteries/fixie.js"
+      },
+      "require": {
+        "default": "./dist/cjs/batteries/fixie.cjs"
+      }
+    },
     "./batteries/docs": {
       "import": {
         "types": "./dist/esm/batteries/docs.d.ts",

packages/ai-jsx/src/batteries/docs.tsx

@@ -18,6 +18,8 @@ import * as AI from '../index.js';
 import { Node } from '../index.js';
 import { getEnvVar } from '../lib/util.js';
 import { JsonChatCompletion } from './constrained-output.js';
+import { DEFAULT_API_CONFIGURATION, FixieAPIConfiguration, FixieAPIContext } from './fixie.js';
+import { Tool } from './use-tools.js';
 
 /**
  * A raw document loaded from an arbitrary source that has not yet been parsed.
@@ -616,31 +618,29 @@ export class LocalCorpus<
 
 /** A fully mananged {@link Corpus} served by Fixie. */
 export class FixieCorpus<ChunkMetadata extends Jsonifiable = Jsonifiable> implements Corpus<ChunkMetadata> {
-  private static readonly DEFAULT_FIXIE_API_URL = 'https://api.fixie.ai';
-
-  private readonly fixieApiUrl: string;
-
-  constructor(private readonly corpusId: string, private readonly fixieApiKey?: string) {
-    if (!fixieApiKey) {
-      this.fixieApiKey = getEnvVar('FIXIE_API_KEY', false);
-      if (!this.fixieApiKey) {
-        throw new AIJSXError(
-          'You must provide a Fixie API key to access Fixie corpora. Find yours at https://api.fixie.ai/profile.',
-          ErrorCode.MissingFixieAPIKey,
-          'user'
-        );
-      }
+  private readonly fixieApiConfiguration: FixieAPIConfiguration;
+
+  constructor(
+    private readonly corpusId: string,
+    fixieApiConfiguration: FixieAPIConfiguration = DEFAULT_API_CONFIGURATION
+  ) {
+    this.fixieApiConfiguration = fixieApiConfiguration;
+    if (!this.fixieApiConfiguration.authToken) {
+      throw new AIJSXError(
+        `You must provide a Fixie API key to access Fixie corpora. Find yours at ${this.fixieApiConfiguration.url}/profile.`,
+        ErrorCode.MissingFixieAPIKey,
+        'user'
+      );
     }
-    this.fixieApiUrl = getEnvVar('FIXIE_API_URL', false) ?? FixieCorpus.DEFAULT_FIXIE_API_URL;
   }
 
   async search(query: string, params?: { limit?: number }): Promise<ScoredChunk<ChunkMetadata>[]> {
-    const url = `${this.fixieApiUrl}/api/v1/corpora/${this.corpusId}:query`;
+    const url = new URL(`/api/v1/corpora/${this.corpusId}:query`, this.fixieApiConfiguration.url);
     const response = await fetch(url, {
       method: 'POST',
       headers: {
         'Content-Type': 'application/json',
-        Authorization: `Bearer ${this.fixieApiKey}`,
+        Authorization: `Bearer ${this.fixieApiConfiguration.authToken}`,
       },
       body: JSON.stringify({
         corpus_id: this.corpusId,
@@ -663,6 +663,36 @@ export class FixieCorpus<ChunkMetadata extends Jsonifiable = Jsonifiable> implem
       score: result.score,
     }));
   }
+
+  static createTool(corpusId: string, description: string): Tool {
+    return {
+      description,
+      parameters: {
+        query: {
+          description: 'The search query. It will be embedded and used in a vector search against the corpus.',
+          type: 'string',
+          required: true,
+        },
+      },
+      func: async function FixieCorpusQuery({ query }: { query: string }, { getContext }) {
+        const corpus = new FixieCorpus(corpusId, getContext(FixieAPIContext));
+
+        const results = await corpus.search(query);
+
+        /**
+         * Reverse the array so the closest chunk is listed last.
+         *
+         * N.B.: The results will not be sorted by `score` because the Fixie API reranks the chunks.
+         */
+        results.reverse();
+
+        return JSON.stringify({
+          kind: 'docs',
+          results: _.uniqBy(results, (result) => result.chunk.content),
+        });
+      },
+    };
+  }
 }
 
 /** A {@link Corpus} backed by a LangChain {@link VectorStore}. */

packages/ai-jsx/src/batteries/fixie.tsx

@@ -0,0 +1,14 @@
+import { createContext } from '../core/render.js';
+import { getEnvVar } from '../lib/util.js';
+
+export interface FixieAPIConfiguration {
+  url: string;
+  authToken?: string;
+}
+
+export const DEFAULT_API_CONFIGURATION: FixieAPIConfiguration = {
+  url: getEnvVar('FIXIE_API_URL', false) ?? 'https://api.fixie.ai',
+  authToken: getEnvVar('FIXIE_API_KEY', false),
+};
+
+export const FixieAPIContext = createContext<FixieAPIConfiguration>(DEFAULT_API_CONFIGURATION);

packages/ai-jsx/src/batteries/sidekick/platform/conversation.tsx

@@ -12,7 +12,7 @@ import {
   renderToConversation,
   SystemMessage,
 } from '../../../core/conversation.js';
-import { UseToolsProps } from '../../use-tools.js';
+import { ExecuteFunction, UseToolsProps } from '../../use-tools.js';
 
 /**
  * This function defines the shrinking policy. It's activated when the conversation history overflows the context
@@ -73,7 +73,7 @@ export function present(conversationElement: ConversationMessage) {
  * without using tools. For example, this is how the model is able to call `listConversations`, followed by
  * `getConversation`, and then finally write a response.
  */
-export async function getNextConversationStep(
+export function getNextConversationStep(
   messages: ConversationMessage[],
   fullConversation: ConversationMessage[],
   finalSystemMessageBeforeResponse: AI.Node,
@@ -84,15 +84,7 @@ export async function getNextConversationStep(
   switch (lastMessage.type) {
     case 'functionCall': {
       const { name, args } = lastMessage.element.props;
-      try {
-        return <FunctionResponse name={name}>{await tools[name].func(args)}</FunctionResponse>;
-      } catch (e: any) {
-        return (
-          <FunctionResponse failed name={name}>
-            {e.message}
-          </FunctionResponse>
-        );
-      }
+      return <ExecuteFunction func={tools[name].func} name={name} args={args} />;
     }
     case 'functionResponse':
       return (

packages/ai-jsx/src/batteries/sidekick/platform/sidekick.tsx

@@ -1,7 +1,6 @@
 import { present } from './conversation.js';
 import { UseTools } from './use-tools-eject.js';
 import { SidekickSystemMessage } from './system-message.js';
-import _ from 'lodash';
 import { OpenAI } from '../../../lib/openai.js';
 import { UseToolsProps } from '../../use-tools.js';
 import * as AI from '../../../index.js';
@@ -52,31 +51,12 @@ export interface SidekickProps {
   role: string;
 }
 
-function makeObservedTools(tools: UseToolsProps['tools'], logger: AI.ComponentContext['logger']) {
-  return _.mapValues(tools, (tool, toolName) => ({
-    ...tool,
-    func: async (...args: Parameters<typeof tool.func>) => {
-      logger.info({ toolName, args }, 'Calling tool');
-      try {
-        const result = await Promise.resolve(tool.func(...args));
-        logger.info({ toolName, args, result }, 'Got result from tool');
-        return typeof result === 'string' ? result : JSON.stringify(result);
-      } catch (e) {
-        logger.error({ toolName, args, e }, 'Got error calling tool');
-        throw e;
-      }
-    },
-  }));
-}
-
-export function Sidekick(props: SidekickProps, { logger }: AI.ComponentContext) {
-  const observedTools = makeObservedTools(props.tools ?? {}, logger);
-
+export function Sidekick(props: SidekickProps) {
   return (
     <ModelProvider model="gpt-4-32k" modelProvider="openai">
       <ShowConversation present={present}>
         <UseTools
-          tools={observedTools}
+          tools={props.tools ?? {}}
           showSteps
           finalSystemMessageBeforeResponse={props.finalSystemMessageBeforeResponse}
         >

packages/ai-jsx/src/batteries/use-tools.tsx

@@ -5,7 +5,7 @@
  */
 
 import { ChatCompletion, FunctionParameters, FunctionResponse } from '../core/completion.js';
-import { Node, RenderContext } from '../index.js';
+import { Component, ComponentContext, Node, RenderContext } from '../index.js';
 import { Converse, renderToConversation } from '../core/conversation.js';
 
 /**
@@ -23,12 +23,16 @@ export interface Tool {
   parameters: FunctionParameters;
 
   /**
-   * A function to invoke the tool.
+   * A function that invokes the tool.
+   *
+   * @remarks
+   * The function will be treated as an AI.JSX component: the tool parameters
+   * will be passed as fields on the first argument (props) and the function
+   * can return a `string` or any AI.JSX {@link Node}, synchronously or
+   * asynchronously.
    */
-  // Can we use Zod to do better than any[]?
-  func: (
-    ...args: any[]
-  ) => string | number | boolean | null | undefined | Promise<string | number | boolean | null | undefined>;
+  // Can we use Zod to do better than any?
+  func: Component<any>;
 }
 
 /**
@@ -58,6 +62,29 @@ export interface UseToolsProps {
   userData?: string;
 }
 
+/**
+ * Executes a function during rendering and wraps the result in a `<FunctionResponse>`.
+ */
+export async function ExecuteFunction<T>(
+  { name, func, args }: { name: string; func: Component<T>; args: T },
+  { render }: ComponentContext
+) {
+  if (typeof func !== 'function') {
+    return (
+      <FunctionResponse failed name={name}>
+        Error: unknown function {name}
+      </FunctionResponse>
+    );
+  }
+
+  try {
+    const Func = func;
+    return <FunctionResponse name={name}>{await render(<Func {...args} />)}</FunctionResponse>;
+  } catch (e) {
+    return <FunctionResponse failed name={name}>{`${e}`}</FunctionResponse>;
+  }
+}
+
 /**
  * Give a model tools it can use, like a calculator, or ability to call an API.
  *
@@ -108,20 +135,12 @@ export interface UseToolsProps {
 export async function UseTools(props: UseToolsProps, { render }: RenderContext) {
   const converse = (
     <Converse
-      reply={async (messages, fullConversation) => {
+      reply={(messages, fullConversation) => {
         const lastMessage = messages[messages.length - 1];
         switch (lastMessage.type) {
           case 'functionCall': {
             const { name, args } = lastMessage.element.props;
-            try {
-              return <FunctionResponse name={name}>{await props.tools[name].func(args)}</FunctionResponse>;
-            } catch (e: any) {
-              return (
-                <FunctionResponse failed name={name}>
-                  {e.message}
-                </FunctionResponse>
-              );
-            }
+            return <ExecuteFunction func={props.tools[name].func} name={name} args={args} />;
           }
           case 'functionResponse':
           case 'user':

packages/docs/docs/changelog.md

@@ -1,6 +1,13 @@
 # Changelog
 
-## 0.16.0
+## 0.17.0
+
+- Changed `<UseTools>` to allow AI.JSX components to be tools.
+- Added `FixieAPIConfiguration` context.
+- Changed `FixieCorpus` to take a `FixieAPIConfiguration`.
+- Added the `FixieCorpus.createTool` helper to create a tool that consults a Fixie corpus.
+
+## [0.16.0](https://github.com/fixie-ai/ai-jsx/commit/c951b4695c97230016b1ae2763649f67089adf92)
 
 - Updated default URL for `<FixieCorpus>` to `api.fixie.ai`.