app/page.tsx

"use client";

import styles from './page.module.css';
import { useEffect, useState } from 'react';
import { useChat } from 'ai/react';
import { FunctionCallHandler, Message, nanoid } from 'ai';
import ReactMarkdown from "react-markdown";
import { Bot, User } from "lucide-react";
import { toast } from 'sonner';
import { FunctionIcon } from './icons';
import { updateBackground } from './util';
import Input from './input';

const Page: React.FC = () => {
  useEffect(() => {
    updateBackground();
    const interval = setInterval(updateBackground, 600);

    return () => clearInterval(interval);
  }, []);
  const functionCallHandler: FunctionCallHandler = async (
    chatMessages,
    functionCall,
  ) => {
    let result;
    const { name, arguments: args } = functionCall;
    const response = await fetch("/api/functions", {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
        args: args,
        name: name
      })
    } as any);

    if (!response.ok) {
      const errorText = await response.text();
      toast.error(`Something went wrong: ${errorText}`);
      return;
    }    

    result = await response.text();

    return {
      messages: [
        ...chatMessages,
        {
          id: nanoid(),
          name: functionCall.name,
          role: "function" as const,
          content: result,
        },
      ],
    };
  };
  const { messages, input, setInput, handleSubmit, isLoading } = useChat({
    experimental_onFunctionCall: functionCallHandler,
    onError: (error: any) => {
      console.log(error);
    },
  });
  const [isExpanded, setIsExpanded] = useState(false);
  const toggleExpand = () => {
    setIsExpanded(!isExpanded);
  };

  const roleUIConfig: {
    [key: string]: {
      avatar: JSX.Element;
      bgColor: string;
      avatarColor: string;
      // eslint-disable-next-line no-unused-vars
      dialogComponent: (message: Message) => JSX.Element;
    };
  } = {
    user: {
      avatar: <User width={20} />,
      bgColor: "bg-white",
      avatarColor: "bg-black",
      dialogComponent: (message: Message) => (
        <ReactMarkdown
          className=""
          components={{
            a: (props) => (
              <a {...props} target="_blank" rel="noopener noreferrer" />
            ),
          }}
        >
          {message.content}
        </ReactMarkdown>
      ),
    },
    assistant: {
      avatar: <Bot width={20} />,
      bgColor: "bg-gray-100",
      avatarColor: "bg-green-500",
      dialogComponent: (message: Message) => (
        <ReactMarkdown
          className=""
          components={{
            a: (props) => (
              <a {...props} target="_blank" rel="noopener noreferrer" />
            ),
          }}
        >
          {message.content}
        </ReactMarkdown>
      ),
    },
    function: {
      avatar: <div className="cursor-pointer" onClick={toggleExpand}><FunctionIcon /></div>,
      bgColor: "bg-gray-200",
      avatarColor: "bg-blue-500",
      dialogComponent: (message: Message) => {
        return (
          <div className="flex flex-col">
            {isExpanded && (
              <div className="py-1">{message.content}</div>
            )}
          </div>
        );
      },
    }    
  };


  return (
    <main className={styles.main}>
      <div id="bg" className={styles.background}></div>
      <div className={styles.messages}>
        {messages.length > 0 ? (
          messages.map((message, i) => {
            const messageClass = `${styles.message} ${message.role === 'user' ? styles['message-user'] : ''}`;
            return (
              <div key={i} className={messageClass} style={{ display: 'flex', alignItems: 'center' }}>
                <div className={styles.avatar}>
                  {roleUIConfig[message.role].avatar}
                </div>
                <div style={{ flex: 1 }}>
                  {message.content === "" && message.function_call != undefined ? (
                    typeof message.function_call === "object" ? (
                      <div style={{ display: 'flex', flexDirection: 'column' }}>
                        <div>
                          Using{" "}
                          <span className="font-bold">
                            {message.function_call.name}
                          </span>{" "}
                          ...
                        </div>
                        <div>
                          {message.function_call.arguments}
                        </div>
                      </div>
                    ) : (
                      <div className="function-call">{message.function_call}</div>
                    )
                  ) : (
                    roleUIConfig[message.role].dialogComponent(message)
                  )}
                </div>
              </div>
            );
          })
        ) : null}
      </div>
      <Input handleSubmit={handleSubmit as any} setInput={setInput} input={input} />
    </main>
  );
}

export default Page;