Remove magic boolean value (e.g. messages: true and context: true)

statelyai · Jun 29, 2024 · 13ec4ca · 13ec4ca
1 parent 094bb8a
commit 13ec4ca
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 26 deletions.
diff --git a/examples/example.ts b/examples/example.ts
@@ -32,11 +32,11 @@ const machine = setup({
       invoke: [
         {
           src: 'summarizer',
-          input: {
-            context: true,
+          input: (x) => ({
+            context: x.context,
             prompt:
               'Summarize the patient visit in a single sentence. The summary should be in English.',
-          },
+          }),
           onDone: {
             actions: assign({
               englishSummary: ({ event }) => event.output.text,
@@ -45,11 +45,11 @@ const machine = setup({
         },
         {
           src: 'summarizer',
-          input: {
-            context: true,
+          input: (x) => ({
+            context: x.context,
             prompt:
               'Summarize the patient visit in a single sentence. The summary should be in Spanish.',
-          },
+          }),
           onDone: {
             actions: assign({
               spanishSummary: ({ event }) => event.output.text,

diff --git a/examples/wiki.ts b/examples/wiki.ts
@@ -20,7 +20,7 @@ async function main() {
   console.log(response1.text);
 
   const response2 = await agent.generateText({
-    messages: true,
+    messages: (x) => x.select((ctx) => ctx.messages),
     prompt: 'What about the first one?',
   });
 

diff --git a/src/text.ts b/src/text.ts
@@ -35,9 +35,7 @@ async function getMessages(
   options: AgentStreamTextOptions
 ): Promise<CoreMessage[]> {
   let messages: CoreMessage[] = [];
-  if (options.messages === true) {
-    messages = agent.select((s) => s.messages);
-  } else if (typeof options.messages === 'function') {
+  if (typeof options.messages === 'function') {
     messages = await options.messages(agent);
   } else if (options.messages) {
     messages = options.messages;
@@ -164,15 +162,10 @@ export function fromTextStream<T extends Agent<any>>(
 ): ObservableActorLogic<
   { textDelta: string },
   Omit<AgentStreamTextOptions, 'context'> & {
-    context?: AgentStreamTextOptions['context'] | boolean;
+    context?: AgentStreamTextOptions['context'];
   }
 > {
-  return fromObservable(({ input, self }) => {
-    const context =
-      input.context === true
-        ? (self._parent?.getSnapshot() as AnyMachineSnapshot).context
-        : input.context;
-
+  return fromObservable(({ input }) => {
     const observers = new Set<Observer<{ textDelta: string }>>();
 
     // TODO: check if messages was provided instead
@@ -181,7 +174,7 @@ export function fromTextStream<T extends Agent<any>>(
       const result = await agentStreamText(agent, {
         ...defaultOptions,
         ...input,
-        context,
+        context: input.context,
       });
 
       for await (const part of result.fullStream) {
@@ -214,18 +207,14 @@ export function fromText<T extends Agent<any>>(
 ): PromiseActorLogic<
   GenerateTextResult<Record<string, CoreTool<any, any>>>,
   Omit<AgentGenerateTextOptions, 'context'> & {
-    context?: AgentGenerateTextOptions['context'] | boolean;
+    context?: AgentGenerateTextOptions['context'];
   }
 > {
-  return fromPromise(async ({ input, self }) => {
-    const context =
-      input.context === true
-        ? (self._parent?.getSnapshot() as AnyMachineSnapshot).context
-        : input.context;
+  return fromPromise(async ({ input }) => {
     return await agentGenerateText(agent, {
       ...input,
       ...defaultOptions,
-      context,
+      context: input.context,
     });
   });
 }
diff --git a/src/types.ts b/src/types.ts
@@ -325,7 +325,7 @@ export interface CommonTextOptions {
   prompt: FromAgent<string>;
   model?: LanguageModel;
   context?: Record<string, any>;
-  messages?: FromAgent<CoreMessage[]> | true;
+  messages?: FromAgent<CoreMessage[]>;
   template?: PromptTemplate<any>;
 }