34 lines
1.2 KiB
TypeScript
34 lines
1.2 KiB
TypeScript
import { DynamicTool, DynamicToolInput } from '@langchain/core/tools'
|
|
import { BaseChain } from 'langchain/chains'
|
|
import { handleEscapeCharacters } from '../../../src/utils'
|
|
import { CustomChainHandler } from '../../../src'
|
|
|
|
export interface ChainToolInput extends Omit<DynamicToolInput, 'func'> {
|
|
chain: BaseChain
|
|
}
|
|
|
|
export class ChainTool extends DynamicTool {
|
|
chain: BaseChain
|
|
|
|
constructor({ chain, ...rest }: ChainToolInput) {
|
|
super({
|
|
...rest,
|
|
func: async (input, runManager) => {
|
|
const childManagers = runManager?.getChild()
|
|
const handlers = childManagers?.handlers?.filter((handler) => !(handler instanceof CustomChainHandler)) || []
|
|
if (childManagers) childManagers.handlers = handlers
|
|
|
|
if ((chain as any).prompt && (chain as any).prompt.promptValues) {
|
|
const promptValues = handleEscapeCharacters((chain as any).prompt.promptValues, true)
|
|
const values = await chain.call(promptValues, childManagers)
|
|
return values?.text
|
|
}
|
|
|
|
const values = chain.run(input, childManagers)
|
|
return values
|
|
}
|
|
})
|
|
this.chain = chain
|
|
}
|
|
}
|