Devlog Feb 8: AI Agent Improvements

This week focused on enhancing the AI Agent v3 with multi-message streaming, improved progress tracking, and module validation.

Multi-Message Streaming

Message Segments

We implemented multi-message streaming where messages are split at tool boundaries:

// src/lib/server/ai/agent/v3/core/molos-agent.ts
async *stream(prompt: string): AsyncGenerator<StreamChunk> {
  // Planning phase
  yield { type: 'message', role: 'assistant', content: 'Planning...' };

  const plan = await this.createPlan(prompt);

  // Execute each step
  for (const step of plan.steps) {
    yield {
      type: 'message',
      role: 'assistant',
      content: `Executing: ${step.description}`
    };

    const result = await this.executeStep(step);

    yield {
      type: 'message',
      role: 'tool',
      content: JSON.stringify(result)
    };

    yield {
      type: 'message',
      role: 'assistant',
      content: await this.processResult(result)
    };
  }

  // Final response
  yield {
    type: 'message',
    role: 'assistant',
    content: this.formatFinalResponse()
  };
}

Step Number Tracking

Each progress event includes step numbers for better tracking:

interface ProgressEvent {
  type: ProgressEventType;
  stepNumber: number;
  totalSteps: number;
  message: string;
  details?: Record<string, unknown>;
  timestamp: Date;
}

export interface ProgressEventType {
  'step_start': { stepNumber: number; description: string };
  'step_complete': { stepNumber: number; result: unknown };
  'step_error': { stepNumber: number; error: string };
  'planning': { stepNumber: 0; steps: PlanStep[] };
  'complete': { stepNumber: number; finalResponse: string };
}

Message Persistence

Each message segment is saved separately for proper persistence:

// src/lib/server/ai/agent/v3/message-persistence.ts
export async function saveMessageSegments(
  sessionId: string,
  segments: MessageSegment[]
): Promise<void> {
  for (const segment of segments) {
    await db.insert(aiMessages).values({
      id: crypto.randomUUID(),
      sessionId,
      role: segment.role,
      content: segment.content,
      type: segment.type,
      stepNumber: segment.stepNumber,
      createdAt: new Date()
    });
  }
}

UI Integration

Streaming Progress Events

The chat UI handles streaming progress events:

// src/routes/ui/(modules)/ai/chat/+page.svelte
<script lang="ts">
  import { onMount } from 'svelte';
  import type { ProgressEvent } from '$lib/types/ai';

  let progressEvents: ProgressEvent[] = [];
  let currentStep = 0;

  async function streamResponse(prompt: string) {
    const response = await fetch('/api/agent/stream', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ prompt })
    });

    const reader = response.body.getReader();
    const decoder = new TextDecoder();

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;

      const chunk = decoder.decode(value);
      const event = JSON.parse(chunk) as ProgressEvent;

      progressEvents = [...progressEvents, event];

      if (event.type !== 'planning' && event.type !== 'complete') {
        currentStep = event.stepNumber;
      }
    }
  }
</script>

UIMessage Format

We updated the chat UI to use the UIMessage format from @ai-sdk/svelte:

// src/lib/components/ai/chat/AiChatWorkspace.svelte
<script lang="ts">
  import type { UIMessage } from 'ai';

  let messages: UIMessage[] = [];

  function addProgressToMessages(event: ProgressEvent): void {
    const message: UIMessage = {
      id: crypto.randomUUID(),
      role: 'assistant',
      content: event.message,
      createdAt: new Date()
    };

    messages = [...messages, message];
  }
</script>

Progress Display

Progress is displayed in an accordion component:

<!-- src/lib/components/ai/progress/ProgressAccordion.svelte -->
<script lang="ts">
  import type { ProgressEvent } from '$lib/types/ai';

  export let events: ProgressEvent[];

  let expandedStep: number | null = null;

  $: completedSteps = events.filter(e => e.type === 'step_complete');
  $: currentStep = events.length > 0 ? events[events.length - 1].stepNumber : 0;
</script>

<div class="progress-accordion">
  {#each completedSteps as event, i}
    <div class="step" class:current={i === completedSteps.length - 1}>
      <div class="step-header" on:click={() => expandedStep = i}>
        <span class="step-number">Step {event.stepNumber}</span>
        <span class="step-message">{event.message}</span>
        <span class="expand-icon">{expandedStep === i ? '▼' : '▶'}</span>
      </div>

      {#if expandedStep === i}
        <div class="step-details">
          <pre>{JSON.stringify(event.details, null, 2)}</pre>
        </div>
      {/if}
    </div>
  {/each}
</div>

Module Validation

Safe Module Loading

We implemented safe external module loading with error handling:

// src/lib/server/modules/external-loader.ts
export async function loadExternalModule(
  modulePath: string,
  options: LoadOptions
): Promise<ExternalModule> {
  try {
    const module = await import(modulePath);

    // Validate module structure
    if (!module.default || typeof module.default !== 'function') {
      throw new InvalidModuleError('Module must export a default function');
    }

    return {
      path: modulePath,
      module: module.default,
      metadata: module.metadata || {}
    };
  } catch (error) {
    if (error instanceof InvalidModuleError) {
      throw error;
    }

    await disableModule(modulePath, {
      reason: 'Failed to load',
      error: error instanceof Error ? error.message : 'Unknown error'
    });

    throw new ModuleLoadError(
      `Failed to load module ${modulePath}: ${error}`,
      { cause: error }
    );
  }
}

Module Validation

During initialization, modules are validated:

// src/lib/server/modules/validator.ts
export async function validateModule(
  module: Module
): Promise<ValidationResult> {
  const errors: ValidationError[] = [];

  // Check required properties
  if (!module.id) {
    errors.push({ field: 'id', message: 'Module ID is required' });
  }

  if (!module.version) {
    errors.push({ field: 'version', message: 'Module version is required' });
  }

  // Validate tools
  if (module.tools) {
    for (const tool of module.tools) {
      const toolErrors = validateTool(tool);
      errors.push(...toolErrors);
    }
  }

  // Validate configuration schema
  if (module.configSchema) {
    try {
      // Test schema with example data
      module.configSchema.parse({});
    } catch (error) {
      errors.push({
        field: 'configSchema',
        message: 'Invalid configuration schema'
      });
    }
  }

  return {
    isValid: errors.length === 0,
    errors
  };
}

Auto-Disable Failed Modules

Failed modules are automatically disabled:

// src/lib/server/modules/auto-disable.ts
export async function disableModule(
  moduleId: string,
  reason: DisableReason
): Promise<void> {
  await db.update(modules)
    .set({
      enabled: false,
      disabledAt: new Date(),
      disableReason: reason.reason,
      disableError: reason.error
    })
    .where(eq(modules.id, moduleId));

  logger.warn(`Module ${moduleId} auto-disabled`, reason);
}

export async function linkModules(): Promise<void> {
  const modules = await getExternalModules();

  for (const module of modules) {
    try {
      await validateAndInitializeModule(module);
    } catch (error) {
      await disableModule(module.id, {
        reason: 'Validation failed',
        error: error instanceof Error ? error.message : 'Unknown error'
      });
    }
  }
}

Provider Enhancements

OpenAI-Compatible Providers

We added support for OpenAI-compatible providers like Z.AI:

// src/lib/server/ai/agent/v3/providers/factory.ts
export function createProvider(
  config: ProviderConfig
): CoreMessage[] {
  switch (config.type) {
    case 'openai':
      return createOpenAI(config);

    case 'anthropic':
      return createAnthropic(config);

    case 'openai-compatible':
      if (config.useCompatibilityMode) {
        return createOpenAICompatibleProvider(config);
      }
      return createOpenAI(config);

    default:
      throw new Error(`Unsupported provider: ${config.type}`);
  }
}

function createOpenAICompatibleProvider(
  config: ProviderConfig
): OpenAIProvider {
  const client = new OpenAI({
    apiKey: config.apiKey,
    baseURL: config.baseUrl,
    dangerouslyAllowBrowser: false
  });

  return {
    stream: async (messages, options) => {
      return client.chat.completions.create({
        model: config.model,
        messages,
        stream: true,
        ...options
      });
    }
  };
}

.chat() Method

For OpenAI-compatible providers, we use the .chat() method to force chat completions:

// Use .chat() to ensure chat completions
const completion = await client.chat.completions.create({
  model: 'zai-v3',
  messages: convertToOpenAIMessages(messages),
  stream: true,
  temperature: 0.7
});

Testing

Unit Tests

We added comprehensive unit tests for v3 agent components:

// tests/agent/v3/schema-converter.test.ts
import { describe, it, expect } from 'vitest';
import { convertZodToJsonSchema } from '$lib/server/ai/agent/v3/tools/schema-converter';

describe('Schema Converter', () => {
  it('should convert Zod string schema', () => {
    const schema = z.string().describe('A test string');
    const jsonSchema = convertZodToJsonSchema(schema);

    expect(jsonSchema).toEqual({
      type: 'string',
      description: 'A test string'
    });
  });

  it('should convert Zod object schema', () => {
    const schema = z.object({
      name: z.string(),
      age: z.number()
    });

    const jsonSchema = convertZodToJsonSchema(schema);

    expect(jsonSchema.type).toBe('object');
    expect(jsonSchema.properties).toHaveProperty('name');
    expect(jsonSchema.properties).toHaveProperty('age');
  });
});

Tool Wrapper Tests

// tests/agent/v3/tool-wrapper.test.ts
import { describe, it, expect, vi } from 'vitest';
import { wrapTool } from '$lib/server/ai/agent/v3/tools/tool-wrapper';

describe('Tool Wrapper', () => {
  it('should wrap tool with error handling', async () => {
    const mockTool = vi.fn().mockResolvedValue({ success: true });
    const wrapped = wrapTool(mockTool, 'test-tool');

    const result = await wrapped({ param: 'value' });

    expect(mockTool).toHaveBeenCalledWith({ param: 'value' });
    expect(result).toEqual({ success: true });
  });

  it('should handle tool errors', async () => {
    const mockTool = vi.fn().mockRejectedValue(new Error('Test error'));
    const wrapped = wrapTool(mockTool, 'test-tool');

    await expect(wrapped({})).rejects.toThrow('Tool execution failed');
  });
});

View commits on GitHub

What's Next

Upcoming improvements include:

• Enhanced tool discovery and suggestion
• Multi-agent collaboration protocols
• Custom tool registration API
• Agent memory and conversation persistence

---

Check out the AI Agent v3 Launch announcement for the feature overview.