@worktango/ai-assistant 0.0.21 → 0.0.22

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();
+  });
+});

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

@@ -0,0 +1,104 @@
+import React from "react";
+
+import { micromark } from "micromark";
+
+import {
+  Flex,
+  Spacer,
+  UtilityText,
+  Text,
+  Pill,
+  CardPadding,
+  Icon,
+} from "@kazoohr/confetti";
+
+import { type ChatItemWithTimestamp } from "../../tools/useAiAssistantSession";
+
+export function AssistantMessage({
+  message,
+}: {
+  message: ChatItemWithTimestamp;
+}) {
+  const aggregatedToolCalls = message.parts.filter(
+    (part) => part.type === "toolCall"
+  );
+  const aggregatedTextParts = message.parts
+    .filter((part) => part.type === "text")
+    .map((text) => text.text)
+    .join("");
+
+  const streamingTextResizingContainerRef = React.useRef<HTMLDivElement>(null);
+  const streamingTextRef = React.useRef<HTMLDivElement>(null);
+
+  const componentsToRender: React.ReactNode[] = [];
+
+  React.useEffect(() => {
+    if (streamingTextRef.current && streamingTextResizingContainerRef.current) {
+      const textHeight = streamingTextRef.current.clientHeight;
+      streamingTextResizingContainerRef.current.style.height = `${textHeight}px`;
+    }
+  }, [aggregatedTextParts]);
+
+  if (aggregatedTextParts) {
+    componentsToRender.push(
+      <div
+        ref={streamingTextResizingContainerRef}
+        className="streaming-text-resizing-container"
+        key="text"
+      >
+        <div ref={streamingTextRef}>
+          <CardPadding>
+            <Flex gap="16px" alignItems="center">
+              <Icon icon="teams" />
+              <Text bold size="small">
+                assistant
+              </Text>
+              <UtilityText size="extraSmall">
+                {message.timestamp?.toLocaleString()}
+              </UtilityText>
+            </Flex>
+            <Spacer />
+            <div
+              className="markdown"
+              dangerouslySetInnerHTML={{
+                __html: micromark(aggregatedTextParts),
+              }}
+            />
+          </CardPadding>
+        </div>
+      </div>
+    );
+  }
+
+  if (aggregatedToolCalls.length > 0) {
+    componentsToRender.push(
+      <div key="tools">
+        {aggregatedToolCalls.map((toolCall) => {
+          return (
+            <CardPadding key={toolCall.toolName}>
+              <div
+                style={{ textAlign: "right" }}
+                title={`Tool call: ${toolCall.toolName}`}
+              >
+                <div style={{ display: "inline-block" }}>
+                  <Pill
+                    icon="info"
+                    size="large"
+                    text="Fetch data"
+                    variant="outline"
+                  />
+                </div>
+              </div>
+            </CardPadding>
+          );
+        })}
+      </div>
+    );
+  }
+
+  if (componentsToRender.length === 0) {
+    return null;
+  }
+
+  return <>{componentsToRender}</>;
+}

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

@@ -0,0 +1,103 @@
+import React from "react";
+
+import { render, screen } from "@kazoohr/test";
+
+import { type ChatItemWithTimestamp } from "../../tools/useAiAssistantSession";
+
+import { UserMessage } from "./UserMessage";
+
+jest.mock("micromark", () => ({
+  micromark: (text: string) => text,
+}));
+
+describe("UserMessage", () => {
+  const mockTimestamp = new Date("2024-01-01T12:00:00Z");
+
+  it("renders text message correctly", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "user",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "text",
+          text: "Hello world",
+        },
+      ],
+    };
+
+    render(<UserMessage message={message} />);
+
+    expect(screen.getByText("you")).toBeInTheDocument();
+    expect(screen.getByText("Hello world")).toBeInTheDocument();
+    expect(
+      screen.getByText(mockTimestamp.toLocaleString())
+    ).toBeInTheDocument();
+  });
+
+  it("renders markdown text correctly", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "user",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "text",
+          text: "# Hello\n\n**Bold text**",
+        },
+      ],
+    };
+
+    render(<UserMessage message={message} />);
+
+    const content = document.querySelector(".markdown");
+    expect(content).toBeInTheDocument();
+    expect(content?.innerHTML).toContain("# Hello\n\n**Bold text**");
+  });
+
+  it("returns null for empty message parts", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "user",
+      timestamp: mockTimestamp,
+      parts: [],
+    };
+
+    const { container } = render(<UserMessage message={message} />);
+    expect(container.firstChild).toBeNull();
+  });
+
+  it("aggregates multiple text parts", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "user",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "text",
+          text: "Hello",
+        },
+        {
+          type: "text",
+          text: " world",
+        },
+      ],
+    };
+
+    render(<UserMessage message={message} />);
+    expect(screen.getByText("Hello world")).toBeInTheDocument();
+  });
+
+  it("ignores non-text parts", () => {
+    const message: ChatItemWithTimestamp = {
+      role: "user",
+      timestamp: mockTimestamp,
+      parts: [
+        {
+          type: "toolCall",
+          toolName: "testTool",
+          args: { test: true },
+        },
+      ],
+    };
+
+    const { container } = render(<UserMessage message={message} />);
+    expect(container.firstChild).toBeNull();
+  });
+});

package/src/frontend/components/messageRenderers/UserMessage.tsx

@@ -0,0 +1,43 @@
+import React from "react";
+
+import { micromark } from "micromark";
+
+import { Card, Flex, Spacer, UtilityText, Text, Icon } from "@kazoohr/confetti";
+
+import { ChatItemWithTimestamp } from "../../tools/useAiAssistantSession";
+
+export function UserMessage({ message }: { message: ChatItemWithTimestamp }) {
+  const aggregatedTextParts = message.parts
+    .filter((part) => part.type === "text")
+    .map((text) => text.text)
+    .join("");
+
+  const componentsToRender: React.ReactNode[] = [];
+
+  if (aggregatedTextParts) {
+    componentsToRender.push(
+      <Card key="text" noShadow withBorder>
+        <Flex gap="16px" alignItems="center">
+          <Icon icon="feedback" />
+          <Text bold size="small">
+            you
+          </Text>
+          <UtilityText size="extraSmall">
+            {message.timestamp?.toLocaleString()}
+          </UtilityText>
+        </Flex>
+        <Spacer />
+        <div
+          className="markdown"
+          dangerouslySetInnerHTML={{ __html: micromark(aggregatedTextParts) }}
+        />
+      </Card>
+    );
+  }
+
+  if (componentsToRender.length === 0) {
+    return null;
+  }
+
+  return <>{componentsToRender}</>;
+}

package/src/frontend/reactComponent.ts

@@ -1,3 +1,5 @@
 // istanbul ignore file - Main package entrypoint
 
 export { AiAssistant } from "./AiAssistant";
+
+export * as tools from "./tools/llm-tools";

package/src/frontend/tools/generatedSdks.ts

@@ -0,0 +1,43 @@
+// istanbul ignore file - light wrapper around the generated SDK
+
+import { DocumentNode, print } from "@kazoohr/graphql";
+
+import { getSdk as getStitchedSdk } from "../../../generated/stitchedSdk";
+
+export const stitchedSdk = {
+  ...getStitchedSdk((query, variables) => {
+    return fetcher({
+      url: "/graphql",
+      query,
+      variables,
+    });
+  }),
+};
+
+export async function fetcher({
+  query,
+  url,
+  variables,
+}: {
+  url: string;
+  query: DocumentNode;
+  variables?: any;
+}): Promise<any> {
+  const response = await fetch(url, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      Accept: "application/graphql-response+json",
+    },
+    body: JSON.stringify({
+      query: print(query),
+      variables,
+    }),
+  });
+
+  if (!response.ok) {
+    throw new Error("Network response returned non-OK status");
+  }
+
+  return await response.json();
+}

package/src/frontend/tools/llm-tools/getCurrentUser.test.ts

@@ -0,0 +1,100 @@
+import { getCurrentUserTool, resetCurrentUser } from "./getCurrentUser";
+
+const mockFetch = jest.fn();
+global.fetch = mockFetch;
+
+const mockConsoleLog = jest.fn();
+global.console.log = mockConsoleLog;
+
+describe("getCurrentUserTool", () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    resetCurrentUser();
+  });
+
+  it("fetches and caches user data", async () => {
+    const mockUserData = {
+      data: {
+        directory_me: {
+          id: "123",
+          name: "Test User",
+        },
+      },
+    };
+
+    mockFetch.mockResolvedValueOnce({
+      json: () => Promise.resolve(mockUserData),
+    });
+
+    // First call should make the fetch request
+    const result1 = await getCurrentUserTool.fulfill({});
+    expect(result1).toBe(JSON.stringify(mockUserData));
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+    expect(mockFetch).toHaveBeenCalledWith("/graphql", {
+      method: "POST",
+      credentials: "include",
+      headers: {
+        "Content-Type": "application/json",
+      },
+      body: expect.any(String),
+    });
+
+    // Second call should use cached value
+    const result2 = await getCurrentUserTool.fulfill({});
+    expect(result2).toBe(JSON.stringify(mockUserData));
+    expect(mockFetch).toHaveBeenCalledTimes(1); // Should not make another fetch
+  });
+
+  it("handles fetch errors", async () => {
+    const errorMessage = "Network error";
+    const error = new Error(errorMessage);
+    mockFetch.mockRejectedValueOnce(error);
+
+    const result = await getCurrentUserTool.fulfill({});
+    expect(result).toBe(JSON.stringify({ error: errorMessage }));
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+    expect(mockConsoleLog).not.toHaveBeenCalled();
+
+    // Next call should try fetching again since error cleared the cache
+    const mockUserData = { data: { directory_me: { id: "123" } } };
+    mockFetch.mockResolvedValueOnce({
+      json: () => Promise.resolve(mockUserData),
+    });
+    const result2 = await getCurrentUserTool.fulfill({});
+    expect(result2).toBe(JSON.stringify(mockUserData));
+    expect(mockFetch).toHaveBeenCalledTimes(2);
+  });
+
+  it("handles non-Error fetch rejections", async () => {
+    mockFetch.mockRejectedValueOnce("Non-error rejection");
+
+    const result = await getCurrentUserTool.fulfill({});
+    expect(result).toBe(JSON.stringify({ error: "Non-error rejection" }));
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+  });
+
+  it("handles JSON parse errors", async () => {
+    mockFetch.mockResolvedValueOnce({
+      json: () => Promise.reject(new Error("Invalid JSON")),
+    });
+
+    const result = await getCurrentUserTool.fulfill({});
+    expect(result).toBe(JSON.stringify({ error: "Invalid JSON" }));
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+  });
+
+  it("handles non-Error JSON parse errors", async () => {
+    mockFetch.mockResolvedValueOnce({
+      json: () => Promise.reject("Invalid JSON string"),
+    });
+
+    const result = await getCurrentUserTool.fulfill({});
+    expect(result).toBe(JSON.stringify({ error: "Invalid JSON string" }));
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+  });
+
+  it("has correct tool metadata", () => {
+    expect(getCurrentUserTool.name).toBe("getCurrentUser");
+    expect(getCurrentUserTool.description).toBe("Get the current user");
+  });
+});

package/src/frontend/tools/llm-tools/getCurrentUser.ts

@@ -0,0 +1,158 @@
+import { handleForegroundPromise } from "@kazoohr/helpers";
+import { createTool } from "@kazoohr/llm/createTool";
+
+export let currentUser: Promise<string> | undefined = undefined;
+
+export function resetCurrentUser() {
+  currentUser = undefined;
+}
+
+export const getCurrentUserTool = createTool({
+  name: "getCurrentUser",
+  description: "Get the current user",
+  fulfill: async function fulfillGetCurrentUser() {
+    if (currentUser) {
+      return await currentUser;
+    }
+    currentUser = new Promise((resolve) => {
+      handleForegroundPromise(async (): Promise<unknown> => {
+        try {
+          // First, introspect `/api` endpoint to get the schema
+          const response = await fetch("/graphql", {
+            method: "POST",
+            credentials: "include",
+            headers: {
+              "Content-Type": "application/json",
+            },
+            body: JSON.stringify({
+              query: currentUserQuery,
+            }),
+          }).catch((error: Error | string) => error);
+          if (response instanceof Error || typeof response === "string") {
+            currentUser = undefined;
+            resolve(
+              JSON.stringify({
+                error: response instanceof Error ? response.message : response,
+              })
+            );
+            return;
+          }
+          const schema = await response.json();
+          resolve(JSON.stringify(schema));
+        } catch (error) {
+          currentUser = undefined;
+          resolve(
+            JSON.stringify({
+              error: error instanceof Error ? error.message : String(error),
+            })
+          );
+        }
+      });
+    });
+    return currentUser;
+  },
+});
+
+const currentUserQuery = `
+query GetCurrentUser {
+  directory_me {
+    avatar
+    adminOf {
+      departments {
+        id
+        name
+      }
+      locations {
+        id
+        name
+      }
+      groups {
+        id
+        name
+      }
+    }
+    birthday {
+      month
+      day
+    }
+    companyInfo {
+      customEula
+      customAcceptTermsMessage
+      name
+      allowUsersToShareInfo
+      allowProfilePictureUploads
+      kwProfileSettings {
+        canShareBirthday
+        canShareAnniversary
+        canShareBio
+        canChangeAvatar
+        canChangePreferredName
+      }
+      timezone
+    }
+    location
+    department
+    departmentInfo {
+      id
+      name
+    }
+    firstName
+    preferredName
+    id
+    ids {
+      kazooId
+      rrId
+      yeiImpersonatingUserId
+      yeiCompanyId
+    }
+    email
+    image30
+    image70
+    image160
+    avatar
+    initials
+    title
+    jobTitle: title
+    name
+    roles
+    isManager
+    startDate
+    groupFeatureFlags {
+      performance {
+        adhocGoal
+        brandedEmail
+        enablePerform
+        goalLibrary
+        syncUp
+        continuousFeedback
+        feedback360
+        requestFeedback
+        giveFeedback
+        recognition2
+        enableFeedback
+        enableGoals
+        enablePoints
+        enableLanguages
+        enableLocation
+        enablePoints
+        enablePolling
+        feedbackCheckin
+      }
+      recognition
+    }
+    permissionsInGroup {
+      performance
+    }
+    points {
+      givingBalance
+      redemptionBalance
+    }
+    showTermsConditionsModal
+    status
+    following
+    enabledModules
+    source
+    lastName
+    secondaryManagerOf
+  }
+}`;