handler.ts

import { CallbackHandler } from 'langfuse-langchain';
import { APIGatewayProxyEvent, APIGatewayProxyEventV2 } from 'aws-lambda';
import { v7 as uuidv7 } from 'uuid';

import { logger, selectLlm } from '@llm-ts-example/common-backend';
import { selectEmbeddings } from './embeddings-models.js';
import { createTool } from './tool.js';
import { createVectorStore } from './vector-store.js';
import { createAgent } from 'langchain';

export async function handle(
  event: APIGatewayProxyEvent | APIGatewayProxyEventV2,
  output: NodeJS.WritableStream,
) {
  logger.info('event', { event });
  const { question, model, embeddingType, sessionId } = event.body ? JSON.parse(event.body) : { question: undefined, model: undefined, sessionId: uuidv7() };

  const modelType = model || 'gpt-5-nano';

  const langfuse = {
    publicKey: process.env.LANGFUSE_PUBLIC_KEY,
    secretKey: process.env.LANGFUSE_SECRET_KEY,
  };

  try {
    const { indexName, model: embeddings } = selectEmbeddings({
      type: embeddingType ?? 'titan', dataSource: process.env.PINECONE_INDEX ?? '',
    });
    const vectorStore = await createVectorStore({ embeddings, indexName });
    const tool = createTool(vectorStore);
    const systemPrompt =
      'あなたは文書からコンテキストを取得する専門家です。ツールを使用してユーザーの質問に答える手助けをしてください。必ずユーザーの質問と同じ言語で答えてください。';

    const { platform, model, modelName } = selectLlm(modelType);
    // Initialize Langfuse callback handler
    const langfuseHandler = langfuse.publicKey && langfuse.secretKey ? new CallbackHandler({
      sessionId,
      flushInterval: 0,
      flushAt: 1,
      tags: [modelName],
    }) : undefined;

    const agent = createAgent({
      model,
      tools: [tool],
      systemPrompt,
    });

    logger.debug(`Langfuse: ${langfuseHandler ? 'enable' : 'disable'}`);

    const threadId = uuidv7();

    const stream = await agent.streamEvents(
      { messages: [{ role: 'user', content: question }] },
      {
        version: 'v2',
        configurable: {
          sessionId,
          thread_id: threadId,
        },
        //@ts-expect-error LangChain.js の型定義誤り？
        callbacks: langfuseHandler ? [langfuseHandler] : [],
      },
    );
    for await (const sEvent of stream) {
      logger.trace('event', sEvent);
      if (sEvent.event === 'on_chat_model_stream') {
        const chunk = sEvent.data.chunk;
        if (platform === 'aws') {
          output.write(chunk.content ?? '');
        } else {
          output.write(chunk.text ?? '');
        }
      }
    }
    output.write('\n');
  } catch (e) {
    logger.error('', JSON.parse(JSON.stringify(e)));
    output.write(`Error: ${(e as Error).message}\n`);
  }
}