112 lines
3.6 KiB
TypeScript
112 lines
3.6 KiB
TypeScript
import { tool } from '@langchain/core/tools';
|
|
import { z } from 'zod';
|
|
import { RunnableConfig } from '@langchain/core/runnables';
|
|
import { withStructuredOutput } from '@/lib/utils/structuredOutput';
|
|
import { PromptTemplate } from '@langchain/core/prompts';
|
|
import { taskBreakdownPrompt } from '@/lib/prompts/taskBreakdown';
|
|
|
|
// Schema for task manager tool input
|
|
const TaskManagerToolSchema = z.object({
|
|
query: z.string().describe('The user query to break down into smaller tasks'),
|
|
context: z
|
|
.string()
|
|
.optional()
|
|
.describe('Additional context about the query or current situation'),
|
|
});
|
|
|
|
// Schema for structured output
|
|
const TaskBreakdownSchema = z.object({
|
|
tasks: z
|
|
.array(z.string())
|
|
.describe(
|
|
'Array of specific, focused tasks broken down from the original query',
|
|
),
|
|
reasoning: z
|
|
.string()
|
|
.describe(
|
|
'Explanation of how and why the query was broken down into these tasks',
|
|
),
|
|
});
|
|
|
|
/**
|
|
* TaskManagerTool - Breaks down complex queries into manageable task lists
|
|
*
|
|
* This tool takes a user query and returns a list of specific, actionable tasks
|
|
* that can help answer the original question. The tasks are returned as natural
|
|
* language instructions that the main agent can follow.
|
|
*/
|
|
export const taskManagerTool = tool(
|
|
async (
|
|
input: z.infer<typeof TaskManagerToolSchema>,
|
|
config?: RunnableConfig,
|
|
): Promise<{ tasks: string[]; reasoning: string }> => {
|
|
try {
|
|
console.log(
|
|
'TaskManagerTool: Starting task breakdown for query:',
|
|
input.query,
|
|
);
|
|
const { query, context = '' } = input;
|
|
|
|
// Get LLM from config
|
|
if (!config?.configurable?.llm) {
|
|
throw new Error('LLM not available in config');
|
|
}
|
|
|
|
const llm = config.configurable.llm;
|
|
|
|
// Create structured LLM for task breakdown
|
|
const structuredLLM = withStructuredOutput(llm, TaskBreakdownSchema, {
|
|
name: 'task_breakdown',
|
|
includeRaw: false,
|
|
});
|
|
|
|
// Create the prompt template
|
|
const template = PromptTemplate.fromTemplate(taskBreakdownPrompt);
|
|
|
|
// Format the prompt with the query and context
|
|
const prompt = await template.format({
|
|
systemInstructions:
|
|
config.configurable?.systemInstructions ||
|
|
'You are a helpful AI assistant.',
|
|
fileContext: context || 'No additional context provided.',
|
|
query: query,
|
|
currentTasks: 0,
|
|
taskHistory: 'No previous tasks.',
|
|
});
|
|
|
|
// Get the task breakdown from the LLM
|
|
const response = await structuredLLM.invoke(prompt, {
|
|
signal: config?.signal,
|
|
});
|
|
|
|
if (!response?.tasks || response.tasks.length === 0) {
|
|
// If no breakdown is needed, return the original query as a single task
|
|
return {
|
|
tasks: [query],
|
|
reasoning:
|
|
'The query is straightforward and does not require breaking down into smaller tasks.',
|
|
};
|
|
}
|
|
|
|
return {
|
|
tasks: response.tasks,
|
|
reasoning: response.reasoning,
|
|
};
|
|
} catch (error) {
|
|
console.error('Error in TaskManagerTool:', error);
|
|
// Fallback: return the original query as a single task
|
|
const errorMessage =
|
|
error instanceof Error ? error.message : 'Unknown error';
|
|
return {
|
|
tasks: [input.query],
|
|
reasoning: `Error occurred during task breakdown: ${errorMessage}. Proceeding with the original query.`,
|
|
};
|
|
}
|
|
},
|
|
{
|
|
name: 'task_manager',
|
|
description:
|
|
'Breaks down complex user queries into a list of specific, manageable tasks that can be executed to answer the original question',
|
|
schema: TaskManagerToolSchema,
|
|
},
|
|
);
|