refactor: Refactored codeblock and fixed storybook

backend/dataline/services/llm_flow/llm_calls/chart_generator.py

@@ -1,7 +1,9 @@
 from enum import Enum
 from typing import Type
 
+from langsmith import wrappers
 from mirascope import tags
+from mirascope.base import BaseConfig
 from mirascope.openai import OpenAICallParams, OpenAIExtractor
 from pydantic import BaseModel, Field
 
@@ -179,6 +181,7 @@ class ShouldGenerateChart(BaseModel):
 
 @tags(["version:0001"])
 class ShouldGenerateChartCall(OpenAIExtractor[ShouldGenerateChart]):
+ configuration = BaseConfig(client_wrappers=[wrappers.wrap_openai])
  extract_schema: Type[ShouldGenerateChart] = ShouldGenerateChart
  call_params = OpenAICallParams(model="gpt-3.5-turbo")
  api_key: str | None
@@ -198,6 +201,7 @@ class GeneratedChart(BaseModel):
 
 @tags(["version:0001"])
 class GenerateChartCall(OpenAIExtractor[GeneratedChart]):
+ configuration = BaseConfig(client_wrappers=[wrappers.wrap_openai])
  extract_schema: Type[GeneratedChart] = GeneratedChart
  call_params = OpenAICallParams(model="gpt-3.5-turbo")
  api_key: str | None

backend/pyproject.toml

@@ -29,7 +29,7 @@ langchain = "^0.1.20"
 langgraph = "^0.0.48"
 langsmith = "^0.1.57"
 langchain-openai = "^0.1.6"
-mirascope = "^0.12.3"
+mirascope = "^0.18.2"
 tenacity = "^8.2.3"
 pandas = "^2.2.2"
 pyreadstat = "^1.2.7"

frontend/.storybook/preview.tsx

@@ -0,0 +1,97 @@
+import type { Preview, StoryContext } from "@storybook/react";
+import "../src/index.css";
+import {
+ RouterProvider,
+ createMemoryHistory,
+ createRootRoute,
+ createRoute,
+ createRouter,
+} from "@tanstack/react-router";
+import { PartialStoryFn } from "storybook/internal/types";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React from "react";
+import { closeSnackbar, SnackbarKey, SnackbarProvider } from "notistack";
+function withRouter(Story: PartialStoryFn, { parameters }: StoryContext) {
+ const {
+ initialEntries = ["/"],
+ initialIndex,
+ routes = ["/"],
+ routeParams = {},
+ } = parameters?.router || {};
+
+ const rootRoute = createRootRoute();
+
+ const children = routes.map((path: string) =>
+ createRoute({
+ path,
+ getParentRoute: () => rootRoute,
+ component: Story,
+ })
+ );
+
+ rootRoute.addChildren(children);
+
+ // Ensure initialEntries are strings
+ const formattedInitialEntries = initialEntries.map((entry: string) => {
+ // If the entry includes parameters, replace them with the provided values
+ return Object.keys(routeParams).reduce((acc, key) => {
+ return acc.replace(`:${key}`, routeParams[key]);
+ }, entry);
+ });
+
+ const router = createRouter({
+ history: createMemoryHistory({
+ initialEntries: formattedInitialEntries,
+ initialIndex,
+ }),
+ routeTree: rootRoute,
+ context: routeParams,
+ });
+
+ return <RouterProvider router={router} />;
+}
+
+function withQueryClient(Story: PartialStoryFn) {
+ const action = (snackbarId: SnackbarKey | undefined) => (
+ <>
+ <button className="px-2" onClick={() => closeSnackbar(snackbarId)}>
+ Dismiss
+ </button>
+ </>
+ );
+ return (
+ <QueryClientProvider client={new QueryClient()}>
+ <SnackbarProvider
+ autoHideDuration={5000}
+ maxSnack={5}
+ anchorOrigin={{ horizontal: "right", vertical: "bottom" }}
+ action={action}
+ >
+ <Story />
+ </SnackbarProvider>
+ </QueryClientProvider>
+ );
+}
+
+const preview: Preview = {
+ globalTypes: {
+ darkMode: {
+ defaultValue: true, // Enable dark mode by default on all stories
+ },
+ },
+ decorators: [withRouter, withQueryClient],
+ parameters: {
+ backgrounds: {
+ default: "dark",
+ },
+ actions: { argTypesRegex: "^on[A-Z].*" },
+ controls: {
+ matchers: {
+ color: /(background|color)$/i,
+ date: /Date$/,
+ },
+ },
+ },
+};
+
+export default preview;

frontend/src/components/Conversation/CodeBlock.stories.tsx

@@ -1,5 +1,6 @@
 import type { Meta, StoryObj } from "@storybook/react";
 import { CodeBlock } from "@components/Conversation/CodeBlock";
+import { Dialect } from "../Library/types";
 
 const meta: Meta<typeof CodeBlock> = {
  component: CodeBlock,
@@ -14,5 +15,18 @@ type Story = StoryObj<typeof CodeBlock>;
  * to learn how to use render functions.
  */
 export const Primary: Story = {
- args: {},
+ args: {
+ code: `SELECT * FROM table JOIN table2 ON table.id = table2.id WHERE table.id = 1;`,
+ language: Dialect.Postgres,
+ },
+ parameters: {
+ // Place in conversation context
+ router: {
+ routes: ["/", "/_app/chat/$conversationId"],
+ initialEntries: ["/_app/chat/dummy-conversation-id"],
+ routeParams: {
+ conversationId: "dummy-conversation-id",
+ },
+ },
+ },
 };