@worktango/ai-assistant 0.0.21 → 0.0.24

package/src/frontend/components/AiAssistantComponent.test.tsx

@@ -1,11 +1,270 @@
 import React from "react";
 
-import { render } from "@kazoohr/test";
+import { ToastProvider } from "@kazoohr/confetti";
+import { act, fireEvent, render, screen } from "@kazoohr/test";
+
+import { useAiAssistantSession } from "../tools/useAiAssistantSession";
 
 import { AiAssistantComponent } from "./AiAssistantComponent";
 
+jest.mock("../tools/useAiAssistantSession", () => ({
+  useAiAssistantSession: jest.fn(),
+}));
+
+// Need to mock micromark because it uses ES modules
+jest.mock("micromark", () => ({
+  micromark: (text: string) => `<div>${text}</div>`,
+}));
+
 describe("AiAssistantComponent", () => {
-  it("should render", () => {
-    render(<AiAssistantComponent />);
+  const mockSendUserMessage = jest.fn();
+  const mockSetMessages = jest.fn();
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    jest.useFakeTimers();
+
+    (useAiAssistantSession as jest.Mock).mockReturnValue({
+      messages: [],
+      sendUserMessage: mockSendUserMessage,
+      isLoading: false,
+      debugInfo: {},
+      setMessages: mockSetMessages,
+    });
+  });
+
+  afterEach(() => {
+    jest.useRealTimers();
+  });
+
+  const renderWithToast = (ui: React.ReactNode) => {
+    return render(<ToastProvider>{ui}</ToastProvider>);
+  };
+
+  it("renders initial empty state", () => {
+    renderWithToast(<AiAssistantComponent />);
+
+    expect(
+      screen.getByText(
+        "No messages yet! Get started by asking a question below."
+      )
+    ).toBeInTheDocument();
+    expect(screen.getByText("Send")).toBeInTheDocument();
+    expect(screen.getByText("Clear current chat")).toBeInTheDocument();
+  });
+
+  it("handles message sending", async () => {
+    renderWithToast(<AiAssistantComponent />);
+
+    const textarea = screen.getByRole("textbox");
+    const sendButton = screen.getByText("Send");
+
+    // Clear the default text first
+    fireEvent.change(textarea, { target: { value: "" } });
+    expect(sendButton).toBeDisabled();
+
+    fireEvent.change(textarea, { target: { value: "This is a test message" } });
+    expect(sendButton).not.toBeDisabled();
+
+    // Send the message
+    fireEvent.click(sendButton);
+    expect(mockSendUserMessage).toHaveBeenCalledWith("This is a test message");
+  });
+
+  it("handles ctrl+enter submission", () => {
+    renderWithToast(<AiAssistantComponent />);
+
+    const textarea = screen.getByRole("textbox");
+    const sendButton = screen.getByText("Send");
+
+    fireEvent.change(textarea, { target: { value: "This is a test message" } });
+    expect(sendButton).not.toBeDisabled();
+
+    fireEvent.keyDown(textarea, { key: "Space", ctrlKey: false });
+
+    expect(mockSendUserMessage).not.toHaveBeenCalled();
+
+    fireEvent.keyDown(textarea, { key: "Enter", ctrlKey: true });
+
+    expect(mockSendUserMessage).toHaveBeenCalledWith("This is a test message");
+  });
+
+  it("displays loading state", () => {
+    (useAiAssistantSession as jest.Mock).mockReturnValue({
+      messages: [
+        {
+          role: "user",
+          parts: [{ type: "text", text: "Hello" }],
+        },
+      ],
+      sendUserMessage: mockSendUserMessage,
+      isLoading: true,
+      debugInfo: {},
+      setMessages: mockSetMessages,
+    });
+
+    const { baseElement } = renderWithToast(<AiAssistantComponent />);
+
+    expect(baseElement.innerHTML).toContain("loading");
+  });
+
+  it("display the LLM loading state", () => {
+    (useAiAssistantSession as jest.Mock).mockReturnValue({
+      messages: [
+        {
+          role: "user",
+          parts: [{ type: "text", text: "Hello" }],
+        },
+      ],
+      sendUserMessage: mockSendUserMessage,
+      isLoading: true,
+      debugInfo: {},
+      setMessages: mockSetMessages,
+    });
+
+    renderWithToast(<AiAssistantComponent />);
+
+    expect(screen.getByText("Thinking...")).toBeInTheDocument();
+  });
+
+  it("displays messages correctly", () => {
+    (useAiAssistantSession as jest.Mock).mockReturnValue({
+      messages: [
+        {
+          role: "user",
+          parts: [{ type: "text", text: "Hello" }],
+        },
+        {
+          role: "assistant",
+          parts: [{ type: "text", text: "Hi there!" }],
+        },
+        {
+          role: "assistant",
+          parts: [
+            { type: "toolCall", toolName: "testTool", args: { test: true } },
+            {
+              type: "toolResult",
+              toolName: "testTool",
+              result: { success: true },
+            },
+          ],
+        },
+        {
+          role: "tool",
+          parts: [
+            {
+              type: "UNKNOWN!",
+              toolName: "testTool",
+              result: { success: true },
+            },
+          ],
+        },
+      ],
+      sendUserMessage: mockSendUserMessage,
+      isLoading: false,
+      debugInfo: {},
+      setMessages: mockSetMessages,
+    });
+
+    const { baseElement } = renderWithToast(<AiAssistantComponent />);
+    const html = baseElement.innerHTML;
+
+    expect(html).toContain("Hi there!");
+  });
+
+  it("handles debug mode toggle", () => {
+    (useAiAssistantSession as jest.Mock).mockReturnValue({
+      messages: [
+        { role: "user", parts: [{ type: "text", text: "Hello world!" }] },
+      ],
+      sendUserMessage: mockSendUserMessage,
+      isLoading: false,
+      debugInfo: { test: true },
+      setMessages: mockSetMessages,
+    });
+
+    renderWithToast(<AiAssistantComponent />);
+
+    // Initially debug info is hidden
+    const debugButton = screen.getByTitle("Show debug info");
+    expect(debugButton).toBeInTheDocument();
+
+    // Show debug info
+    fireEvent.click(debugButton);
+
+    // Find the tabs
+    const tabs = screen.getAllByRole("tab");
+    expect(tabs[0]).toHaveTextContent("History");
+    expect(tabs[1]).toHaveTextContent("Debug");
+
+    // Click the messages tab
+    fireEvent.click(tabs[0]);
+
+    expect(screen.getByText("Hello world!")).toBeInTheDocument();
+
+    // Click the debug tab
+    fireEvent.click(tabs[1]);
+
+    // Check that the debug info is displayed
+    expect(screen.getByTestId("debug-textarea")).toHaveTextContent("test");
+  });
+
+  it("handles expand/collapse functionality", () => {
+    (useAiAssistantSession as jest.Mock).mockReturnValue({
+      messages: [{ role: "user", parts: [{ type: "text", text: "test" }] }],
+      sendUserMessage: mockSendUserMessage,
+      isLoading: false,
+      debugInfo: {},
+      setMessages: mockSetMessages,
+    });
+
+    renderWithToast(<AiAssistantComponent />);
+
+    const collapseButton = screen.getByTitle("Collapse chat window");
+    expect(collapseButton).toBeInTheDocument();
+    fireEvent.click(collapseButton);
+
+    const expandButton = screen.getByTitle("Expand chat window");
+    expect(expandButton).toBeInTheDocument();
+    fireEvent.click(expandButton);
+
+    expect(screen.getByTitle("Collapse chat window")).toBeInTheDocument();
+  });
+
+  it("handles clearing chat", () => {
+    renderWithToast(<AiAssistantComponent />);
+
+    fireEvent.click(screen.getByText("Clear current chat"));
+    expect(mockSetMessages).toHaveBeenCalledWith([]);
+  });
+
+  it("handles errors", () => {
+    renderWithToast(<AiAssistantComponent />);
+
+    const error = new Error("Test error");
+    const { onError } = (useAiAssistantSession as jest.Mock).mock.calls[0][0];
+
+    act(() => {
+      onError(error);
+    });
+
+    expect(screen.getByText("Test error")).toBeInTheDocument();
+  });
+
+  it("handles scroll to bottom on message send", () => {
+    const mockRef = { current: { scrollTop: 0, scrollHeight: 100 } };
+    jest.spyOn(React, "useRef").mockReturnValue(mockRef);
+
+    renderWithToast(<AiAssistantComponent />);
+
+    const { onBeforeLlmCall } = (useAiAssistantSession as jest.Mock).mock
+      .calls[0][0];
+
+    act(() => {
+      onBeforeLlmCall();
+      jest.advanceTimersByTime(10); // Advance past the setTimeout
+    });
+
+    expect(mockRef.current.scrollTop).toBe(mockRef.current.scrollHeight);
   });
 });

package/src/frontend/components/AiAssistantComponent.tsx

@@ -1,18 +1,246 @@
 import React from "react";
 
-import { Card, TextArea } from "@kazoohr/confetti";
-import { setState } from "@kazoohr/helpers";
+import {
+  BRANDING_CLASSES,
+  Button,
+  Card,
+  Flex,
+  Illustration,
+  Loader,
+  PageHeader,
+  Spacer,
+  Tab,
+  Tabs,
+  Text,
+  TextArea,
+  useToast,
+} from "@kazoohr/confetti";
+import {
+  composeSimpleFunctions,
+  getClosureFor,
+  setState,
+  setStateValue,
+} from "@kazoohr/helpers";
+import { ToolDeclaration } from "@kazoohr/llm/types";
+
+import { useAiAssistantSession } from "../tools/useAiAssistantSession";
+
+import styles from "./AiAssistantComponent.scss";
+import { AssistantMessage } from "./messageRenderers/AssistantMessage";
+import { UserMessage } from "./messageRenderers/UserMessage";
 
 /**
- * This is a placeholder 😎
+ * Our AI Assistant component, for managing the LLM session and displaying a
+ * single chat session with the LLM.
  */
-export function AiAssistantComponent() {
-  // TODO: AI assistant go here.
-  const [value, setValue] = React.useState("");
+// eslint-disable-next-line complexity
+export function AiAssistantComponent({ tools }: { tools?: ToolDeclaration[] }) {
+  const { error: errorToast } = useToast();
+  const [showDebugInfo, setShowDebugInfo] = React.useState(false);
+  const [prompt, setPrompt] = React.useState(
+    'This is a test! Give me a "Hello world! From <your model>". Then, use your tool calls to tell me about myself. Be comprehensive in your answer.'
+  );
+  const textareaRef = React.useRef<HTMLTextAreaElement>(null);
+  const scrollWindowRef = React.useRef<HTMLDivElement>(null);
+  const scrollWindowToBottom = React.useCallback(
+    // istanbul ignore next - UI-related stuff
+    () => {
+      setTimeout(() => {
+        if (scrollWindowRef.current) {
+          scrollWindowRef.current.scrollTop =
+            scrollWindowRef.current.scrollHeight;
+        }
+      }, 10);
+    },
+    []
+  );
+  const { messages, sendUserMessage, isLoading, debugInfo, setMessages } =
+    useAiAssistantSession({
+      onBeforeLlmCall: scrollWindowToBottom,
+      onAfterLlmCall: composeSimpleFunctions(
+        setStateValue(setPrompt, ""),
+        scrollWindowToBottom,
+        textareaRef.current?.focus
+      ),
+      onError: (error) => errorToast(error.message),
+      onChunkAdded: scrollWindowToBottom,
+      tools,
+    });
+  const [fullHeight, setFullHeight] = React.useState(true);
+  const [tab, setTab] = React.useState<"messages" | "debug">("messages");
+  React.useEffect(() => {
+    // istanbul ignore else - noop
+    if (textareaRef.current) {
+      // Ctrl + enter submits the form
+      const handleKeyDown = (event: KeyboardEvent) => {
+        if (event.ctrlKey && event.key === "Enter") {
+          event.preventDefault();
+          sendUserMessage(prompt);
+        }
+      };
+      textareaRef.current.addEventListener("keydown", handleKeyDown);
+      return () => {
+        textareaRef.current?.removeEventListener("keydown", handleKeyDown);
+      };
+    }
+  }, [prompt, sendUserMessage]);
+
+  const canSubmit = prompt.trim().length > 5;
+
   return (
-    <Card background="confetti">
-      Hello world! This is the AI assistant.
-      <TextArea id="id" value={value} onChange={setState(setValue)} />
-    </Card>
+    <div>
+      <style type="text/css">
+        {`
+          .streaming-text-resizing-container {
+            overflow-y: hidden;
+            transition: height 0.1s ease-in-out;
+          }
+
+          .markdown {
+            line-height: 1.4;
+          }
+
+          .markdown * + p {
+            margin-top: 12px;
+          }
+
+          .markdown ul, .markdown ol {
+            margin-top: 12px;
+            margin-left: 24px;
+          }
+
+          .scroll-window {
+            padding: 12px 4px 16px 0;
+            overflow-y: scroll;
+            overflow-x: hidden;
+            height: min(100vh, max(50vh, 400px));
+          }
+`}
+      </style>
+      <PageHeader
+        header="AI Assistant"
+        withBottomBorder={false}
+        renderRightContent={() => (
+          <Flex gap="12px">
+            <Button
+              icon={fullHeight ? "caret_up" : "caret_dn"}
+              onClick={setStateValue(setFullHeight, !fullHeight)}
+              title={fullHeight ? "Collapse chat window" : "Expand chat window"}
+              variant="transparent"
+              size="extra_small"
+            >
+              {fullHeight ? "Collapse" : "Expand"}
+            </Button>
+            <Button
+              icon="eye"
+              onClick={composeSimpleFunctions(
+                setStateValue(setShowDebugInfo, !showDebugInfo),
+                setStateValue(setTab, "messages")
+              )}
+              title={showDebugInfo ? "Hide debug info" : "Show debug info"}
+              variant="transparent"
+              size="extra_small"
+            >
+              {showDebugInfo ? "Debugging" : "Debug"}
+            </Button>
+            <Button
+              icon="close"
+              variant="transparent"
+              onClick={setStateValue(setMessages, [])}
+              size="extra_small"
+            >
+              Clear current chat
+            </Button>
+          </Flex>
+        )}
+      />
+      {showDebugInfo && (
+        <>
+          <Tabs value={tab} onChange={setState(setTab)}>
+            <Tab value="messages" data-testid="messages-tab">
+              History
+            </Tab>
+            <Tab value="debug" data-testid="debug-tab">
+              Debug
+            </Tab>
+          </Tabs>
+          <Spacer />
+        </>
+      )}
+      {tab === "debug" && (
+        <Card>
+          <textarea
+            value={JSON.stringify(debugInfo, null, 2)}
+            className={styles["ai-assistant-debug-textarea"]}
+            readOnly
+            data-testid="debug-textarea"
+          />
+        </Card>
+      )}
+      {tab === "messages" &&
+        (messages.length === 0 ? (
+          <Flex direction="column" alignItems="center" gap="12px">
+            <Illustration illustration="feedback" size={196} />
+            <Text slim size="large">
+              No messages yet! Get started by asking a question below.
+            </Text>
+          </Flex>
+        ) : (
+          <>
+            <div
+              className={fullHeight ? "" : "scroll-window"}
+              ref={scrollWindowRef}
+            >
+              <Flex direction="column" gap="16px" justifyContent="flex-end">
+                {messages.map((message, index) => {
+                  if (message.role === "user") {
+                    return <UserMessage key={index} message={message} />;
+                  }
+                  return <AssistantMessage key={index} message={message} />;
+                })}
+                {isLoading && (
+                  <div style={{ height: "100px" }}>
+                    <Loader />
+                  </div>
+                )}
+              </Flex>
+            </div>
+            <Spacer size="large" />
+          </>
+        ))}
+      <Spacer />
+      <Card
+        className={BRANDING_CLASSES.PRIMARY_COLOR_BACKGROUND}
+        background="two-tone"
+        data-testid="toggle-panel-controls"
+      >
+        <TextArea
+          id="id"
+          value={prompt}
+          onChange={setState(setPrompt)}
+          fullWidth
+          className={styles["ai-assistant-textarea"]}
+          placeholder="Ask a question..."
+          inputRef={textareaRef}
+          autoFocus
+          disabled={isLoading}
+        />
+        <Spacer size="small" />
+        <Flex justifyContent="flex-end" gap="8px">
+          <Button
+            onClick={getClosureFor(sendUserMessage, prompt)}
+            disabled={!canSubmit || isLoading}
+            variant="feature_inverted"
+            size="large"
+            icon="people"
+            style={{
+              cursor: !canSubmit || isLoading ? "not-allowed" : "pointer",
+            }}
+          >
+            {isLoading ? "Thinking..." : "Send"}
+          </Button>
+        </Flex>
+      </Card>
+    </div>
   );
 }

package/src/frontend/components/messageRenderers/AssistantMessage.test.tsx

@@ -0,0 +1,154 @@
+import React from "react";
+
+import { render, screen } from "@kazoohr/test";
+
+import { type ChatItemWithTimestamp } from "../../tools/useAiAssistantSession";
+
+import { AssistantMessage } from "./AssistantMessage";
+
+jest.mock("micromark", () => ({
+  micromark: (text: string) => text,
+}));
+
+describe("AssistantMessage", () => {
+  const mockTimestamp = new Date("2024-01-01T12:00:00Z");
+
+  it("renders text message correctly", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "model",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "text",
+          text: "Hello world",
+        },
+      ],
+    };
+
+    render(<AssistantMessage message={message} />);
+
+    expect(screen.getByText("assistant")).toBeInTheDocument();
+    expect(screen.getByText("Hello world")).toBeInTheDocument();
+    expect(
+      screen.getByText(mockTimestamp.toLocaleString())
+    ).toBeInTheDocument();
+  });
+
+  it("renders markdown text correctly", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "model",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "text",
+          text: "# Hello\n\n**Bold text**",
+        },
+      ],
+    };
+
+    render(<AssistantMessage message={message} />);
+
+    const content = document.querySelector(".markdown");
+    expect(content).toBeInTheDocument();
+    expect(content?.innerHTML).toContain("# Hello\n\n**Bold text**");
+  });
+
+  it("renders tool calls correctly", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "model",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "toolCall",
+          toolName: "testTool",
+          args: { test: true },
+        },
+      ],
+    };
+
+    render(<AssistantMessage message={message} />);
+
+    expect(screen.getByTitle("Tool call: testTool")).toBeInTheDocument();
+    expect(screen.getByText("Fetch data")).toBeInTheDocument();
+  });
+
+  it("renders both text and tool calls together", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "model",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "text",
+          text: "Hello world",
+        },
+        {
+          type: "toolCall",
+          toolName: "testTool",
+          args: { test: true },
+        },
+      ],
+    };
+
+    render(<AssistantMessage message={message} />);
+
+    expect(screen.getByText("Hello world")).toBeInTheDocument();
+    expect(screen.getByTitle("Tool call: testTool")).toBeInTheDocument();
+    expect(screen.getByText("Fetch data")).toBeInTheDocument();
+  });
+
+  it("renders multiple tool calls", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "model",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "toolCall",
+          toolName: "tool1",
+          args: { test: true },
+        },
+        {
+          type: "toolCall",
+          toolName: "tool2",
+          args: { test: true },
+        },
+      ],
+    };
+
+    render(<AssistantMessage message={message} />);
+
+    expect(screen.getByTitle("Tool call: tool1")).toBeInTheDocument();
+    expect(screen.getByTitle("Tool call: tool2")).toBeInTheDocument();
+    expect(screen.getAllByText("Fetch data")).toHaveLength(2);
+  });
+
+  it("returns null for empty message parts", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "model",
+      timestamp: mockTimestamp,
+      parts: [],
+    };
+
+    const { container } = render(<AssistantMessage message={message} />);
+    expect(container.firstChild).toBeNull();
+  });
+
+  it("aggregates multiple text parts", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "model",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "text",
+          text: "Hello",
+        },
+        {
+          type: "text",
+          text: " world",
+        },
+      ],
+    };
+
+    render(<AssistantMessage message={message} />);
+    expect(screen.getByText("Hello world")).toBeInTheDocument();
+  });
+});