Chapter 10: Learning

type Task = { type: string; input: string };
type Result = { output: string; score: number };

export class LearningWrapper {
  // System-level state: things that can change
  @field promptLibrary: Map<string, string> = new Map();
  @field successfulExamples: Task[] = [];
  @field routingHints: Map<string, string[]> = new Map();

  constructor(private readonly llm: LLMClient) {}

  async handle(task: Task): Promise<Result> {
    const prompt = this.buildPrompt(task);
    const { output } = await this.llm.complete(prompt);

    const score = await this.evaluate(task, output);

    // This is where learning actually happens
    await this.updateScaffolding(task, output, score);

    return { output, score };
  }

  private buildPrompt(task: Task): string {
    const base = this.promptLibrary.get(task.type) ?? "You are a helpful assistant.";
    const hints = this.routingHints.get(task.type) ?? [];
    return [
      base,
      hints.length ? `Guidelines:\n- ${hints.join("\n- ")}` : "",
      `Task:\n${task.input}`
    ].join("\n\n");
  }

  private async evaluate(task: Task, output: string): Promise<number> {
    // Chapter 8/9: any evaluation mechanism you like
    return await externalJudge(task, output);
  }

  private async updateScaffolding(task: Task, output: string, score: number) {
    if (score > 0.9) {
      // Treat this as an example of “what good looks like”
      this.successfulExamples.push(task);
    }

    if (score < 0.5) {
      // Ask the model how to avoid this kind of failure next time
      const hint = await this.llm.complete(
        `We got a low score on this ${task.type} task.

Task:
${task.input}

Bad output:
${output}

In one concise bullet, state a concrete guideline that would reduce this failure mode in the future.`
      );

      const existing = this.routingHints.get(task.type) ?? [];
      this.routingHints.set(task.type, [...existing, hint.output.trim()]);
    }
  }
}

@field rawLog: { task: Task; output: string; score: number }[] = [];

async afterTask(task: Task, output: string, score: number) {
  this.rawLog.push({ task, output, score });
}

async buildContextFor(task: Task): Promise<string> {
  const similar = await this.retrieveSimilarTasks(task);

  const examples = similar
    .slice(0, 3)
    .map(
      (rec) =>
        `Task:\n${rec.task.input}\n\nOutput:\n${rec.output}\n\nScore: ${rec.score.toFixed(2)}`
    )
    .join("\n\n---\n\n");

  return `You are solving a new task. Here are some past tasks and outputs:

${examples}

Now solve this task:\n${task.input}`;
}

@field lessons: string[] = [];

async afterTaskWithFeedback(task: Task, output: string, score: number) {
  // Memory: keep the raw record if you want
  // ...

  // Learning: extract a portable lesson
  const lesson = await this.llm.complete(
    `We just handled a task of type "${task.type}".

Task:
${task.input}

Output:
${output}

Score: ${score}

In one sentence, describe a general rule that would improve future handling of similar tasks. Focus on *general* patterns, not this specific example.`
  );

  const text = lesson.output.trim();
  if (text && !this.lessons.includes(text)) {
    this.lessons.push(text);
  }
}

async buildContextFor(task: Task): Promise<string> {
  const relevantLessons = this.lessons.filter((l) => l.includes(task.type)).slice(0, 5);

  return [
    `You are handling a "${task.type}" task.`,
    relevantLessons.length
      ? `Here are lessons learned from past similar tasks:\n- ${relevantLessons.join("\n- ")}`
      : "",
    `Now handle this task:\n${task.input}`
  ]
    .filter(Boolean)
    .join("\n\n");
}

type Example = {
  id: string;
  taskType: string;
  input: string;
  output: string;
  score: number;
  timestamp: number;
};

export class ExampleManager {
  @field examples: Example[] = [];

  async considerForStorage(task: Task, output: string, score: number) {
    if (score < 0.9) return; // Only store strong successes

    this.examples.push({
      id: crypto.randomUUID(),
      taskType: task.type,
      input: task.input,
      output,
      score,
      timestamp: Date.now()
    });
  }

  async selectForPrompt(task: Task, maxExamples = 3): Promise<Example[]> {
    // Start from candidates of the same type
    const candidates = this.examples.filter((e) => e.taskType === task.type);

    // Rank by mixed objective: similarity, recency, quality
    const scored = await Promise.all(
      candidates.map(async (ex) => {
        const similarity = await this.semanticSimilarity(task.input, ex.input);
        const recency = 1 - this.agePenalty(ex.timestamp);
        const quality = ex.score;
        return { ex, score: 0.6 * similarity + 0.2 * recency + 0.2 * quality };
      })
    );

    return scored
      .sort((a, b) => b.score - a.score)
      .slice(0, maxExamples)
      .map((s) => s.ex);
  }

  private agePenalty(timestamp: number): number {
    const days = (Date.now() - timestamp) / (1000 * 60 * 60 * 24);
    return Math.min(1, days / 30); // Anything older than 30 days gets full penalty
  }

  private async semanticSimilarity(a: string, b: string): Promise<number> {
    // Use embeddings or a classifier; 0–1 similarity
    return await embeddingSimilarity(a, b);
  }
}

async summarizeAndCull(taskType: string) {
  const ofType = this.examples.filter((e) => e.taskType === taskType);

  const summary = await this.llm.complete(
    `You are managing a library of examples for task type "${taskType}".

Here are some examples:
${ofType
  .slice(0, 20)
  .map(
    (e, i) => `Example ${i + 1}:
Input: ${e.input}
Output: ${e.output}
Score: ${e.score}
`
  )
  .join("\n")}

1. Group these examples into 3–5 distinct patterns.
2. For each pattern, write a short description and pick *one* example that best represents it.
Return JSON with keys "patterns" (array of {description, representativeExampleIndex}).`
  );

  const patterns = JSON.parse(summary.output).patterns;

  // Keep only the representatives; mark others as archival
  const keepIds = new Set(
    patterns.map((p: { representativeExampleIndex: number }) => ofType[p.representativeExampleIndex].id)
  );

  this.examples = this.examples.filter(
    (e) => e.taskType !== taskType || keepIds.has(e.id)
  );
}

async beforeTask(task: Task): Promise<AugmentedTask> {
  // Learning applied: prompts, examples, knowledge, routing
}

async afterTask(task: Task, result: Result): Promise<void> {
  // Feedback captured: evaluation, critique, metadata
}

@schedule("0 2 * * *")
async periodicReflection(): Promise<void> {
  // Patterns extracted; wrapper updated
}

// System-level: collect structured training signals
type TrainingExample = {
  input: string;
  idealOutput: string;
};

@field trainingBuffer: TrainingExample[] = [];

async afterTask(task: Task, result: Result) {
  const { output, score } = result;

  // If user corrected the output, treat that as "ideal"
  if (score < 0.7 && result.correctedOutput) {
    this.trainingBuffer.push({
      input: this.buildRawModelInput(task), // what we sent to the base model
      idealOutput: result.correctedOutput
    });
  }

  // ...usual system-level learning here...
}

// Offline: periodically export data for fine-tuning
async exportForFinetune() {
  const batch = this.trainingBuffer.splice(0, this.trainingBuffer.length);
  await writeToBlobStorage(batch);
}

export class ImprovingAgent extends AgenticSystem {
  // 1. Things that can change
  @field promptLibrary: Map<string, string> = new Map();
  @field examples: Example[] = [];
  @field lessons: string[] = [];
  @field knowledgeBase: VectorStore<Knowledge>;
  @field history: { task: Task; result: Result }[] = [];
  @field improvementLog: { timestamp: number; change: string }[] = [];

  // 2. Before each task: apply what we’ve learned so far
  async beforeTask(task: Task): Promise<AugmentedTask> {
    const basePrompt =
      this.promptLibrary.get(task.type) ?? "You are a careful, honest assistant.";

    const relevantExamples = await this.exampleManager.selectForPrompt(task, 2);
    const relevantLessons = this.lessons.filter((l) => l.includes(task.type)).slice(0, 3);
    const knowledge = await this.knowledgeBase.search(task.input, 5);

    const promptParts = [
      basePrompt,
      relevantLessons.length
        ? `Lessons learned from past "${task.type}" tasks:\n- ${relevantLessons.join("\n- ")}`
        : "",
      relevantExamples.length
        ? `Here are successful past examples:\n\n${relevantExamples
            .map(
              (ex) =>
                `Input:\n${ex.input}\n\nOutput:\n${ex.output}\n(Score: ${ex.score.toFixed(2)})`
            )
            .join("\n\n---\n\n")}`
        : "",
      knowledge.length
        ? `Relevant reference knowledge:\n${knowledge.map((k) => "- " + k.content).join("\n")}`
        : "",
      `Now handle this new task:\n${task.input}`
    ].filter(Boolean);

    return { ...task, prompt: promptParts.join("\n\n") };
  }

  // 3. After each task: capture experience
  async afterTask(task: AugmentedTask, result: Result) {
    this.history.push({ task, result });

    const { score } = result;

    // Memory: keep high-quality examples
    await this.exampleManager.considerForStorage(task, result.output, score);

    // Learning: distill explicit lessons
    if (score < 0.6 || score > 0.9) {
      const lesson = await this.llm.complete(
        `We just handled a "${task.type}" task with score ${score}.

Prompt used:
${task.prompt}

Output:
${result.output}

In one sentence, state a general guideline that would either:
- prevent this kind of failure in the future (if score < 0.6), or
- preserve this success pattern (if score > 0.9).

Return only the guideline.`
      );

      const text = lesson.output.trim();
      if (text && !this.lessons.includes(text)) {
        this.lessons.push(text);
      }
    }

    // Knowledge: optionally extract and store new information
    await this.maybeExtractKnowledge(task, result);
  }

  // 4. Periodically: reflect and update the scaffolding
  @schedule("0 3 * * *") // every day at 3am
  async nightlyReflection() {
    const recent = this.history.slice(-200); // recent tasks only

    const analysis = await this.llm.complete(
      `You are analyzing recent performance of an agent.

Here are some tasks, outputs, and scores:
${recent
  .map(
    ({ task, result }, i) => `Case ${i + 1}:
Type: ${task.type}
Prompt: ${task.prompt.slice(0, 500)}...
Output: ${result.output.slice(0, 500)}...
Score: ${result.score}
`
  )
  .join("\n")}

1. Identify task types whose scores are trending down.
2. For each, propose a small prompt change or additional instruction that would likely improve performance.
Return JSON: [{ taskType, proposal, rationale }].`
    );

    const proposals = JSON.parse(analysis.output);

    for (const p of proposals) {
      const prev = this.promptLibrary.get(p.taskType) ?? "You are a careful, honest assistant.";
      const newPrompt = `${prev}\n\nAdditional guidance:\n${p.proposal}`;

      // (You could test this on a held-out set before adopting it.)
      this.promptLibrary.set(p.taskType, newPrompt);

      this.improvementLog.push({
        timestamp: Date.now(),
        change: `Updated prompt for ${p.taskType}: ${p.proposal}`
      });
    }

    // Maintenance: compress examples, prune stale knowledge, etc.
    await this.exampleManager.summarizeAndCullForAllTypes();
    await this.pruneStaleKnowledge();
  }
}