69 lines
1.7 KiB
TypeScript
69 lines
1.7 KiB
TypeScript
import { ChatOpenAI } from "@langchain/openai";
|
|
import { HumanMessage, AIMessage, SystemMessage, BaseMessage } from "@langchain/core/messages";
|
|
|
|
export interface AgentConfig
|
|
{
|
|
baseURL: string;
|
|
model: string;
|
|
apiKey: string;
|
|
headers?: Record<string, string>;
|
|
systemPrompt?: string;
|
|
}
|
|
|
|
type MessageRole = "user" | "assistant" | "system";
|
|
type MessageData = { role: MessageRole; content: string };
|
|
type SessionData = { id: string; messages: MessageData[] };
|
|
|
|
const sessions = new Map<string, SessionData>();
|
|
|
|
function toBaseMessages( messages: MessageData[] ): BaseMessage[]
|
|
{
|
|
return messages.map( m =>
|
|
{
|
|
if ( m.role === "user" ) return new HumanMessage( m.content );
|
|
if ( m.role === "assistant" ) return new AIMessage( m.content );
|
|
return new SystemMessage( m.content );
|
|
} );
|
|
}
|
|
|
|
export async function getAgentStream( id: string, userMessage: string, config: AgentConfig )
|
|
{
|
|
const model = new ChatOpenAI(
|
|
{
|
|
model: config.model,
|
|
apiKey: config.apiKey,
|
|
configuration:
|
|
{
|
|
baseURL: config.baseURL,
|
|
defaultHeaders: config.headers ?? {},
|
|
},
|
|
} );
|
|
|
|
if ( !sessions.has( id ) )
|
|
{
|
|
const messages: MessageData[] = [];
|
|
|
|
if ( config.systemPrompt?.trim() )
|
|
{
|
|
messages.push( { role: "system", content: config.systemPrompt } );
|
|
}
|
|
|
|
sessions.set( id, { id, messages } );
|
|
}
|
|
|
|
const session = sessions.get( id )!;
|
|
session.messages.push( { role: "user", content: userMessage } );
|
|
|
|
return model.stream( toBaseMessages( session.messages ) );
|
|
}
|
|
|
|
export function updateAgentConversation( id: string, assistantMessage: string ): void
|
|
{
|
|
const session = sessions.get( id );
|
|
|
|
if ( session )
|
|
{
|
|
session.messages.push( { role: "assistant", content: assistantMessage } );
|
|
}
|
|
}
|