LLM Chatbot with Multi-Agent System

import { Mastra } from '@mastra/core/mastra';
import { MongoDBStore } from '@mastra/mongodb';
import { chatRoute } from '@mastra/ai-sdk';
import { recruitingAgent } from './agents/recruiting-agent';

export const mastra = new Mastra({
  agents: { recruitingAgent },
  storage: new MongoDBStore({
    url: process.env.MONGODB_URL,
    dbName: process.env.MONGODB_DATABASE,
  }),
  server: {
    apiRoutes: [
      chatRoute({
        path: '/recruiter/recruiting-chat',
        agent: 'recruitingAgent',
      }),
    ],
  },
});

// Root Agent - Handles routing and simple conversations
export const recruitingAgent = new Agent({
  name: 'Recruiting Agent',
  instructions: `You're a routing agent for JobHai recruiters. 
  Your role is to handle simple conversations directly and 
  route job-related queries to the appropriate tool.
  
  Handle directly (no tool):
  - Greetings and acknowledgments
  - General conversation
  
  Route to handleRecruiterQueryTool:
  - Job posting issues
  - Account/job verification questions
  - Document issues
  - Any query requiring job-specific information`,
  model: azure('gpt-4.1-mini'),
  tools: {
    handleRecruiterQueryTool, // Tool that calls sub-agent
  },
  memory: createMemory(),
});

// Tool that invokes sub-agent
export const handleRecruiterQueryTool = createTool({
  id: 'handle-recruiter-query',
  description: 'Handles recruiter queries about jobs, accounts, verification',
  execute: async ({ context, runtimeContext }) => {
    const { query } = context;
    
    // Extract thread/resource from runtime context
    const thread = runtimeContext?.get('thread');
    const resource = runtimeContext?.get('resource');
    
    // Call sub-agent with same thread/resource for context continuity
    const handlerAgent = new Agent({
      name: 'Recruiter Query Handler Agent',
      tools: {
        getJobDetailsTool,
        checkJobStatusTool,
        getUserProfileTool,
        checkVerificationStatusTool,
      },
      memory: createMemory(),
    });
    
    const response = await handlerAgent.generate(query, {
      memory: { thread, resource },
    });
    
    return { response: response.text };
  },
});

// Root Agent calls tool, tool invokes sub-agent
export const handleRecruiterQueryTool = createTool({
  id: 'handle-recruiter-query',
  description: 'Handles recruiter queries by calling sub-agent',
  execute: async ({ context, runtimeContext, writer }) => {
    const { query } = context;
    
    // Extract thread/resource from runtime context
    const thread = runtimeContext?.get('thread');
    const resource = runtimeContext?.get('resource');
    
    // Create sub-agent instance
    const recruiterQueryHandlerAgent = new Agent({
      name: 'Recruiter Query Handler Agent',
      instructions: `You handle recruiter queries about jobs, accounts, 
      and verification. Use tools to get real-time data.`,
      tools: {
        getJobDetailsTool,
        checkJobStatusTool,
        getUserProfileTool,
        checkVerificationStatusTool,
      },
      memory: createMemory(),
    });
    
    // Call sub-agent with shared thread/resource for context continuity
    const response = await recruiterQueryHandlerAgent.stream(query, {
      memory: {
        thread: thread,
        resource: resource,
      },
    });
    
    // Stream response back to root agent
    if (writer) {
      await response.fullStream.pipeTo(writer);
    }
    
    return {
      response: await response.text,
      agentName: 'Recruiter Query Handler Agent',
    };
  },
});

import { useChat } from "@ai-sdk/react";
import { DefaultChatTransport } from "ai";

export function useChatbotTransport({
  threadIdRef,
  resourceIdRef,
}) {
  const { messages, sendMessage, status } = useChat({
    transport: new DefaultChatTransport({
      api: API_ENDPOINT,
      // Custom request preparation
      prepareSendMessagesRequest: ({ messages: msgs }) => {
        return {
          messages: msgs,
          memory: {
            thread: threadIdRef.current,
            resource: resourceIdRef.current,
          },
        };
      },
    }),
    
    // Handle tool calls
    onToolCall({ toolCall }) {
      // Client-side tool handling
      if (toolCall.toolName === "colorChangeTool") {
        changeBgColor(toolCall.input.color);
      }
    },
  });

  return { messages, sendMessage, status };
}

// MessageBubble Component - Renders different message parts
export function MessageBubble({ message, sendMessage }) {
  return (
    <div className="message-bubble">
      {message.parts?.map((part, index) => {
        // Text parts - formatted text
        if (isTextPart(part)) {
          return <FormattedText text={part.text} />;
        }
        
        // Tool invocation parts
        if (part.type?.startsWith("tool-")) {
          const toolName = getToolName(part);
          
          // Show tool status
          if (part.state === "call") {
            return (
              <ToolInvocation part={part} />
              // Shows: "Processing your request..." with pulsing dot
            );
          }
          
          // Human-in-the-loop: Interactive selectors
          if (isTool(part, "askJobTool") && isToolWaitingForInput(part)) {
            return (
              <JobSelector 
                jobs={part.input.jobs}
                onSelect={(jobId) => sendMessage({ text: jobId })}
              />
            );
          }
          
          if (isTool(part, "askCityTool") && isToolWaitingForInput(part)) {
            return (
              <CitySelector 
                cities={CITIES}
                onSelect={(city) => sendMessage({ text: city })}
              />
            );
          }
        }
        
        // Nested agent parts (from tool calling agent)
        if (isNestedAgentPart(part)) {
          return <NestedAgentDisplay part={part} />;
        }
      })}
    </div>
  );
}

// ToolInvocation Component - User-friendly tool status
export function ToolInvocation({ part }) {
  const toolName = getToolName(part);
  const userMessage = getToolMessage(toolName);
  
  // Show simple message during tool execution
  return (
    <div className="text-xs text-slate-500 italic flex items-center gap-1.5">
      <span className="w-1.5 h-1.5 bg-slate-400 rounded-full animate-pulse"></span>
      {userMessage} {/* e.g., "Processing your request..." */}
    </div>
  );
}