model-switch-driver.mjs 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. /** Test-only driver that selects another model after the first step's tool call. */
  2. import { installModelSelection } from '@deepseek-ai/dsh-agent'
  3. const SELECTED = { provider: 'deepseek-official', model: 'deepseek-v4-pro' }
  4. const selections = new WeakMap()
  5. export const name = 'model-switch-driver'
  6. export const inject = ['agents']
  7. /**
  8. * Install the real selection helper and change its input after `todo_write`.
  9. * @param {import('@deepseek-ai/cordis').Context} ctx - composition context.
  10. */
  11. export function apply(ctx) {
  12. ctx.on('agent/created', ({ agent }) => {
  13. const selection = { current: undefined, assembled: undefined }
  14. selections.set(agent.session, selection)
  15. installModelSelection(agent.ctx, selection)
  16. })
  17. ctx.on('session/event', (session, event) => {
  18. if (event.type !== 'todo/write') return
  19. const selection = selections.get(session)
  20. if (selection === undefined) throw new Error('model-switch driver requires an installed selection')
  21. selection.current = SELECTED
  22. })
  23. // Headless also fixes the original selection. These root waterfalls make the
  24. // driver authoritative; ending after step two avoids a reverse notice.
  25. ctx.on('system-prompt/assemble', async (_assembly, context, next) => {
  26. const assembled = await next()
  27. if (context.agent === undefined) return assembled
  28. const selected = selections.get(context.agent.session)?.assembled
  29. if (selected === undefined) return assembled
  30. return {
  31. ...assembled,
  32. variables: { ...assembled.variables, provider: selected.provider, model: selected.model },
  33. }
  34. })
  35. ctx.on('agent/request', async ({ agent }, next) => {
  36. const resolved = await next()
  37. const selected = selections.get(agent.session)?.assembled
  38. if (selected === undefined) return resolved
  39. const { reasoningEffort: _inheritedEffort, ...withoutInheritedEffort } = resolved
  40. return { ...withoutInheritedEffort, ...selected }
  41. })
  42. }