loader-composition.spec.ts 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  1. import { afterEach, describe, expect, it } from 'vitest'
  2. import { Context } from 'cordis'
  3. import Loader from '@cordisjs/plugin-loader'
  4. import Include from '@cordisjs/plugin-include'
  5. import { mkdtemp, rm, writeFile } from 'node:fs/promises'
  6. import { tmpdir } from 'node:os'
  7. import { join } from 'node:path'
  8. import { pathToFileURL } from 'node:url'
  9. import LlmService, { createUserMessage, LlmAdapter } from '@deepseek-ai/dsh-llm'
  10. import type { GenerateOptions, StreamChunk } from '@deepseek-ai/dsh-llm'
  11. import SessionStore, { SessionId } from '@deepseek-ai/dsh-session'
  12. import SessionTitleService from '@deepseek-ai/dsh-session-title'
  13. import * as providerPlugin from '@deepseek-ai/dsh-session-title-first-message-llm'
  14. let root: string | undefined
  15. let context: Context | undefined
  16. class LoaderAdapter extends LlmAdapter {
  17. readonly requests: GenerateOptions[] = []
  18. override async * stream(options: GenerateOptions): AsyncIterable<StreamChunk> {
  19. this.requests.push(options)
  20. yield { type: 'text-delta', index: 0, text: 'Loader composed title' }
  21. yield { type: 'finish', reason: { kind: 'stop' } }
  22. }
  23. }
  24. afterEach(async () => {
  25. await context?.fiber.dispose()
  26. context = undefined
  27. if (root !== undefined) await rm(root, { recursive: true, force: true })
  28. root = undefined
  29. })
  30. async function loadComposition(): Promise<Context> {
  31. root = await mkdtemp(join(tmpdir(), 'dsh-title-loader-'))
  32. const configPath = join(root, 'cordis.yml')
  33. await writeFile(configPath, [
  34. "- name: '@deepseek-ai/dsh-llm'",
  35. "- name: '@deepseek-ai/dsh-session'",
  36. "- name: '@deepseek-ai/dsh-session-title'",
  37. ' config:',
  38. ' fallbackMaxWords: 5',
  39. ' fallbackMaxBytes: 40',
  40. ' maxTitleBytes: 80',
  41. "- name: '@deepseek-ai/dsh-session-title-first-message-llm'",
  42. ' config:',
  43. ' targetWords: 5',
  44. ' targetCjkCharacters: 10',
  45. ' maxInputBytes: 1000',
  46. ' maxOutputTokens: 32',
  47. ' timeoutMs: 1000',
  48. " provider: 'title-route'",
  49. " model: 'title-model'",
  50. '',
  51. ].join('\n'))
  52. context = new Context()
  53. context.baseUrl = pathToFileURL(root).href + '/'
  54. await context.plugin(Loader)
  55. context.loader.builtins.include = Include
  56. const modules = new Map<string, unknown>([
  57. ['@deepseek-ai/dsh-llm', LlmService],
  58. ['@deepseek-ai/dsh-session', SessionStore],
  59. ['@deepseek-ai/dsh-session-title', SessionTitleService],
  60. ['@deepseek-ai/dsh-session-title-first-message-llm', providerPlugin],
  61. ])
  62. context.loader.internal = {
  63. version: 'v2',
  64. async import(specifier: string) {
  65. if (!modules.has(specifier)) throw new Error(`unexpected Loader import: ${specifier}`)
  66. return modules.get(specifier)
  67. },
  68. } as unknown as NonNullable<typeof context.loader.internal>
  69. await context.loader.create({
  70. name: 'cordis:include',
  71. config: { path: pathToFileURL(configPath).href },
  72. })
  73. await context.loader.await()
  74. return context
  75. }
  76. describe('session-title Loader composition', () => {
  77. it('loads the service and one model provider with required deployment policy', async () => {
  78. const ctx = await loadComposition()
  79. const unloaded = [...ctx.loader.entries()]
  80. .filter(entry => entry.fiber === undefined && !entry.disabled)
  81. .map(entry => entry.options.name)
  82. expect(unloaded).toEqual([])
  83. const adapter = new LoaderAdapter()
  84. ctx.llm.registerAdapter(['title-route'], adapter)
  85. const session = ctx.sessions.create(SessionId('loader-title'))
  86. session.append('turn/start', {
  87. turn: 1,
  88. })
  89. const message = session.append('user/message', createUserMessage({
  90. content: [{ type: 'text', text: 'Compose a title through Loader' }],
  91. source: { kind: 'user' },
  92. }), { surfaceOp: 'append' })
  93. await new Promise(resolve => setTimeout(resolve, 0))
  94. session.append('request/header', {
  95. header: { config: { provider: 'main-route', model: 'main-model' } },
  96. reason: 'initial',
  97. })
  98. await new Promise(resolve => setTimeout(resolve, 0))
  99. expect(adapter.requests[0]).toMatchObject({ provider: 'title-route', model: 'title-model' })
  100. expect(ctx.sessionTitle.get(session)).toMatchObject({
  101. title: 'Loader composed title',
  102. messageSeqs: [message.seq],
  103. source: {
  104. kind: 'provider',
  105. provider: 'session-title-first-message-llm',
  106. model: { provider: 'title-route', model: 'title-model' },
  107. },
  108. })
  109. })
  110. })