71 lines
2.0 KiB
TypeScript
71 lines
2.0 KiB
TypeScript
import { AIMessage, BaseMessage, HumanMessage, SystemMessage } from '@langchain/core/messages'
|
|
|
|
export interface Config {
|
|
baseUrl: string
|
|
configurationId: string
|
|
}
|
|
|
|
export class ClientConfig implements Config {
|
|
baseUrl: string
|
|
configurationId: string
|
|
|
|
constructor(baseUrl: string, configurationId: string) {
|
|
this.baseUrl = baseUrl
|
|
this.configurationId = configurationId
|
|
}
|
|
}
|
|
|
|
export class NemoClient {
|
|
private readonly config: Config
|
|
|
|
constructor(baseUrl: string, configurationId: string) {
|
|
this.config = new ClientConfig(baseUrl, configurationId)
|
|
}
|
|
|
|
getRoleFromMessage(message: BaseMessage): string {
|
|
if (message instanceof HumanMessage || message instanceof SystemMessage) {
|
|
return 'user'
|
|
}
|
|
|
|
//AIMessage, ToolMessage, FunctionMessage
|
|
return 'assistant'
|
|
}
|
|
|
|
getContentFromMessage(message: BaseMessage): string {
|
|
return message.content.toString()
|
|
}
|
|
|
|
buildBody(messages: BaseMessage[], configurationId: string): any {
|
|
const bodyMessages = messages.map((message) => {
|
|
return {
|
|
role: this.getRoleFromMessage(message),
|
|
content: this.getContentFromMessage(message)
|
|
}
|
|
})
|
|
|
|
const body = {
|
|
config_id: configurationId,
|
|
messages: bodyMessages
|
|
}
|
|
|
|
return body
|
|
}
|
|
|
|
async chat(messages: BaseMessage[]): Promise<AIMessage[]> {
|
|
const headers = new Headers()
|
|
headers.append('Content-Type', 'application/json')
|
|
|
|
const body = this.buildBody(messages, this.config.configurationId)
|
|
|
|
const requestOptions = {
|
|
method: 'POST',
|
|
body: JSON.stringify(body),
|
|
headers: headers
|
|
}
|
|
|
|
return await fetch(`${this.config.baseUrl}/v1/chat/completions`, requestOptions)
|
|
.then((response) => response.json())
|
|
.then((body) => body.messages.map((message: any) => new AIMessage(message.content)))
|
|
}
|
|
}
|