import { ThoughtData } from '../models/interfaces.js';
import { thoughtDataSchema } from '../schemas/trace.js';
import { ZodError } from 'zod';
import chalk from 'chalk';
export class TraceServer {
private thoughtHistory: ThoughtData[] = [];
private branches: Record<string, ThoughtData[]> = {};
private validateThoughtData(input: unknown): ThoughtData {
try {
return thoughtDataSchema.parse(input);
} catch (error) {
if (error instanceof ZodError) {
const errorMessages = error.errors.map(err => `${err.path.join('.')}: ${err.message}`).join(', ');
throw new Error(`Invalid thought data: ${errorMessages}`);
}
throw error;
}
}
private formatThoughtOutput(data: ThoughtData): string {
const {
thought,
thoughtNumber,
totalThoughts,
nextThoughtNeeded,
isRevision,
revisesThought,
branchFromThought,
branchId,
needsMoreThoughts
} = data;
let output = '';
// Add header based on thought type
if (isRevision && revisesThought) {
output += `\n${chalk.bold.yellow(`Thought ${thoughtNumber}/${totalThoughts} (Revising Thought ${revisesThought}):`)} `;
} else if (branchFromThought && branchId) {
output += `\n${chalk.bold.magenta(`Thought ${thoughtNumber}/${totalThoughts} (Branch "${branchId}" from Thought ${branchFromThought}):`)} `;
} else {
output += `\n${chalk.bold.blue(`Thought ${thoughtNumber}/${totalThoughts}:`)} `;
}
// Add the thought content
output += `${thought}\n`;
// Add footer with status
if (nextThoughtNeeded) {
output += chalk.green(`\nContinuing to next thought...\n`);
if (needsMoreThoughts) {
output += chalk.yellow(`More thoughts needed than initially estimated.\n`);
}
} else {
output += chalk.cyan(`\nThinking process complete.\n`);
}
return output;
}
private storeThought(thought: ThoughtData): void {
// If this is a branch, store in the appropriate branch collection
if (thought.branchId) {
if (!this.branches[thought.branchId]) {
this.branches[thought.branchId] = [];
}
this.branches[thought.branchId].push(thought);
} else {
// Otherwise store in main thought history
this.thoughtHistory.push(thought);
}
}
public processThought(input: unknown): ThoughtData {
const validatedInput = this.validateThoughtData(input);
// Store the thought for future reference
this.storeThought(validatedInput);
// Log formatted output to console
const formattedOutput = this.formatThoughtOutput(validatedInput);
console.error(formattedOutput);
return validatedInput;
}
}
