Optimizing Agent Reliability: Debugging Trajectories and Prompt Engineering

// agent.ts
 
interface Tool {
    name: string;
    description: string;
    execute: (input: string) => Promise<string>;
}
 
interface AgentState {
    thought: string;
    action?: { tool: string; input: string };
    observation?: string;
    history: { role: 'user' | 'assistant' | 'system'; content: string }[];
}
 
class LeanAgent {
    private tools: Record<string, Tool>;
    private llmClient: any; // Raw OpenAI, Anthropic, or Agno client
    private systemPrompt: string;
 
    constructor(tools: Tool[], systemPrompt: string, llmClient: any) {
        this.tools = tools.reduce((acc, tool) => ({ ...acc, [tool.name]: tool }), {});
        this.llmClient = llmClient;
        this.systemPrompt = systemPrompt;
    }
 
    private formatToolsForPrompt(): string {
        return Object.values(this.tools)
            .map(tool => {
                let example = '';
                // Add specific few-shot examples directly in tool description for clarity
                if (tool.name === 'search') {
                    example = `Example: search("latest AI research papers")`;
                } else if (tool.name === 'calculator') {
                    example = `Example: calculator("2+2*3")`;
                }
                return `Tool: ${tool.name}(input: string) -> string\nDescription: ${tool.description}\n${example}`;
            })
            .join('\n\n');
    }
 
    private async callLLM(messages: { role: 'user' | 'assistant' | 'system'; content: string }[]): Promise<string> {
        // Direct call to LLM API (e.g., OpenAI chat completions)
        // Add robust retry logic, timeout, and error handling here
        try {
            const response = await this.llmClient.chat.completions.create({
                model: "gpt-4o", // Or your preferred fast, capable model
                messages: messages,
                temperature: 0.1, // Keep it low for reliability
                max_tokens: 512,
            });
            return response.choices[0].message.content || '';
        } catch (error) {
            console.error("LLM API error:", error);
            // Implement circuit breaker or specific fallback
            throw new Error("Failed to get LLM response");
        }
    }
 
    public async run(initialQuery: string, maxIterations: number = 5): Promise<string> {
        let state: AgentState = {
            thought: '',
            history: [{ role: 'system', content: this.systemPrompt + '\n\n' + this.formatToolsForPrompt() }],
        };
        state.history.push({ role: 'user', content: initialQuery });
 
        for (let i = 0; i < maxIterations; i++) {
            console.log(`--- Agent Iteration ${i + 1} ---`);
            const llmResponse = await this.callLLM(state.history);
 
            // This is where custom parsing comes in
            const parsedOutput = this.parseAgentOutput(llmResponse);
 
            if (parsedOutput && 'answer' in parsedOutput) {
                console.log(`FINAL ANSWER: ${parsedOutput.answer}`);
                return parsedOutput.answer;
            } else if (parsedOutput && 'tool' in parsedOutput) {
                const { tool, input } = parsedOutput;
                console.log(`Action: ${tool}, Input: ${input}`);
 
                const toolExecutor = this.tools[tool];
                if (!toolExecutor) {
                    const observation = `Error: Tool '${tool}' not found. Available tools: ${Object.keys(this.tools).join(', ')}`;
                    state.history.push({ role: 'assistant', content: llmResponse });
                    state.history.push({ role: 'user', content: `Observation: ${observation}` });
                } else {
                    try {
                        const observation = await toolExecutor.execute(input);
                        state.history.push({ role: 'assistant', content: llmResponse });
                        state.history.push({ role: 'user', content: `Observation: ${observation}` });
                    } catch (error) {
                        const observation = `Tool '${tool}' failed with error: ${error instanceof Error ? error.message : String(error)}`;
                        state.history.push({ role: 'assistant', content: llmResponse });
                        state.history.push({ role: 'user', content: `Observation: ${observation}` });
                    }
                }
            } else {
                console.warn("LLM provided malformed or unparseable output. Retrying with explicit instruction.");
                state.history.push({ role: 'assistant', content: llmResponse });
                state.history.push({ role: 'user', content: "Your previous output was malformed. Please ensure you either use 'Action: [tool]\nAction Input: [input]' or 'FINAL ANSWER: [answer]' format. Do not use any other format." });
            }
        }
        return "Agent failed to find a final answer within the maximum iterations.";
    }
 
    // Custom output parser (implementation below)
    private parseAgentOutput(llmOutput: string): { tool: string; input: string } | { answer: string } | null {
        // ... (implementation detailed in section 5)
        return null; // Placeholder
    }
}

// Minimal LangSmith integration for observability
import { Client, RunTree } from "langsmith";
 
// Assuming LANGCHAIN_API_KEY and LANGCHAIN_TRACING_V2 are set
const langsmithClient = new Client();
 
async function wrapLLMCallWithTrace(
    llmClient: any,
    messages: { role: 'user' | 'assistant' | 'system'; content: string }[],
    parentRun?: RunTree
): Promise<{ output: string; run: RunTree }> {
    const run = new RunTree({
        name: "LLM_Call",
        run_type: "llm",
        inputs: { messages },
        parent_run: parentRun,
        // tags: ["agent-trace"] // Add custom tags for filtering
    });
    try {
        await run.post();
        const response = await llmClient.chat.completions.create({ /* ... */ });
        const output = response.choices[0].message.content || '';
        run.outputs = { output };
        await run.end();
        return { output, run };
    } catch (error) {
        run.error = String(error);
        await run.end();
        throw error;
    }
}
 
// In the Agent's run method:
// let currentRun: RunTree | undefined; // Pass this through iterations
// ...
// const { output: llmResponse, run: llmRun } = await wrapLLMCallWithTrace(this.llmClient, state.history, currentRun);
// currentRun = llmRun; // Update current run for next step

const systemPrompt = `
You are a highly efficient and accurate problem-solving agent.
Your goal is to precisely answer the user's query by breaking it down into steps, using available tools, and providing a concise FINAL ANSWER.
 
Strictly follow this thought-action-observation loop:
1.  **Thought**: You must always first reflect on the current state, what you need to do next, and which tool to use.
2.  **Action**: If you need to use a tool, output 'Action: [tool_name]\nAction Input: [tool_input]'.
3.  **Observation**: This will be provided by the system after your action.
4.  **FINAL ANSWER**: Once you have fully solved the user's request and have the complete answer, output 'FINAL ANSWER: [your_answer]'.
 
Do NOT elaborate or provide conversational text outside of your 'Thought' or 'FINAL ANSWER'.
You MUST ONLY use the tools provided. DO NOT hallucinate tool names or arguments.
If you get stuck or receive an error, try to recover or state a FINAL ANSWER if appropriate.
`;

// Inside formatToolsForPrompt() method or similar utility
const tools: Tool[] = [
    {
        name: 'search',
        description: 'Searches the web for factual information. Use this when you need current data or to verify facts.',
        execute: async (query: string) => `Simulated web search for "${query}" resulted in "LLMs are complex neural networks."`,
    },
    {
        name: 'calculator',
        description: 'Evaluates mathematical expressions. Input must be a valid arithmetic string.',
        execute: async (expression: string) => `Simulated calculation of "${expression}" resulted in "10"`,
    },
    // ... more tools
];
 
// Example of how they render in the prompt:
// Tool: search(input: string) -> string
// Description: Searches the web for factual information. Use this when you need current data or to verify facts.
// Example: search("latest AI research papers")
//
// Tool: calculator(input: string) -> string
// Description: Evaluates mathematical expressions. Input must be a valid arithmetic string.
// Example: calculator("2+2*3")

// Inside the LeanAgent class
private parseAgentOutput(llmOutput: string): { tool: string; input: string } | { answer: string } | null {
    llmOutput = llmOutput.trim(); // Always trim whitespace
 
    // Attempt to parse Final Answer first
    const finalAnswerMatch = llmOutput.match(/^FINAL ANSWER:\s*(.*)/is);
    if (finalAnswerMatch) {
        return { answer: finalAnswerMatch[1].trim() };
    }
 
    // Attempt to parse Action
    // Regex is designed to be robust against slight variations in newlines/spacing
    const actionMatch = llmOutput.match(/Action:\s*([a-zA-Z0-9_]+)\s*\nAction Input:\s*(.*)/is);
    if (actionMatch) {
        return {
            tool: actionMatch[1].trim(),
            input: actionMatch[2].trim(),
        };
    }
 
    // Fallback: If neither matches, try to salvage
    console.warn("LLM output malformed, attempting heuristic recovery:", llmOutput);
 
    // Heuristic 1: If it starts with "Thought:" and then has "Action:"
    const thoughtActionMatch = llmOutput.match(/Thought:\s*.*?\s*Action:\s*([a-zA-Z0-9_]+)\s*\nAction Input:\s*(.*)/is);
    if (thoughtActionMatch) {
         console.warn("Recovered from Thought-Action format.");
         return {
            tool: thoughtActionMatch[1].trim(),
            input: thoughtActionMatch[2].trim(),
        };
    }
 
    // Heuristic 2: Simple case where Action/Input might be on same line or separated differently
    const simpleActionMatch = llmOutput.match(/Action:\s*([a-zA-Z0-9_]+)\s*Input:\s*(.*)/is);
    if (simpleActionMatch) {
        console.warn("Recovered from simple Action-Input format.");
        return {
            tool: simpleActionMatch[1].trim(),
            input: simpleActionMatch[2].trim(),
        };
    }
 
    // If all recovery attempts fail
    console.error("Failed to parse LLM output after all attempts:", llmOutput);
    return null; // Signal a parsing failure, which the agent loop should handle
}