maxSteps config for chains and agents

docs/guides/prompt-manager/configuration.mdx

@@ -13,7 +13,8 @@ This section is enclosed between three dashes (`---`) and written in YAML format
 
 ```yaml
 ---
-model: gpt-4o
+provider: OpenAI
+model: gpt-4o-mini
 temperature: 0.6
 top_p: 0.9
 ---
@@ -76,6 +77,13 @@ It is recommended to set either temperature or topP, but not both.
   [Learn more about caching in Latitude](/guides/prompt-manager/cache)
 </Note>
 
+#### **`maxSteps`**
+
+Maximum number of steps to execute.
+
+When using [Chains](/promptl/advanced/chains), you can configure the maximum number of steps to execute. Reaching this limit will stop the execution of the chain and return an error.
+Defaults to 20.
+
 #### **`maxTokens`**
 
 Maximum number of tokens to generate.

docs/promptl/advanced/chains.mdx

@@ -8,6 +8,7 @@ description: Chains and Steps are used to create multi-step prompts that can int
 Chains in PromptL allow you to break complex workflows into smaller, manageable steps. Each step generates a response, which can then be used in subsequent steps. This approach improves the model's ability to perform complex tasks and provides greater control over dynamic conversations.
 
 With Chains, you can:
+
 - Process tasks incrementally to guide the model step-by-step.
 - Store and reuse intermediate results dynamically.
 - Customize step-specific configurations for more efficient execution.
@@ -20,6 +21,7 @@ With Chains, you can:
 Define steps in your prompt using the `<step>` tag. The engine pauses after each step, waits for the model's response, and adds it as an assistant message before continuing.
 
 ### Basic Syntax
+
 ```xml
 <step>
   Analyze the following text and identify the main idea:
@@ -35,7 +37,9 @@ Define steps in your prompt using the `<step>` tag. The engine pauses after each
 ```
 
 ### Step with Custom Configuration
+
 Override the default configuration by adding attributes to the `<step>` tag:
+
 ```tsx
 <step model="gpt-3.5-turbo" temperature={{0.5}}>
   Rewrite the following paragraph in simpler language:
@@ -51,6 +55,7 @@ Override the default configuration by adding attributes to the `<step>` tag:
 ## Advanced Features
 
 ### Storing Step Responses
+
 You can store the text response of a step in a variable using the `as` attribute. This allows you to reuse the response in later steps or logic blocks.
 
 ```xml
@@ -67,6 +72,7 @@ You can store the text response of a step in a variable using the `as` attribute
 ```
 
 #### Parse Step Responses as JSON
+
 The response of a step will be automatically parsed as JSON if the JSON output schema is defined.
 
 ```xml
@@ -85,6 +91,7 @@ The response of a step will be automatically parsed as JSON if the JSON output s
 [Learn more about JSON Output](/guides/prompt-manager/json-output).
 
 ### Storing Raw Messages
+
 Some providers will return additional metadata along with the response. To store the entire message object, instead of just the text response (e.g., role, content, and additional metadata), use the `raw` attribute:
 
 ```xml
@@ -100,6 +107,7 @@ The `content` attribute will always be defined as an array of content objects, w
 ---
 
 ### Isolating Steps
+
 Use the `isolated` attribute to prevent a step from inheriting context from previous steps. This can reduce unnecessary costs or confusion for the model.
 
 ```xml
@@ -124,10 +132,34 @@ In this example, the final step does not need to conside the full texts used in
 
 ---
 
+### Limiting the number of steps
+
+<Note>This feature is only available on the Latitude platform.</Note>
+
+You can limit the number of steps executed in a chain by setting a `maxSteps` attribute on the main configuration section. This can help prevent infinite loops or excessive processing in long chains when creating complex workflows with steps within loops.
+
+```xml
+---
+maxSteps: 5
+---
+
+{{ for item in list }}
+  <step>
+    ...
+  </step>
+{{ endfor }}
+```
+
+Read more about this configuration in the [Latitude Prompt Configuration](/guides/prompt-manager/configuration#maxsteps) guide.
+
+---
+
 ## Real-World Use Cases
 
 ### Multi-Step Workflow
+
 Chains are ideal for breaking down tasks like:
+
 1. Analyzing data.
 2. Generating intermediate results.
 3. Combining results for a final output.
@@ -146,7 +178,9 @@ Chains are ideal for breaking down tasks like:
 ```
 
 ### Decision Trees
+
 Use logic to adapt workflows based on intermediate results:
+
 ```xml
 <step as="classification">
   Classify this document into one of the following categories: A, B, or C.
@@ -176,6 +210,7 @@ To execute chains, use the `Chain` class. The chain evaluates the prompt step-by
 To run a step, execute the `step` method of the chain instance. The first time `step` is called, it should not include any arguments. Subsequent calls must always pass the model response message from the previous step.
 
 ### Example: Using the Chain Class
+
 ```javascript
 import { Chain } from 'promptl-ai';
 import OpenAI from 'openai';

packages/constants/src/errors.ts

@@ -24,6 +24,7 @@ export enum RunErrorCodes {
   EvaluationRunMissingWorkspaceError = 'ev_run_missing_workspace_error',
   EvaluationRunUnsupportedResultTypeError = 'ev_run_unsupported_result_type_error',
   EvaluationRunResponseJsonFormatError = 'ev_run_response_json_format_error',
+  MaxStepCountExceededError = 'max_step_count_exceeded_error',
 }
 // NOTE: If you add a new error code, please add it to the pg enum in models/runErrors.ts
 

packages/core/src/constants.ts

@@ -448,3 +448,6 @@ export type DraftChange = {
 }
 
 export const AGENT_RETURN_TOOL_NAME = 'agent_finish_task'
+export const MAX_STEPS_CONFIG_NAME = 'maxSteps'
+export const DEFAULT_MAX_STEPS = 20
+export const ABSOLUTE_MAX_STEPS = 150

packages/core/src/helpers.ts

@@ -1,5 +1,5 @@
 import { z } from 'zod'
-import { CsvData, ParameterType } from './constants'
+import { CsvData, MAX_STEPS_CONFIG_NAME, ParameterType } from './constants'
 import { ProviderApiKey, ProviderLogDto } from './schema/types'
 
 import {
@@ -60,6 +60,7 @@ export function promptConfigSchema({
         }),
       )
       .optional(),
+    [MAX_STEPS_CONFIG_NAME]: z.number().min(1).optional(),
   })
 }
 

packages/core/src/schema/models/runErrors.ts

@@ -20,6 +20,7 @@ export const errorCodeEnum = latitudeSchema.enum('run_error_code_enum', [
   RunErrorCodes.EvaluationRunUnsupportedResultTypeError,
   RunErrorCodes.EvaluationRunResponseJsonFormatError,
   RunErrorCodes.DefaultProviderInvalidModel,
+  RunErrorCodes.MaxStepCountExceededError,
 ])
 
 export const runErrorEntities = latitudeSchema.enum('run_error_entity_enum', [

packages/core/src/services/chains/agents/run.ts

@@ -5,13 +5,19 @@ import {
   ChainResponse,
   CachedApiKeys,
   createChainRunError,
+  stepLimitExceededErrorMessage,
 } from '../run'
 import { generateUUIDIdentifier } from '../../../lib/generateUUID'
 import {
   LatitudeChainCompleteEventData,
   LogSources,
 } from '@latitude-data/constants'
-import { ErrorableEntity } from '../../../constants'
+import {
+  MAX_STEPS_CONFIG_NAME,
+  ErrorableEntity,
+  DEFAULT_MAX_STEPS,
+  ABSOLUTE_MAX_STEPS,
+} from '../../../constants'
 import {
   ChainEvent,
   ChainEventTypes,
@@ -80,6 +86,7 @@ export async function runAgent<T extends boolean, C extends SomeChain>({
   })
 
   let conversation: Conversation
+  let stepCount = 0
 
   const chainStartTime = Date.now()
   const stream = new ReadableStream<ChainEvent>({
@@ -100,6 +107,16 @@ export async function runAgent<T extends boolean, C extends SomeChain>({
             } else {
               // Forward all events that are not the ChainComplete event
               controller.enqueue(value)
+
+              if (value.data.type === ChainEventTypes.StepComplete) {
+                stepCount++
+              }
+            }
+
+            if (value.data.type === ChainEventTypes.Error) {
+              return responseResolve(
+                Result.error(value.data.error as ChainError<RunErrorCodes>),
+              )
             }
           }
 
@@ -114,6 +131,7 @@ export async function runAgent<T extends boolean, C extends SomeChain>({
             controller,
             errorableUuid,
             errorableType,
+            stepCount,
             previousCount: conversation.messages.length - 1,
           })
             .then((okResponse) => {
@@ -156,6 +174,7 @@ async function runAgentStep({
   previousResponse,
   errorableUuid,
   errorableType,
+  stepCount,
 }: {
   workspace: Workspace
   source: LogSources
@@ -166,6 +185,7 @@ async function runAgentStep({
   errorableType?: ErrorableEntity
   previousCount?: number
   previousResponse?: ChainStepResponse<StreamType>
+  stepCount: number
 }) {
   const prevText = previousResponse?.text
   const streamConsumer = new ChainStreamConsumer({
@@ -191,6 +211,18 @@ async function runAgentStep({
       return previousResponse
     }
 
+    const maxSteps = Math.min(
+      (conversation.config[MAX_STEPS_CONFIG_NAME] as number | undefined) ??
+        DEFAULT_MAX_STEPS,
+      ABSOLUTE_MAX_STEPS,
+    )
+    if (maxSteps && stepCount >= maxSteps) {
+      throw new ChainError({
+        message: stepLimitExceededErrorMessage(maxSteps),
+        code: RunErrorCodes.MaxStepCountExceededError,
+      })
+    }
+
     const { messageCount, stepStartTime } = streamConsumer.setup(step)
 
     const cachedResponse = await getCachedResponse({
@@ -248,6 +280,7 @@ async function runAgentStep({
           controller,
           previousCount: messageCount,
           previousResponse: response,
+          stepCount: stepCount + 1,
         })
       }
     }
@@ -323,6 +356,7 @@ async function runAgentStep({
       controller,
       previousCount: messageCount,
       previousResponse: response,
+      stepCount: stepCount + 1,
     })
   } catch (e: unknown) {
     const error = streamConsumer.chainError(e)

packages/core/src/services/chains/run.ts

@@ -4,10 +4,13 @@ import { Chain as PromptlChain } from 'promptl-ai'
 
 import { ProviderApiKey, Workspace } from '../../browser'
 import {
+  ABSOLUTE_MAX_STEPS,
   ChainEvent,
   ChainStepResponse,
+  DEFAULT_MAX_STEPS,
   ErrorableEntity,
   LogSources,
+  MAX_STEPS_CONFIG_NAME,
   StreamType,
 } from '../../constants'
 import { Result, TypedResult } from '../../lib'
@@ -29,6 +32,9 @@ import {
 export type CachedApiKeys = Map<string, ProviderApiKey>
 export type SomeChain = LegacyChain | PromptlChain
 
+export const stepLimitExceededErrorMessage = (maxSteps: number) =>
+  `Limit of ${maxSteps} steps exceeded. Configure the '${MAX_STEPS_CONFIG_NAME}' setting in your prompt configuration to allow for more steps.`
+
 export async function createChainRunError({
   error,
   errorableUuid,
@@ -154,6 +160,7 @@ async function runStep({
   errorableType,
   configOverrides,
   removeSchema,
+  stepCount = 0,
 }: {
   workspace: Workspace
   source: LogSources
@@ -167,6 +174,7 @@ async function runStep({
   previousResponse?: ChainStepResponse<StreamType>
   configOverrides?: ConfigOverrides
   removeSchema?: boolean
+  stepCount?: number
 }) {
   const prevText = previousResponse?.text
   const streamConsumer = new ChainStreamConsumer({
@@ -195,6 +203,22 @@ async function runStep({
       return previousResponse!
     }
 
+    const maxSteps = Math.min(
+      (step.conversation.config[MAX_STEPS_CONFIG_NAME] as number | undefined) ??
+        DEFAULT_MAX_STEPS,
+      ABSOLUTE_MAX_STEPS,
+    )
+    const exceededMaxSteps =
+      chain instanceof PromptlChain
+        ? stepCount >= maxSteps
+        : stepCount > maxSteps
+    if (exceededMaxSteps) {
+      throw new ChainError({
+        message: stepLimitExceededErrorMessage(maxSteps),
+        code: RunErrorCodes.MaxStepCountExceededError,
+      })
+    }
+
     const { messageCount, stepStartTime } = streamConsumer.setup(step)
 
     const cachedResponse = await getCachedResponse({
@@ -248,6 +272,7 @@ async function runStep({
           previousCount: messageCount,
           previousResponse: response,
           configOverrides,
+          stepCount: ++stepCount,
         })
       }
     }
@@ -327,6 +352,7 @@ async function runStep({
         previousCount: messageCount,
         previousResponse: response,
         configOverrides,
+        stepCount: ++stepCount,
       })
     }
   } catch (e: unknown) {