stack/apps/api/src/llm/llm.service.spec.ts

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { ServiceUnavailableException } from "@nestjs/common";
import { LlmService } from "./llm.service";
import { LlmManagerService } from "./llm-manager.service";
import { LlmTelemetryTrackerService } from "./llm-telemetry-tracker.service";
import type { ChatRequestDto, EmbedRequestDto, ChatResponseDto, EmbedResponseDto } from "./dto";
import type {
  LlmProviderInterface,
  LlmProviderHealthStatus,
} from "./providers/llm-provider.interface";

describe("LlmService", () => {
  let service: LlmService;
  let mockManagerService: {
    getDefaultProvider: ReturnType<typeof vi.fn>;
  };
  let mockTelemetryTracker: {
    trackLlmCompletion: ReturnType<typeof vi.fn>;
  };
  let mockProvider: {
    chat: ReturnType<typeof vi.fn>;
    chatStream: ReturnType<typeof vi.fn>;
    embed: ReturnType<typeof vi.fn>;
    listModels: ReturnType<typeof vi.fn>;
    checkHealth: ReturnType<typeof vi.fn>;
    name: string;
    type: string;
  };

  beforeEach(async () => {
    // Create mock provider
    mockProvider = {
      chat: vi.fn(),
      chatStream: vi.fn(),
      embed: vi.fn(),
      listModels: vi.fn(),
      checkHealth: vi.fn(),
      name: "Test Provider",
      type: "ollama",
    };

    // Create mock manager service
    mockManagerService = {
      getDefaultProvider: vi.fn().mockResolvedValue(mockProvider),
    };

    // Create mock telemetry tracker
    mockTelemetryTracker = {
      trackLlmCompletion: vi.fn(),
    };

    const module: TestingModule = await Test.createTestingModule({
      providers: [
        LlmService,
        {
          provide: LlmManagerService,
          useValue: mockManagerService,
        },
        {
          provide: LlmTelemetryTrackerService,
          useValue: mockTelemetryTracker,
        },
      ],
    }).compile();

    service = module.get<LlmService>(LlmService);
  });

  it("should be defined", () => {
    expect(service).toBeDefined();
  });

  describe("checkHealth", () => {
    it("should delegate to provider and return healthy status", async () => {
      const healthStatus: LlmProviderHealthStatus = {
        healthy: true,
        provider: "ollama",
        endpoint: "http://localhost:11434",
        models: ["llama3.2"],
      };
      mockProvider.checkHealth.mockResolvedValue(healthStatus);

      const result = await service.checkHealth();

      expect(mockManagerService.getDefaultProvider).toHaveBeenCalled();
      expect(mockProvider.checkHealth).toHaveBeenCalled();
      expect(result).toEqual(healthStatus);
    });

    it("should return unhealthy status on error", async () => {
      mockProvider.checkHealth.mockRejectedValue(new Error("Connection failed"));

      const result = await service.checkHealth();

      expect(result.healthy).toBe(false);
      expect(result.error).toContain("Connection failed");
    });

    it("should handle manager service failure", async () => {
      mockManagerService.getDefaultProvider.mockRejectedValue(new Error("No provider configured"));

      const result = await service.checkHealth();

      expect(result.healthy).toBe(false);
      expect(result.error).toContain("No provider configured");
    });
  });

  describe("listModels", () => {
    it("should delegate to provider and return models", async () => {
      const models = ["llama3.2", "mistral"];
      mockProvider.listModels.mockResolvedValue(models);

      const result = await service.listModels();

      expect(mockManagerService.getDefaultProvider).toHaveBeenCalled();
      expect(mockProvider.listModels).toHaveBeenCalled();
      expect(result).toEqual(models);
    });

    it("should throw ServiceUnavailableException on error", async () => {
      mockProvider.listModels.mockRejectedValue(new Error("Failed to fetch models"));

      await expect(service.listModels()).rejects.toThrow(ServiceUnavailableException);
    });
  });

  describe("chat", () => {
    const request: ChatRequestDto = {
      model: "llama3.2",
      messages: [{ role: "user", content: "Hi" }],
    };

    it("should delegate to provider and return response", async () => {
      const response: ChatResponseDto = {
        model: "llama3.2",
        message: { role: "assistant", content: "Hello" },
        done: true,
        totalDuration: 1000,
      };
      mockProvider.chat.mockResolvedValue(response);

      const result = await service.chat(request);

      expect(mockManagerService.getDefaultProvider).toHaveBeenCalled();
      expect(mockProvider.chat).toHaveBeenCalledWith(request);
      expect(result).toEqual(response);
    });

    it("should track telemetry on successful chat", async () => {
      const response: ChatResponseDto = {
        model: "llama3.2",
        message: { role: "assistant", content: "Hello" },
        done: true,
        promptEvalCount: 10,
        evalCount: 20,
      };
      mockProvider.chat.mockResolvedValue(response);

      await service.chat(request, "chat");

      expect(mockTelemetryTracker.trackLlmCompletion).toHaveBeenCalledWith(
        expect.objectContaining({
          model: "llama3.2",
          providerType: "ollama",
          operation: "chat",
          inputTokens: 10,
          outputTokens: 20,
          callingContext: "chat",
          success: true,
        })
      );
    });

    it("should track telemetry on failed chat", async () => {
      mockProvider.chat.mockRejectedValue(new Error("Chat failed"));

      await expect(service.chat(request)).rejects.toThrow(ServiceUnavailableException);

      expect(mockTelemetryTracker.trackLlmCompletion).toHaveBeenCalledWith(
        expect.objectContaining({
          model: "llama3.2",
          operation: "chat",
          success: false,
        })
      );
    });

    it("should throw ServiceUnavailableException on error", async () => {
      mockProvider.chat.mockRejectedValue(new Error("Chat failed"));

      await expect(service.chat(request)).rejects.toThrow(ServiceUnavailableException);
    });
  });

  describe("chatStream", () => {
    const request: ChatRequestDto = {
      model: "llama3.2",
      messages: [{ role: "user", content: "Hi" }],
      stream: true,
    };

    it("should delegate to provider and yield chunks", async () => {
      async function* mockGenerator(): AsyncGenerator<ChatResponseDto> {
        yield {
          model: "llama3.2",
          message: { role: "assistant", content: "Hello" },
          done: false,
        };
        yield {
          model: "llama3.2",
          message: { role: "assistant", content: " world" },
          done: true,
        };
      }

      mockProvider.chatStream.mockReturnValue(mockGenerator());

      const chunks: ChatResponseDto[] = [];
      for await (const chunk of service.chatStream(request)) {
        chunks.push(chunk);
      }

      expect(mockManagerService.getDefaultProvider).toHaveBeenCalled();
      expect(mockProvider.chatStream).toHaveBeenCalledWith(request);
      expect(chunks.length).toBe(2);
      expect(chunks[0].message.content).toBe("Hello");
      expect(chunks[1].message.content).toBe(" world");
    });

    it("should track telemetry after stream completes", async () => {
      async function* mockGenerator(): AsyncGenerator<ChatResponseDto> {
        yield {
          model: "llama3.2",
          message: { role: "assistant", content: "Hello" },
          done: false,
        };
        yield {
          model: "llama3.2",
          message: { role: "assistant", content: " world" },
          done: true,
          promptEvalCount: 5,
          evalCount: 10,
        };
      }

      mockProvider.chatStream.mockReturnValue(mockGenerator());

      const chunks: ChatResponseDto[] = [];
      for await (const chunk of service.chatStream(request, "brain")) {
        chunks.push(chunk);
      }

      expect(mockTelemetryTracker.trackLlmCompletion).toHaveBeenCalledWith(
        expect.objectContaining({
          model: "llama3.2",
          providerType: "ollama",
          operation: "chatStream",
          inputTokens: 5,
          outputTokens: 10,
          callingContext: "brain",
          success: true,
        })
      );
    });

    it("should estimate tokens when provider does not return counts in stream", async () => {
      async function* mockGenerator(): AsyncGenerator<ChatResponseDto> {
        yield {
          model: "llama3.2",
          message: { role: "assistant", content: "Hello world" },
          done: false,
        };
        yield {
          model: "llama3.2",
          message: { role: "assistant", content: "" },
          done: true,
        };
      }

      mockProvider.chatStream.mockReturnValue(mockGenerator());

      const chunks: ChatResponseDto[] = [];
      for await (const chunk of service.chatStream(request)) {
        chunks.push(chunk);
      }

      // Should use estimated tokens since no actual counts provided
      expect(mockTelemetryTracker.trackLlmCompletion).toHaveBeenCalledWith(
        expect.objectContaining({
          operation: "chatStream",
          success: true,
          // Input estimated from "Hi" -> ceil(2/4) = 1
          inputTokens: 1,
          // Output estimated from "Hello world" -> ceil(11/4) = 3
          outputTokens: 3,
        })
      );
    });

    it("should track telemetry on stream failure", async () => {
      async function* errorGenerator(): AsyncGenerator<ChatResponseDto> {
        throw new Error("Stream failed");
      }

      mockProvider.chatStream.mockReturnValue(errorGenerator());

      const generator = service.chatStream(request);
      await expect(generator.next()).rejects.toThrow(ServiceUnavailableException);

      expect(mockTelemetryTracker.trackLlmCompletion).toHaveBeenCalledWith(
        expect.objectContaining({
          operation: "chatStream",
          success: false,
        })
      );
    });

    it("should throw ServiceUnavailableException on error", async () => {
      async function* errorGenerator(): AsyncGenerator<ChatResponseDto> {
        throw new Error("Stream failed");
      }

      mockProvider.chatStream.mockReturnValue(errorGenerator());

      const generator = service.chatStream(request);
      await expect(generator.next()).rejects.toThrow(ServiceUnavailableException);
    });
  });

  describe("embed", () => {
    const request: EmbedRequestDto = {
      model: "llama3.2",
      input: ["test text"],
    };

    it("should delegate to provider and return embeddings", async () => {
      const response: EmbedResponseDto = {
        model: "llama3.2",
        embeddings: [[0.1, 0.2, 0.3]],
        totalDuration: 500,
      };
      mockProvider.embed.mockResolvedValue(response);

      const result = await service.embed(request);

      expect(mockManagerService.getDefaultProvider).toHaveBeenCalled();
      expect(mockProvider.embed).toHaveBeenCalledWith(request);
      expect(result).toEqual(response);
    });

    it("should track telemetry on successful embed", async () => {
      const response: EmbedResponseDto = {
        model: "llama3.2",
        embeddings: [[0.1, 0.2, 0.3]],
        totalDuration: 500,
      };
      mockProvider.embed.mockResolvedValue(response);

      await service.embed(request, "embed");

      expect(mockTelemetryTracker.trackLlmCompletion).toHaveBeenCalledWith(
        expect.objectContaining({
          model: "llama3.2",
          providerType: "ollama",
          operation: "embed",
          outputTokens: 0,
          callingContext: "embed",
          success: true,
        })
      );
    });

    it("should track telemetry on failed embed", async () => {
      mockProvider.embed.mockRejectedValue(new Error("Embedding failed"));

      await expect(service.embed(request)).rejects.toThrow(ServiceUnavailableException);

      expect(mockTelemetryTracker.trackLlmCompletion).toHaveBeenCalledWith(
        expect.objectContaining({
          operation: "embed",
          success: false,
        })
      );
    });

    it("should throw ServiceUnavailableException on error", async () => {
      mockProvider.embed.mockRejectedValue(new Error("Embedding failed"));

      await expect(service.embed(request)).rejects.toThrow(ServiceUnavailableException);
    });
  });
});