Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Structure task output #174

Merged
merged 10 commits into from
Dec 21, 2024
Merged

Conversation

anthonydevs17
Copy link
Collaborator

@anthonydevs17 anthonydevs17 commented Dec 17, 2024

Add Structured Output Feature

This PR adds Structured Output support to KaibanJS, allowing users to define and enforce specific output structures for their agent tasks using Zod schemas.

Changes

  • Add outputSchema property to Task class
  • Implement schema validation in ReactChampionAgent
  • Add OUTPUT_SCHEMA_VALIDATION_ERROR state
  • Integrate with workflowLogs for error tracking
  • Add schema validation feedback loop in agent processing

Implementation Details

  • Use Zod for schema validation
  • Add schema validation to agent's output processing
  • Integrate with existing state management system
  • Add validation error handling in agent loop
  • Track validation events in workflowLogs

Testing

  • Add tests for schema validation
  • Add tests for error handling
  • Add tests for agent recovery from validation errors

@darielnoel
Copy link
Contributor

Solves #171

* @param {Object} params.outputSchemaError - The error object for the output schema validation.
* @returns {string} The formatted feedback message.
*/
INVALID_JSON_FOR_OUTPUT_SCHEMA_FEEDBACK: ({
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is better to call this variable this way. -> INVALID_OUTPUT_SCHEMA_FEEDBACK.

outputSchema,
outputSchemaError,
}) => {
const prompt = `You returned an invalid JSON object with following error ${outputSchemaError.toString()}. Please format your answer adhere to this JSON schema ${JSON.stringify(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Love this prompt 🤗

src/index.js Outdated
@@ -116,6 +117,12 @@ class Task {
this.interpolatedTaskDescription = null;
this.feedbackHistory = []; // Initialize feedbackHistory as an empty array
this.externalValidationRequired = externalValidationRequired;
this.outputSchema = outputSchema; // Zod Schema
this.expectedOutput = outputSchema
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We cannot overwrite the expectedOutput here... both variables serves different purposes:

expectedOutput and outputSchema serve different purposes:

  • expectedOutput: A human-readable description used in prompts to tell the AI agent what output format you want (used for instruction)
  • outputSchema: A Zod schema that validates the actual output structure (used for validation)

Example:

const task = new Task({
  description: "Extract article metadata",
  expectedOutput: "Get the article's title and list of tags", // For the AI to understand
  outputSchema: z.object({                                    // For validation
    title: z.string(),
    tags: z.array(z.string())
  })
});

@@ -201,6 +201,13 @@ class ReactChampionAgent extends BaseAgent {
output: thinkingResult,
});
break;
case AGENT_STATUS_enum.FINAL_ANSWER_WRONG_STRUCTURED:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's use OUTPUT_SCHEMA_VALIDATION_ERROR as the name for this agent status.

@@ -518,13 +544,34 @@ class ReactChampionAgent extends BaseAgent {
});
return feedbackMessage;
}
handleIssuesParsingSchemaOutput({ agent, task, output }) {
const jSONPArsingError = new Error(
'Received an invalid JSON object from the LLM. JSON object does not match the expected schema.',
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Update error message to: "The output does not match the expected schema structure"

'Received an invalid JSON object from the LLM. JSON object does not match the expected schema.',
output.parsedLLMOutput.outputSchemaErrors
);
agent.store.getState().handleAgentIssuesParsingLLMOutput({
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here we should implement a different handler in the store. Cause they are two separated use cases.

@darielnoel darielnoel merged commit 2d85d7e into kaiban-ai:main Dec 21, 2024
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants