langchain.js

import { HumanMessage, AIMessage, ToolMessage } from '@langchain/core/messages';
import { END, START, Command } from "@langchain/langgraph";
import { StateGraph, MessagesAnnotation } from '@langchain/langgraph';
import { Annotation } from "@langchain/langgraph";
import { v4 } from 'uuid';
import dotenv from 'dotenv';

dotenv.config();

const StateAnnotation = Annotation.Root({
  messages: Annotation({
    reducer: (x, y) => x.concat(y),
  }),
});

const routeMessage = (state) => {
  const { messages } = state;
  const lastMessage = messages[messages.length - 1];
  
  if (!lastMessage?.tool_calls?.length) return END;

  const hasSchedule = lastMessage.tool_calls.some(
    call => call.name === 'schedule_appointment'
  );
  const hasDate = lastMessage.tool_calls.some(
    call => call.name === 'get_current_date'
  );

  // Route to pre-appointment node if scheduling needs date
  if (hasSchedule && !hasDate) return 'pre_appointment';
  
  // Otherwise proceed to tools node
  return 'tools';
};

const preAppointmentNode = (state) => {
  const lastMessage = state.messages[state.messages.length - 1];
  const toolCalls = lastMessage.tool_calls;

  // Create new tool calls array with date first
  const newToolCalls = [
    {
      name: 'get_current_date',
      id: v4(),
      type: 'tool_call',
      args: {}
    },
    ...toolCalls.filter(call => call.name !== 'schedule_appointment')
  ];

  // Return Command to update state and route to tools
  return new Command({
    update: {
      messages: [
        ...state.messages.slice(0, -1),
        {
          ...lastMessage,
          tool_calls: newToolCalls
        }
      ]
    },
    goto: 'tools'
  });
};

class Agent {
  constructor(model, toolNode, tools) {
    this.app = null;
    this.model = model;
    this.toolNode = toolNode;
    this.tools = tools;
    this.initializeWorkflow();
  }

  initializeWorkflow() {
    const callModel = async (state) => {
      const response = await this.model.invoke(state.messages);
      return { messages: [response] };
    };

    const workflow = new StateGraph(MessagesAnnotation)
      .addNode('agent', callModel)
      .addNode('tools', this.toolNode)
      .addNode('pre_appointment', preAppointmentNode)
      .addEdge(START, 'agent')
      .addConditionalEdges(
        'agent',
        routeMessage,
        {
          tools: 'tools',
          pre_appointment: 'pre_appointment',
          [END]: END
        }
      )
      .addEdge('pre_appointment', 'tools')
      .addEdge('tools', 'agent');

    this.app = workflow.compile();
  }

  // Keep your existing sendMessage and disconnect methods
  async sendMessage(state, message, tokenCallBack, eventCallback) {
    const humanMessage = new HumanMessage(message);
    await eventCallback({ event: 'human_message', data: message });
    state.messages.push(humanMessage);

    const stream = await this.app.stream(
      { messages: state.messages },
      { streamMode: ["messages"] }
    );

    let aiMessageContent = "";
    for await (const chunk of stream) {
      if (chunk[0] === "messages") {
        const token = chunk[1][0]?.content;
        if (token) {
          aiMessageContent += token;
          tokenCallBack(token);
        }
      }
    }

    const aiMessage = new AIMessage(aiMessageContent);
    await eventCallback({ event: 'ai_message', data: aiMessageContent });
    state.messages.push(aiMessage);
  }

  async disconnect(eventCallback) {
    await eventCallback({ event: 'end' });
  }
}

export default Agent;