Testing

The best way to test agents is to pass a MockLanguageModel to the agent instead of a real language model.

TypeScript

agent.test.ts

/* eslint-disable @typescript-eslint/no-floating-promises */

import { MockLanguageModel } from "@hoangvvo/llm-sdk/test";
import test, { suite, type TestContext } from "node:test";
import { Agent } from "./agent.ts";

suite("Agent#run", () => {
  test("creates session, runs, and closes", async (t: TestContext) => {
    const model = new MockLanguageModel();
    model.enqueueGenerateResult({
      response: { content: [{ type: "text", text: "Mock response" }] },
    });
    const agent = new Agent({
      name: "test-agent",
      model,
    });

    const response = await agent.run({
      context: {},
      input: [
        {
          type: "message",
          role: "user",
          content: [{ type: "text", text: "Hello" }],
        },
      ],
    });

    t.assert.deepStrictEqual(response, {
      content: [{ type: "text", text: "Mock response" }],
      output: [
        {
          type: "model",
          content: [{ type: "text", text: "Mock response" }],
        },
      ],
    });
  });
});

suite("Agent#runStream", () => {
  test("creates session, streams, and closes", async (t: TestContext) => {
    const model = new MockLanguageModel();
    model.enqueueStreamResult({
      partials: [
        {
          delta: { index: 0, part: { type: "text", text: "Mock" } },
        },
      ],
    });
    const agent = new Agent({
      name: "test-agent",
      model,
    });

    const generator = agent.runStream({
      context: {},
      input: [
        {
          type: "message",
          role: "user",
          content: [{ type: "text", text: "Hello" }],
        },
      ],
    });

    const events = [];
    let current = await generator.next();
    while (!current.done) {
      events.push(current.value);
      current = await generator.next();
    }

    t.assert.deepStrictEqual(events, [
      {
        event: "partial",
        delta: { index: 0, part: { type: "text", text: "Mock" } },
      },
      {
        event: "item",
        index: 0,
        item: {
          type: "model",
          content: [{ type: "text", text: "Mock" }],
        },
      },
      {
        event: "response",
        content: [{ type: "text", text: "Mock" }],
        output: [
          {
            type: "model",
            content: [{ type: "text", text: "Mock" }],
          },
        ],
      },
    ]);
    t.assert.deepStrictEqual(current.value, {
      content: [{ type: "text", text: "Mock" }],
      output: [
        {
          type: "model",
          content: [{ type: "text", text: "Mock" }],
        },
      ],
    });
  });
});

Rust

agent_test.rs

use futures::TryStreamExt;
use llm_agent::{
    Agent, AgentItem, AgentParams, AgentRequest, AgentResponse, AgentStreamEvent,
    AgentStreamItemEvent,
};
use llm_sdk::{
    llm_sdk_test::{MockLanguageModel, MockStreamResult},
    ContentDelta, Message, ModelResponse, Part, PartDelta, PartialModelResponse, TextPartDelta,
    UserMessage,
};
use std::sync::Arc;

#[tokio::test]
async fn agent_run_creates_session_runs_and_finishes() {
    let model = Arc::new(MockLanguageModel::new());
    model.enqueue_generate(ModelResponse {
        content: vec![Part::text("Mock response")],
        ..Default::default()
    });

    let agent = Agent::new(AgentParams::new("test-agent", model.clone()));

    let response = agent
        .run(AgentRequest {
            context: (),
            input: vec![AgentItem::Message(Message::User(UserMessage {
                content: vec![Part::text("Hello")],
            }))],
        })
        .await
        .expect("agent run succeeds");

    let expected = AgentResponse {
        content: vec![Part::text("Mock response")],
        output: vec![AgentItem::Model(ModelResponse {
            content: vec![Part::text("Mock response")],
            ..Default::default()
        })],
    };

    assert_eq!(response, expected);
}

#[tokio::test]
async fn agent_run_stream_creates_session_streams_and_finishes() {
    let model = Arc::new(MockLanguageModel::new());
    model.enqueue_stream(MockStreamResult::partials(vec![PartialModelResponse {
        delta: Some(ContentDelta {
            index: 0,
            part: PartDelta::Text(TextPartDelta {
                text: "Mock".to_string(),
            }),
        }),
        ..Default::default()
    }]));

    let agent = Agent::new(AgentParams::new("test-agent", model.clone()));

    let stream = agent
        .run_stream(AgentRequest {
            context: (),
            input: vec![AgentItem::Message(Message::User(UserMessage {
                content: vec![Part::text("Hello")],
            }))],
        })
        .await
        .expect("agent run_stream succeeds");

    let events = stream
        .map_err(|err| err.to_string())
        .try_collect::<Vec<_>>()
        .await
        .expect("collect stream");

    let expected = vec![
        AgentStreamEvent::Partial(PartialModelResponse {
            delta: Some(ContentDelta {
                index: 0,
                part: PartDelta::Text(TextPartDelta {
                    text: "Mock".to_string(),
                }),
            }),
            ..Default::default()
        }),
        AgentStreamEvent::Item(AgentStreamItemEvent {
            index: 0,
            item: AgentItem::Model(ModelResponse {
                content: vec![Part::text("Mock")],
                ..Default::default()
            }),
        }),
        AgentStreamEvent::Response(AgentResponse {
            content: vec![Part::text("Mock")],
            output: vec![AgentItem::Model(ModelResponse {
                content: vec![Part::text("Mock")],
                ..Default::default()
            })],
        }),
    ];

    assert_eq!(events, expected);
}

Go

agent_test.go

package llmagent_test

import (
  "context"
  "testing"

  "github.com/google/go-cmp/cmp"
  llmagent "github.com/hoangvvo/llm-sdk/agent-go"
  llmsdk "github.com/hoangvvo/llm-sdk/sdk-go"
  "github.com/hoangvvo/llm-sdk/sdk-go/llmsdktest"
)

func TestAgent_Run(t *testing.T) {
  t.Run("creates session, runs, and closes", func(t *testing.T) {
    model := llmsdktest.NewMockLanguageModel()
    model.EnqueueGenerateResult(
      llmsdktest.NewMockGenerateResultResponse(llmsdk.ModelResponse{
        Content: []llmsdk.Part{
          {TextPart: &llmsdk.TextPart{Text: "Mock response"}},
        },
      }),
    )
    agent := llmagent.NewAgent[map[string]interface{}]("test-agent", model)

    response, err := agent.Run(context.Background(), llmagent.AgentRequest[map[string]interface{}]{
      Context: map[string]interface{}{},
      Input: []llmagent.AgentItem{
        llmagent.NewAgentItemMessage(llmsdk.Message{
          UserMessage: &llmsdk.UserMessage{
            Content: []llmsdk.Part{
              {TextPart: &llmsdk.TextPart{Text: "Hello"}},
            },
          },
        }),
      },
    })

    if err != nil {
      t.Fatalf("expected no error, got %v", err)
    }

    expectedResponse := &llmagent.AgentResponse{
      Content: []llmsdk.Part{
        {TextPart: &llmsdk.TextPart{Text: "Mock response"}},
      },
      Output: []llmagent.AgentItem{
        llmagent.NewAgentItemModelResponse(llmsdk.ModelResponse{
          Content: []llmsdk.Part{
            {TextPart: &llmsdk.TextPart{Text: "Mock response"}},
          },
        }),
      },
    }

    if diff := cmp.Diff(expectedResponse, response); diff != "" {
      t.Errorf("response mismatch (-want +got): %s", diff)
    }
  })
}

func TestAgent_RunStream(t *testing.T) {
  t.Run("creates session, streams, and closes", func(t *testing.T) {
    model := llmsdktest.NewMockLanguageModel()
    model.EnqueueStreamResult(
      llmsdktest.NewMockStreamResultPartials([]llmsdk.PartialModelResponse{
        {
          Delta: &llmsdk.ContentDelta{
            Index: 0,
            Part: llmsdk.PartDelta{
              TextPartDelta: &llmsdk.TextPartDelta{Text: "Mock"},
            },
          },
        },
      }),
    )
    agent := llmagent.NewAgent[map[string]interface{}]("test-agent", model)

    stream, err := agent.RunStream(context.Background(), llmagent.AgentRequest[map[string]interface{}]{
      Context: map[string]interface{}{},
      Input: []llmagent.AgentItem{
        llmagent.NewAgentItemMessage(llmsdk.Message{
          UserMessage: &llmsdk.UserMessage{
            Content: []llmsdk.Part{
              {TextPart: &llmsdk.TextPart{Text: "Hello"}},
            },
          },
        }),
      },
    })

    if err != nil {
      t.Fatalf("expected no error, got %v", err)
    }

    events := []*llmagent.AgentStreamEvent{}
    for stream.Next() {
      events = append(events, stream.Current())
    }

    if err := stream.Err(); err != nil {
      t.Fatalf("expected no error, got %v", err)
    }

    expectedEvents := []*llmagent.AgentStreamEvent{
      {
        Partial: &llmsdk.PartialModelResponse{
          Delta: &llmsdk.ContentDelta{
            Index: 0,
            Part: llmsdk.PartDelta{
              TextPartDelta: &llmsdk.TextPartDelta{Text: "Mock"},
            },
          },
        },
      },
      llmagent.NewAgentStreamItemEvent(
        0,
        llmagent.NewAgentItemModelResponse(llmsdk.ModelResponse{
          Content: []llmsdk.Part{
            {TextPart: &llmsdk.TextPart{Text: "Mock"}},
          },
        }),
      ),
      {
        Response: &llmagent.AgentResponse{
          Content: []llmsdk.Part{
            {TextPart: &llmsdk.TextPart{Text: "Mock"}},
          },
          Output: []llmagent.AgentItem{
            llmagent.NewAgentItemModelResponse(llmsdk.ModelResponse{
              Content: []llmsdk.Part{
                {TextPart: &llmsdk.TextPart{Text: "Mock"}},
              },
            }),
          },
        },
      },
    }

    if diff := cmp.Diff(expectedEvents, events); diff != "" {
      t.Errorf("stream events mismatch (-want +got):\n%s", diff)
    }
  })
}