tr]:last:border-b-0', className)}
38 | {...props}
39 | />
40 | );
41 | }
42 |
43 | function TableRow({ className, ...props }: React.ComponentProps<'tr'>) {
44 | return (
45 | |
53 | );
54 | }
55 |
56 | function TableHead({ className, ...props }: React.ComponentProps<'th'>) {
57 | return (
58 | [role=checkbox]]:translate-y-[2px]',
62 | className
63 | )}
64 | {...props}
65 | />
66 | );
67 | }
68 |
69 | function TableCell({ className, ...props }: React.ComponentProps<'td'>) {
70 | return (
71 | | [role=checkbox]]:translate-y-[2px]',
75 | className
76 | )}
77 | {...props}
78 | />
79 | );
80 | }
81 |
82 | function TableCaption({ className, ...props }: React.ComponentProps<'caption'>) {
83 | return (
84 |
89 | );
90 | }
91 |
92 | export { Table, TableHeader, TableBody, TableFooter, TableHead, TableRow, TableCell, TableCaption };
93 |
--------------------------------------------------------------------------------
/app_frontend/src/components/ui/tabs.tsx:
--------------------------------------------------------------------------------
1 | import * as React from 'react';
2 | import * as TabsPrimitive from '@radix-ui/react-tabs';
3 |
4 | import { cn } from '@/lib/utils';
5 |
6 | function Tabs({ className, ...props }: React.ComponentProps) {
7 | return (
8 |
13 | );
14 | }
15 |
16 | function TabsList({ className, ...props }: React.ComponentProps) {
17 | return (
18 |
26 | );
27 | }
28 |
29 | function TabsTrigger({ className, ...props }: React.ComponentProps) {
30 | return (
31 |
39 | );
40 | }
41 |
42 | function TabsContent({ className, ...props }: React.ComponentProps) {
43 | return (
44 |
52 | );
53 | }
54 |
55 | export { Tabs, TabsList, TabsTrigger, TabsContent };
56 |
--------------------------------------------------------------------------------
/app_frontend/src/components/ui/toggle-group.tsx:
--------------------------------------------------------------------------------
1 | import * as React from 'react';
2 | import * as ToggleGroupPrimitive from '@radix-ui/react-toggle-group';
3 | import { type VariantProps } from 'class-variance-authority';
4 |
5 | import { cn } from '@/lib/utils';
6 | import { toggleVariants } from '@/components/ui/toggle';
7 |
8 | const ToggleGroupContext = React.createContext>({
9 | size: 'default',
10 | variant: 'default',
11 | });
12 |
13 | function ToggleGroup({
14 | className,
15 | variant,
16 | size,
17 | children,
18 | ...props
19 | }: React.ComponentProps & VariantProps) {
20 | return (
21 |
31 |
32 | {children}
33 |
34 |
35 | );
36 | }
37 |
38 | function ToggleGroupItem({
39 | className,
40 | children,
41 | variant,
42 | size,
43 | ...props
44 | }: React.ComponentProps & VariantProps) {
45 | const context = React.useContext(ToggleGroupContext);
46 |
47 | return (
48 |
62 | {children}
63 |
64 | );
65 | }
66 |
67 | export { ToggleGroup, ToggleGroupItem };
68 |
--------------------------------------------------------------------------------
/app_frontend/src/components/ui/toggle.tsx:
--------------------------------------------------------------------------------
1 | 'use client';
2 |
3 | import * as React from 'react';
4 | import * as TogglePrimitive from '@radix-ui/react-toggle';
5 | import { cva, type VariantProps } from 'class-variance-authority';
6 |
7 | import { cn } from '@/lib/utils';
8 |
9 | const toggleVariants = cva(
10 | "inline-flex items-center justify-center gap-2 rounded-md text-sm font-medium hover:bg-muted hover:text-muted-foreground disabled:pointer-events-none disabled:opacity-50 data-[state=on]:bg-accent data-[state=on]:text-accent-foreground [&_svg]:pointer-events-none [&_svg:not([class*='size-'])]:size-4 [&_svg]:shrink-0 focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] outline-none transition-[color,box-shadow] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive whitespace-nowrap",
11 | {
12 | variants: {
13 | variant: {
14 | default: 'bg-transparent',
15 | outline:
16 | 'border border-input bg-transparent shadow-xs hover:bg-accent hover:text-accent-foreground',
17 | },
18 | size: {
19 | default: 'h-9 px-2 min-w-9',
20 | sm: 'h-8 px-1.5 min-w-8',
21 | lg: 'h-10 px-2.5 min-w-10',
22 | },
23 | },
24 | defaultVariants: {
25 | variant: 'default',
26 | size: 'default',
27 | },
28 | }
29 | );
30 |
31 | function Toggle({
32 | className,
33 | variant,
34 | size,
35 | ...props
36 | }: React.ComponentProps & VariantProps) {
37 | return (
38 |
43 | );
44 | }
45 |
46 | export { Toggle, toggleVariants }; // eslint-disable-line
47 |
--------------------------------------------------------------------------------
/app_frontend/src/constants/dataSources.ts:
--------------------------------------------------------------------------------
1 | export const DATA_SOURCES = {
2 | DATABASE: 'database',
3 | CATALOG: 'catalog',
4 | FILE: 'file',
5 | };
6 |
--------------------------------------------------------------------------------
/app_frontend/src/constants/dev.ts:
--------------------------------------------------------------------------------
1 | export const VITE_DEFAULT_PORT = 5173;
2 | export const VITE_STATIC_DEFAULT_PORT = 8080;
3 |
--------------------------------------------------------------------------------
/app_frontend/src/global.d.ts:
--------------------------------------------------------------------------------
1 | // global.d.ts
2 |
3 | export {};
4 |
5 | declare global {
6 | interface Window {
7 | ENV: {
8 | APP_BASE_URL?: string;
9 | API_PORT?: string;
10 | DATAROBOT_ENDPOINT?: string;
11 | IS_STATIC_FRONTEND?: boolean;
12 | };
13 | }
14 | }
15 |
--------------------------------------------------------------------------------
/app_frontend/src/jest-dom.d.ts:
--------------------------------------------------------------------------------
1 | // jest-dom.d.ts
2 | import '@testing-library/jest-dom';
3 |
4 | declare global {
5 | namespace jest {
6 | interface Matchers {
7 | toBeInTheDocument(): R;
8 | }
9 | }
10 | }
11 |
12 | export {};
13 |
--------------------------------------------------------------------------------
/app_frontend/src/lib/utils.ts:
--------------------------------------------------------------------------------
1 | import { clsx, type ClassValue } from 'clsx';
2 | import { twMerge } from 'tailwind-merge';
3 |
4 | export function cn(...inputs: ClassValue[]) {
5 | return twMerge(clsx(inputs));
6 | }
7 |
8 | export function isDev() {
9 | return import.meta.env.MODE === 'development';
10 | }
11 |
12 | export function isServedStatic() {
13 | return window.ENV?.IS_STATIC_FRONTEND;
14 | }
15 |
--------------------------------------------------------------------------------
/app_frontend/src/main.tsx:
--------------------------------------------------------------------------------
1 | import { StrictMode } from 'react';
2 | import { createRoot } from 'react-dom/client';
3 | import { BrowserRouter as Router } from 'react-router-dom';
4 | import '@fontsource/dm-sans/300.css';
5 | import '@fontsource/dm-sans/400.css';
6 | import '@fontsource/dm-sans/500.css';
7 | import '@fontsource/dm-sans/600.css';
8 | import '@fontsource/dm-sans/700.css';
9 | import { isServedStatic } from '@/lib/utils.ts';
10 | import './index.css';
11 | import App from './App.tsx';
12 | import { AppStateProvider } from './state';
13 | import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
14 |
15 | let basename = window.ENV?.APP_BASE_URL ?? undefined;
16 | if (
17 | window.ENV?.APP_BASE_URL?.includes('notebook-sessions') &&
18 | window.ENV?.API_PORT &&
19 | isServedStatic()
20 | ) {
21 | basename += `/ports/` + window.ENV.API_PORT;
22 | }
23 |
24 | const queryClient = new QueryClient();
25 |
26 | createRoot(document.getElementById('root')!).render(
27 |
28 |
29 |
30 |
31 |
32 |
33 |
34 |
35 |
36 | );
37 |
--------------------------------------------------------------------------------
/app_frontend/src/pages/Data.tsx:
--------------------------------------------------------------------------------
1 | import React, { useState } from 'react';
2 | import { Separator } from '@radix-ui/react-separator';
3 | import { useGeneratedDictionaries } from '@/api/dictionaries/hooks';
4 |
5 | import {
6 | DatasetCardDescriptionPanel,
7 | DataViewTabs,
8 | SearchControl,
9 | ClearDatasetsButton,
10 | } from '@/components/data';
11 | import { ValueOf } from '@/state/types';
12 | import { DATA_TABS } from '@/state/constants';
13 | import { Loading } from '@/components/ui-custom/loading';
14 |
15 | export const Data: React.FC = () => {
16 | const { data, status } = useGeneratedDictionaries();
17 | const [viewMode, setViewMode] = useState>(DATA_TABS.DESCRIPTION);
18 |
19 | return (
20 |
21 |
22 | Data
23 |
24 |
25 |
26 | View
27 | setViewMode(value as ValueOf)}
30 | />
31 |
32 |
33 |
34 |
35 |
36 |
37 |
38 | {status === 'pending' ? (
39 |
40 |
41 |
42 | ) : (
43 |
44 | {data?.map(dictionary => (
45 |
51 | ))}
52 |
53 | )}
54 | Loading... ;
9 |
10 | const Pages = () => {
11 | return (
12 |
13 | }>
14 |
15 | } />
16 | } />
17 | } />
18 | } />
19 |
20 |
21 |
22 | );
23 | };
24 |
25 | export default Pages;
26 |
--------------------------------------------------------------------------------
/app_frontend/src/pages/routes.ts:
--------------------------------------------------------------------------------
1 | export const ROUTES = {
2 | DATA: '/data',
3 | CHATS: '/chats',
4 | CHAT_WITH_ID: '/chats/:chatId',
5 | };
6 |
7 | export const generateChatRoute = (chatId?: string) => {
8 | if (!chatId) return ROUTES.CHATS;
9 | return `/chats/${chatId}`;
10 | };
11 |
--------------------------------------------------------------------------------
/app_frontend/src/state/AppStateContext.ts:
--------------------------------------------------------------------------------
1 | import { createContext } from 'react';
2 | import { AppState } from './types';
3 |
4 | export const AppStateContext = createContext({} as AppState);
5 |
--------------------------------------------------------------------------------
/app_frontend/src/state/AppStateProvider.tsx:
--------------------------------------------------------------------------------
1 | import React, { useReducer } from 'react';
2 | import { AppState } from './types';
3 | import { reducer, createInitialState, actions } from './reducer';
4 | import { AppStateContext } from './AppStateContext';
5 |
6 | export const AppStateProvider: React.FC<{
7 | children: React.ReactNode;
8 | }> = ({ children }) => {
9 | const [state, dispatch] = useReducer(reducer, createInitialState());
10 |
11 | const hideWelcomeModal = () => {
12 | dispatch(actions.hideWelcomeModal());
13 | };
14 |
15 | const setCollapsiblePanelDefaultOpen = (isOpen: boolean) => {
16 | dispatch(actions.setCollapsiblePanelDefaultOpen(isOpen));
17 | };
18 |
19 | const setEnableChartGeneration = (enabled: boolean) => {
20 | dispatch(actions.setEnableChartGeneration(enabled));
21 | };
22 |
23 | const setEnableBusinessInsights = (enabled: boolean) => {
24 | dispatch(actions.setEnableBusinessInsights(enabled));
25 | };
26 |
27 | const setDataSource = (source: string) => {
28 | dispatch(actions.setDataSource(source));
29 | };
30 |
31 | const contextValue: AppState = {
32 | ...state,
33 | hideWelcomeModal,
34 | setCollapsiblePanelDefaultOpen,
35 | setEnableChartGeneration,
36 | setEnableBusinessInsights,
37 | setDataSource,
38 | };
39 |
40 | return {children};
41 | };
42 |
--------------------------------------------------------------------------------
/app_frontend/src/state/constants.ts:
--------------------------------------------------------------------------------
1 | import { ValueOf } from './types';
2 |
3 | export const ACTION_TYPES = {
4 | HIDE_WELCOME_MODAL: 'HIDE_WELCOME_MODAL',
5 | SET_COLLAPSIBLE_PANEL_DEFAULT_OPEN: 'SET_COLLAPSIBLE_PANEL_DEFAULT_OPEN',
6 | SET_ENABLE_CHART_GENERATION: 'SET_ENABLE_CHART_GENERATION',
7 | SET_ENABLE_BUSINESS_INSIGHTS: 'SET_ENABLE_BUSINESS_INSIGHTS',
8 | SET_DATA_SOURCE: 'SET_DATA_SOURCE',
9 | } as const;
10 |
11 | export type StateActionType = ValueOf;
12 |
13 | export const STORAGE_KEYS = {
14 | HIDE_WELCOME_MODAL: 'HIDE_WELCOME_MODAL',
15 | COLLAPSIBLE_PANEL_DEFAULT_OPEN: 'COLLAPSIBLE_PANEL_DEFAULT_OPEN',
16 | ENABLE_CHART_GENERATION: 'ENABLE_CHART_GENERATION',
17 | ENABLE_BUSINESS_INSIGHTS: 'ENABLE_BUSINESS_INSIGHTS',
18 | DATA_SOURCE: 'DATA_SOURCE',
19 | } as const;
20 |
21 | export const NEW_CHAT_ID = 'new';
22 |
23 | export const DATA_TABS = {
24 | DESCRIPTION: 'description',
25 | RAW: 'raw',
26 | };
27 |
--------------------------------------------------------------------------------
/app_frontend/src/state/hooks.ts:
--------------------------------------------------------------------------------
1 | import { useContext } from 'react';
2 | import { AppStateContext } from './AppStateContext';
3 | import { AppState } from './types';
4 |
5 | export const useAppState = (): AppState => {
6 | const context = useContext(AppStateContext);
7 |
8 | if (context === undefined) {
9 | throw new Error('useAppState must be used within an AppStateProvider');
10 | }
11 |
12 | return context;
13 | };
14 |
--------------------------------------------------------------------------------
/app_frontend/src/state/index.ts:
--------------------------------------------------------------------------------
1 | export { AppStateContext } from './AppStateContext';
2 | export { AppStateProvider } from './AppStateProvider';
3 | export { useAppState } from './hooks';
4 | export type { AppState } from './types';
5 | export { NEW_CHAT_ID } from './constants';
6 |
--------------------------------------------------------------------------------
/app_frontend/src/state/reducer.ts:
--------------------------------------------------------------------------------
1 | import { AppStateData, Action } from './types';
2 | import { ACTION_TYPES, STORAGE_KEYS } from './constants';
3 | import { getStorageItem, setStorageItem } from './storage';
4 |
5 | import { DATA_SOURCES } from '@/constants/dataSources';
6 |
7 | export const createInitialState = (): AppStateData => {
8 | return {
9 | showWelcome: getStorageItem(STORAGE_KEYS.HIDE_WELCOME_MODAL) !== 'true',
10 | collapsiblePanelDefaultOpen:
11 | getStorageItem(STORAGE_KEYS.COLLAPSIBLE_PANEL_DEFAULT_OPEN) === 'true',
12 | enableChartGeneration: getStorageItem(STORAGE_KEYS.ENABLE_CHART_GENERATION) !== 'false', // Enable by default
13 | enableBusinessInsights: getStorageItem(STORAGE_KEYS.ENABLE_BUSINESS_INSIGHTS) !== 'false', // Enable by default
14 | dataSource: getStorageItem(STORAGE_KEYS.DATA_SOURCE) || DATA_SOURCES.FILE, // Default to FILE
15 | };
16 | };
17 |
18 | export const reducer = (state: AppStateData, action: Action): AppStateData => {
19 | switch (action.type) {
20 | case ACTION_TYPES.HIDE_WELCOME_MODAL:
21 | setStorageItem(STORAGE_KEYS.HIDE_WELCOME_MODAL, 'true');
22 | return {
23 | ...state,
24 | showWelcome: false,
25 | };
26 | case ACTION_TYPES.SET_COLLAPSIBLE_PANEL_DEFAULT_OPEN:
27 | setStorageItem(
28 | STORAGE_KEYS.COLLAPSIBLE_PANEL_DEFAULT_OPEN,
29 | action.payload ? 'true' : 'false'
30 | );
31 | return {
32 | ...state,
33 | collapsiblePanelDefaultOpen: action.payload,
34 | };
35 | case ACTION_TYPES.SET_ENABLE_CHART_GENERATION:
36 | setStorageItem(STORAGE_KEYS.ENABLE_CHART_GENERATION, action.payload ? 'true' : 'false');
37 | return {
38 | ...state,
39 | enableChartGeneration: action.payload,
40 | };
41 | case ACTION_TYPES.SET_ENABLE_BUSINESS_INSIGHTS:
42 | setStorageItem(
43 | STORAGE_KEYS.ENABLE_BUSINESS_INSIGHTS,
44 | action.payload ? 'true' : 'false'
45 | );
46 | return {
47 | ...state,
48 | enableBusinessInsights: action.payload,
49 | };
50 | case ACTION_TYPES.SET_DATA_SOURCE:
51 | setStorageItem(STORAGE_KEYS.DATA_SOURCE, action.payload);
52 | return {
53 | ...state,
54 | dataSource: action.payload,
55 | };
56 | default:
57 | return state;
58 | }
59 | };
60 |
61 | export const actions = {
62 | hideWelcomeModal: (): Action => ({
63 | type: ACTION_TYPES.HIDE_WELCOME_MODAL,
64 | }),
65 | setCollapsiblePanelDefaultOpen: (isOpen: boolean): Action => ({
66 | type: ACTION_TYPES.SET_COLLAPSIBLE_PANEL_DEFAULT_OPEN,
67 | payload: isOpen,
68 | }),
69 | setEnableChartGeneration: (enabled: boolean): Action => ({
70 | type: ACTION_TYPES.SET_ENABLE_CHART_GENERATION,
71 | payload: enabled,
72 | }),
73 | setEnableBusinessInsights: (enabled: boolean): Action => ({
74 | type: ACTION_TYPES.SET_ENABLE_BUSINESS_INSIGHTS,
75 | payload: enabled,
76 | }),
77 | setDataSource: (source: string): Action => ({
78 | type: ACTION_TYPES.SET_DATA_SOURCE,
79 | payload: source,
80 | }),
81 | };
82 |
--------------------------------------------------------------------------------
/app_frontend/src/state/storage.ts:
--------------------------------------------------------------------------------
1 | import { STORAGE_KEYS } from './constants';
2 |
3 | const getAppIdFromUrl = (): string => {
4 | const url = window.location.href;
5 | const match = url.match(/\/custom_applications\/([^/]+)/);
6 | return match ? match[1] : '';
7 | };
8 |
9 | /**
10 | * Creates a prefixed key in the format /custom_applications/{id}/{key}
11 | */
12 | const getPrefixedKey = (key: string): string => {
13 | const appId = getAppIdFromUrl();
14 | return appId ? `/custom_applications/${appId}/${key}` : key;
15 | };
16 |
17 | export const getStorageItem = (key: string): string | null => {
18 | return localStorage.getItem(getPrefixedKey(key));
19 | };
20 |
21 | export const setStorageItem = (key: string, value: string): void => {
22 | localStorage.setItem(getPrefixedKey(key), value);
23 | };
24 |
25 | export const isWelcomeModalHidden = (): boolean => {
26 | return getStorageItem(STORAGE_KEYS.HIDE_WELCOME_MODAL) === 'true';
27 | };
28 |
29 | export const isCollapsiblePanelDefaultOpen = (): boolean => {
30 | return getStorageItem(STORAGE_KEYS.COLLAPSIBLE_PANEL_DEFAULT_OPEN) === 'true';
31 | };
32 |
33 | export const setCollapsiblePanelDefaultOpen = (isOpen: boolean): void => {
34 | setStorageItem(STORAGE_KEYS.COLLAPSIBLE_PANEL_DEFAULT_OPEN, isOpen ? 'true' : 'false');
35 | };
36 |
--------------------------------------------------------------------------------
/app_frontend/src/state/types.ts:
--------------------------------------------------------------------------------
1 | export type ValueOf = T[keyof T];
2 |
3 | export interface AppStateData {
4 | showWelcome: boolean;
5 | collapsiblePanelDefaultOpen: boolean;
6 | enableChartGeneration: boolean;
7 | enableBusinessInsights: boolean;
8 | dataSource: string;
9 | }
10 |
11 | export interface AppStateActions {
12 | hideWelcomeModal: () => void;
13 | setCollapsiblePanelDefaultOpen: (isOpen: boolean) => void;
14 | setEnableChartGeneration: (enabled: boolean) => void;
15 | setEnableBusinessInsights: (enabled: boolean) => void;
16 | setDataSource: (source: string) => void;
17 | }
18 |
19 | export type AppState = AppStateData & AppStateActions;
20 |
21 | export type Action =
22 | | { type: 'HIDE_WELCOME_MODAL' }
23 | | { type: 'SET_COLLAPSIBLE_PANEL_DEFAULT_OPEN'; payload: boolean }
24 | | { type: 'SET_ENABLE_CHART_GENERATION'; payload: boolean }
25 | | { type: 'SET_ENABLE_BUSINESS_INSIGHTS'; payload: boolean }
26 | | { type: 'SET_DATA_SOURCE'; payload: string };
27 |
--------------------------------------------------------------------------------
/app_frontend/src/vite-env.d.ts:
--------------------------------------------------------------------------------
1 | ///
2 |
--------------------------------------------------------------------------------
/app_frontend/tailwind.config.js:
--------------------------------------------------------------------------------
1 | /** @type {import('tailwindcss').Config} */
2 | module.exports = {
3 | darkMode: ["class"],
4 | content: ["app/**/*.{ts,tsx}", "components/**/*.{ts,tsx}"],
5 | theme: {
6 | extend: {
7 | colors: {
8 | border: "hsl(var(--border))",
9 | input: "hsl(var(--input))",
10 | ring: "hsl(var(--ring))",
11 | background: "hsl(var(--background))",
12 | foreground: "hsl(var(--foreground))",
13 | primary: {
14 | DEFAULT: "hsl(var(--primary))",
15 | foreground: "hsl(var(--primary-foreground))",
16 | },
17 | secondary: {
18 | DEFAULT: "hsl(var(--secondary))",
19 | foreground: "hsl(var(--secondary-foreground))",
20 | },
21 | destructive: {
22 | DEFAULT: "hsl(var(--destructive))",
23 | foreground: "hsl(var(--destructive-foreground))",
24 | },
25 | muted: {
26 | DEFAULT: "hsl(var(--muted))",
27 | foreground: "hsl(var(--muted-foreground))",
28 | },
29 | accent: {
30 | DEFAULT: "hsl(var(--accent))",
31 | foreground: "hsl(var(--accent-foreground))",
32 | },
33 | popover: {
34 | DEFAULT: "hsl(var(--popover))",
35 | foreground: "hsl(var(--popover-foreground))",
36 | },
37 | card: {
38 | DEFAULT: "hsl(var(--card))",
39 | foreground: "hsl(var(--card-foreground))",
40 | },
41 | },
42 | borderRadius: {
43 | lg: `var(--radius)`,
44 | md: `calc(var(--radius) - 2px)`,
45 | sm: "calc(var(--radius) - 4px)",
46 | },
47 | },
48 | },
49 | plugins: [require("tailwindcss-animate")],
50 | }
51 |
--------------------------------------------------------------------------------
/app_frontend/tests/__mocks__/handlers.ts:
--------------------------------------------------------------------------------
1 | import { appHandlers } from "./handlers/app";
2 |
3 | export const handlers = [
4 | ...appHandlers,
5 | ]
6 |
--------------------------------------------------------------------------------
/app_frontend/tests/__mocks__/handlers/app.ts:
--------------------------------------------------------------------------------
1 | import { http, HttpResponse } from 'msw'
2 |
3 | export const appHandlers = [
4 | http.get('api/v1/welcome', () => {
5 | return HttpResponse.json({
6 | message: 'Welcome Engineer!'
7 | })
8 | }),
9 | ]
10 |
--------------------------------------------------------------------------------
/app_frontend/tests/__mocks__/node.ts:
--------------------------------------------------------------------------------
1 | import { setupServer } from 'msw/node'
2 | import { handlers } from './handlers'
3 |
4 | export const server = setupServer(...handlers)
5 |
--------------------------------------------------------------------------------
/app_frontend/tests/components/LoadingIndicator.test.tsx:
--------------------------------------------------------------------------------
1 | import { render, screen } from '@testing-library/react';
2 | import { test, describe, expect, vi } from 'vitest';
3 | import { LoadingIndicator } from '@/components/chat/LoadingIndicator';
4 |
5 | // Mock the FontAwesomeIcon component
6 | vi.mock('@fortawesome/react-fontawesome', () => ({
7 | FontAwesomeIcon: () =>
8 | }));
9 |
10 | // Mock the loader SVG import
11 | vi.mock('@/assets/loader.svg', () => ({
12 | default: 'mock-loader-path',
13 | }));
14 |
15 | describe('LoadingIndicator Component', () => {
16 | test('renders loader when isLoading is true', () => {
17 | render();
18 |
19 | const loader = screen.getByRole('img');
20 | expect(loader).toBeInTheDocument();
21 | expect(loader).toHaveAttribute('src', 'mock-loader-path');
22 | expect(loader).toHaveAttribute('alt', 'processing');
23 | expect(loader).toHaveClass('animate-spin');
24 |
25 | // Check that the icon is not rendered
26 | expect(screen.queryByTestId('mock-icon')).not.toBeInTheDocument();
27 | });
28 |
29 | test('renders check icon when isLoading is false', () => {
30 | render();
31 |
32 | // Check that the loader is not rendered
33 | expect(screen.queryByRole('img')).not.toBeInTheDocument();
34 |
35 | // Check that the icon is rendered
36 | const icon = screen.getByTestId('mock-icon');
37 | expect(icon).toBeInTheDocument();
38 | });
39 | });
40 |
--------------------------------------------------------------------------------
/app_frontend/tests/components/MessageContainer.test.tsx:
--------------------------------------------------------------------------------
1 | import { render, screen } from '@testing-library/react';
2 | import { test, describe, expect, vi } from 'vitest';
3 | import { MessageContainer } from '@/components/chat/MessageContainer';
4 |
5 | describe('MessageContainer Component', () => {
6 | // Setup and teardown for all tests
7 | let scrollIntoViewMock;
8 | let originalScrollIntoView: typeof HTMLElement.prototype.scrollIntoView;
9 |
10 | beforeEach(() => {
11 | // Store original method before mocking
12 | originalScrollIntoView = HTMLElement.prototype.scrollIntoView;
13 | // Create a mock for scrollIntoView
14 | scrollIntoViewMock = vi.fn();
15 |
16 | // Mock scrollIntoView on the HTMLElement prototype
17 | Object.defineProperty(HTMLElement.prototype, 'scrollIntoView', {
18 | configurable: true,
19 | value: scrollIntoViewMock,
20 | writable: true
21 | });
22 | });
23 |
24 | afterEach(() => {
25 | // Restore the original method after each test
26 | Object.defineProperty(HTMLElement.prototype, 'scrollIntoView', {
27 | configurable: true,
28 | value: originalScrollIntoView,
29 | writable: true
30 | });
31 | });
32 |
33 | test('renders children correctly', () => {
34 | render(
35 |
36 | Test Child Content
37 |
38 | );
39 |
40 | const childElement = screen.getByTestId('test-child');
41 | expect(childElement).toBeInTheDocument();
42 | expect(childElement).toHaveTextContent('Test Child Content');
43 | });
44 |
45 | test('applies correct styling', () => {
46 | render(
47 |
48 | Content
49 |
50 | );
51 |
52 | const container = screen.getByText('Content').parentElement;
53 | expect(container).toHaveClass('p-3');
54 | expect(container).toHaveClass('bg-card');
55 | expect(container).toHaveClass('rounded');
56 | expect(container).toHaveClass('flex-col');
57 | });
58 |
59 | test('calls scrollIntoView when mounted', () => {
60 | render(
61 |
62 | Scroll Test
63 |
64 | );
65 |
66 | // Verify scrollIntoView was called with false
67 | expect(scrollIntoViewMock).toHaveBeenCalledWith(false);
68 | });
69 | });
70 |
--------------------------------------------------------------------------------
/app_frontend/tests/components/NewChatModal.test.tsx:
--------------------------------------------------------------------------------
1 | import { render, screen, fireEvent } from '@testing-library/react';
2 | import { test, describe, expect, vi, type MockInstance, type Mock } from 'vitest';
3 | import { NewChatModal } from '@/components/NewChatModal';
4 | import { useCreateChat } from '@/api/chat-messages/hooks';
5 | import { useNavigate } from 'react-router-dom';
6 | import { generateChatRoute } from '@/pages/routes';
7 |
8 | // Mock the dependencies
9 | vi.mock('@/api/chat-messages/hooks', () => ({
10 | useCreateChat: vi.fn(),
11 | }));
12 |
13 | vi.mock('react-router-dom', () => ({
14 | useNavigate: vi.fn(),
15 | }));
16 |
17 | vi.mock('@/pages/routes', () => ({
18 | generateChatRoute: vi.fn(),
19 | }));
20 |
21 | // Mock FontAwesomeIcon
22 | vi.mock('@fortawesome/react-fontawesome', () => ({
23 | FontAwesomeIcon: () =>
24 | }));
25 |
26 | // Mock lucide-react
27 | vi.mock('lucide-react', () => ({
28 | XIcon: () =>
29 | }));
30 |
31 | describe('NewChatModal Component', () => {
32 | let mockCreateChat: MockInstance;
33 | let mockNavigate: MockInstance;
34 |
35 | // Setup default mocks before each test
36 | beforeEach(() => {
37 | // Mock the createChat mutation
38 | mockCreateChat = vi.fn();
39 | (useCreateChat as Mock).mockReturnValue({
40 | mutate: mockCreateChat,
41 | isPending: false,
42 | });
43 |
44 | // Mock the navigate function
45 | mockNavigate = vi.fn();
46 | (useNavigate as Mock).mockReturnValue(mockNavigate);
47 |
48 | // Mock the generateChatRoute function
49 | (generateChatRoute as Mock).mockImplementation((id) => `/chat/${id}`);
50 | });
51 |
52 | test('renders the trigger button correctly', () => {
53 | render();
54 | const newChatButton = screen.getByRole('button', { name: /new chat/i });
55 | expect(newChatButton).toBeInTheDocument();
56 | expect(screen.getByTestId('mock-icon')).toBeInTheDocument();
57 | });
58 |
59 | test('shows loading state when creating a chat', () => {
60 | // Set isPending to true for loading state
61 | (useCreateChat as Mock).mockReturnValue({
62 | mutate: mockCreateChat,
63 | isPending: true,
64 | });
65 |
66 | render();
67 |
68 | // Find and click the New Chat button to open the dialog
69 | const triggerButton = screen.getByRole('button', { name: /new chat/i });
70 | fireEvent.click(triggerButton);
71 |
72 | // Now the dialog content should have a "Creating..." button
73 | const createButton = screen.getByRole('button', { name: /creating/i });
74 | expect(createButton).toBeInTheDocument();
75 | expect(createButton).toBeDisabled();
76 | });
77 | });
78 |
--------------------------------------------------------------------------------
/app_frontend/tests/components/input.test.tsx:
--------------------------------------------------------------------------------
1 | import { render, screen } from '@testing-library/react';
2 | import { test, describe, expect, vi } from 'vitest';
3 | import userEvent from '@testing-library/user-event';
4 | import { Input } from '@/components/ui/input';
5 |
6 | describe('Input Component', () => {
7 | test('renders input with default props', () => {
8 | render();
9 | const input = screen.getByRole('textbox');
10 |
11 | expect(input).toBeInTheDocument();
12 | expect(input).toHaveClass('h-9');
13 | });
14 |
15 | test('sets the correct input type', () => {
16 | const { rerender } = render();
17 | let input = screen.getByRole('textbox');
18 | expect(input).toHaveAttribute('type', 'text');
19 |
20 | rerender();
21 | input = screen.getByRole('textbox');
22 | expect(input).toHaveAttribute('type', 'email');
23 |
24 | rerender();
25 | expect(screen.queryByRole('textbox')).not.toBeInTheDocument();
26 | const passwordInput = screen.getByDisplayValue('');
27 | expect(passwordInput).toHaveAttribute('type', 'password');
28 | });
29 |
30 | test('applies custom className', () => {
31 | render();
32 | const input = screen.getByRole('textbox');
33 |
34 | expect(input).toHaveClass('test-class');
35 | });
36 |
37 | test('accepts and displays value correctly', async () => {
38 | const user = userEvent.setup();
39 | render();
40 | const input = screen.getByRole('textbox');
41 |
42 | expect(input).toHaveValue('Default text');
43 |
44 | await user.clear(input);
45 | await user.type(input, 'New text');
46 |
47 | expect(input).toHaveValue('New text');
48 | });
49 |
50 | test('handles disabled state correctly', async () => {
51 | const user = userEvent.setup();
52 | render();
53 | const input = screen.getByRole('textbox');
54 |
55 | expect(input).toBeDisabled();
56 |
57 | await user.type(input, 'Test text');
58 | expect(input).not.toHaveValue('Test text');
59 | });
60 |
61 | test('passes additional props to the input element', async () => {
62 | const handleChange = vi.fn();
63 | const user = userEvent.setup();
64 | render();
65 | const input = screen.getByPlaceholderText('Enter text here');
66 |
67 | await user.type(input, 'a');
68 | expect(handleChange).toHaveBeenCalledTimes(1);
69 | });
70 |
71 | test('handles aria-invalid attribute', () => {
72 | render();
73 | const input = screen.getByRole('textbox');
74 |
75 | expect(input).toHaveAttribute('aria-invalid', 'true');
76 | expect(input).toHaveClass('aria-invalid:border-destructive/60');
77 | });
78 | });
79 |
--------------------------------------------------------------------------------
/app_frontend/tests/setupTests.ts:
--------------------------------------------------------------------------------
1 | import '@testing-library/jest-dom'
2 | import { server } from './__mocks__/node.js'
3 |
4 | beforeAll(() => {
5 | server.listen()
6 | })
7 |
8 | afterEach(() => {
9 | server.resetHandlers()
10 | })
11 |
12 | afterAll(() => {
13 | server.close()
14 | })
15 |
--------------------------------------------------------------------------------
/app_frontend/tests/test-utils.tsx:
--------------------------------------------------------------------------------
1 | import { ReactNode } from 'react';
2 | import { render } from '@testing-library/react';
3 | import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
4 | import { MemoryRouter } from 'react-router-dom';
5 | import { AppStateProvider } from '@/state';
6 |
7 | const createTestQueryClient = () =>
8 | new QueryClient({
9 | defaultOptions: {
10 | queries: {
11 | retry: false,
12 | },
13 | },
14 | });
15 |
16 | export function renderWithProviders(children: ReactNode) {
17 | const queryClient = createTestQueryClient();
18 |
19 | return render(
20 |
21 |
22 | {children}
23 |
24 |
25 | );
26 | }
27 |
--------------------------------------------------------------------------------
/app_frontend/tsconfig.app.json:
--------------------------------------------------------------------------------
1 | {
2 | "compilerOptions": {
3 | "tsBuildInfoFile": "./node_modules/.tmp/tsconfig.app.tsbuildinfo",
4 | "target": "ES2020",
5 | "useDefineForClassFields": true,
6 | "lib": ["ES2020", "DOM", "DOM.Iterable"],
7 | "module": "ESNext",
8 | "skipLibCheck": true,
9 |
10 | /* Bundler mode */
11 | "moduleResolution": "bundler",
12 | "allowImportingTsExtensions": true,
13 | "isolatedModules": true,
14 | "moduleDetection": "force",
15 | "noEmit": true,
16 | "jsx": "react-jsx",
17 |
18 | /* Linting */
19 | "strict": true,
20 | "noUnusedLocals": true,
21 | "noUnusedParameters": true,
22 | "noFallthroughCasesInSwitch": true,
23 | "noUncheckedSideEffectImports": true,
24 | "baseUrl": ".",
25 | "paths": {
26 | "@/*": [
27 | "./src/*"
28 | ],
29 | "~/*": [
30 | "./src/*"
31 | ]
32 | }
33 | },
34 | "include": ["src"]
35 | }
36 |
--------------------------------------------------------------------------------
/app_frontend/tsconfig.json:
--------------------------------------------------------------------------------
1 | {
2 | "files": [],
3 | "references": [
4 | { "path": "./tsconfig.app.json" },
5 | { "path": "./tsconfig.node.json" },
6 | ],
7 | "compilerOptions": {
8 | "baseUrl": ".",
9 | "paths": {
10 | "@/*": ["./src/*"],
11 | "~/*": ["./src/*"]
12 | }
13 | }
14 | }
15 |
--------------------------------------------------------------------------------
/app_frontend/tsconfig.node.json:
--------------------------------------------------------------------------------
1 | {
2 | "compilerOptions": {
3 | "tsBuildInfoFile": "./node_modules/.tmp/tsconfig.node.tsbuildinfo",
4 | "target": "ES2022",
5 | "lib": ["ES2023"],
6 | "module": "ESNext",
7 | "skipLibCheck": true,
8 |
9 | /* Bundler mode */
10 | "moduleResolution": "bundler",
11 | "allowImportingTsExtensions": true,
12 | "isolatedModules": true,
13 | "moduleDetection": "force",
14 | "noEmit": true,
15 |
16 | /* Linting */
17 | "strict": true,
18 | "noUnusedLocals": true,
19 | "noUnusedParameters": true,
20 | "noFallthroughCasesInSwitch": true,
21 | "noUncheckedSideEffectImports": true
22 | },
23 | "include": ["vite.config.ts"]
24 | }
25 |
--------------------------------------------------------------------------------
/app_frontend/tsconfig.test.json:
--------------------------------------------------------------------------------
1 | {
2 | "extends": ["./tsconfig.app.json", "./tsconfig.node.json"],
3 | "compilerOptions": {
4 | "lib": ["esnext", "dom"],
5 | "types": ["vitest/globals", "vitest", "node", "@testing-library/jest-dom"],
6 | "module": "ESNext",
7 | "target": "ESNext",
8 | "isolatedModules": true,
9 | "noEmit": true,
10 | "baseUrl": ".",
11 | "paths": {
12 | "@/*": ["./src/*"],
13 | "~/*": ["./src/*"]
14 | }
15 | },
16 | "include": ["setupTests.ts", "src", "tests", "**/*.test.ts", "**/*.test.tsx"]
17 | }
18 |
--------------------------------------------------------------------------------
/app_frontend/vite.config.ts:
--------------------------------------------------------------------------------
1 | import { defineConfig } from 'vitest/config'
2 | import react from '@vitejs/plugin-react'
3 | import tailwindcss from "@tailwindcss/vite";
4 | import { nodePolyfills } from "vite-plugin-node-polyfills";
5 | import path from "path";
6 | import { VITE_DEFAULT_PORT, VITE_STATIC_DEFAULT_PORT } from "./src/constants/dev";
7 |
8 | let base: string = '';
9 | // 1. if NOTEBOOK_ID is set, use /notebook-sessions/${NOTEBOOK_ID}/ports/5173/ for dev server
10 | // 2. if NOTEBOOK_ID and NODE_ENV === 'development' are set, use /notebook-sessions/${NOTEBOOK_ID}/ports/8080/ for static codespace server
11 | if (process.env.NOTEBOOK_ID && process.env.NODE_ENV === 'development') {
12 | const notebookId = process.env.NOTEBOOK_ID;
13 | const defaultPort = process.env.STATIC_CODESPACE ? VITE_STATIC_DEFAULT_PORT : VITE_DEFAULT_PORT;
14 | base = `/notebook-sessions/${notebookId}/ports/${defaultPort}/`;
15 | }
16 | const proxyBase: string = base === '' ? '/' : base;
17 |
18 |
19 | // https://vite.dev/config/
20 | export default defineConfig({
21 | plugins: [
22 | react(),
23 | tailwindcss(),
24 | nodePolyfills({
25 | exclude: [],
26 | // for plotly.js
27 | protocolImports: true,
28 | }),
29 | {
30 | name: 'strip-base',
31 | apply: 'serve',
32 | configureServer({ middlewares }) {
33 | middlewares.use((req, _res, next) => {
34 | if (base !== '' && !req.url?.startsWith(base)) {
35 | req.url = base.slice(0, -1) + req.url;
36 | }
37 | next();
38 | });
39 | },
40 | },
41 | ],
42 | resolve: {
43 | alias: {
44 | "@": path.resolve(__dirname, "./src"),
45 | "~": path.resolve(__dirname, "./src"),
46 | },
47 | },
48 | base: base,
49 | build: {
50 | outDir: '../app_backend/static/',
51 | emptyOutDir: true,
52 | rollupOptions: {
53 | external: ['_dr_env.js'],
54 | },
55 | },
56 | server: {
57 | host: true,
58 | allowedHosts: ["localhost", "127.0.0.1", ".datarobot.com"],
59 | proxy: {
60 | [`${proxyBase}api/`]: {
61 | target: 'http://localhost:8080',
62 | changeOrigin: true,
63 | rewrite: (path) => path.replace(new RegExp(`^${proxyBase}`), ''),
64 | },
65 | [`${proxyBase}_dr_env.js`]: {
66 | target: 'http://localhost:8080',
67 | changeOrigin: true,
68 | rewrite: (path) => path.replace(new RegExp(`^${proxyBase}`), ''),
69 | },
70 | },
71 | },
72 | test: {
73 | globals: true,
74 | environment: 'jsdom',
75 | setupFiles: './tests/setupTests.ts',
76 | typecheck: {
77 | tsconfig: './tsconfig.test.json',
78 | }
79 | },
80 | })
81 |
--------------------------------------------------------------------------------
/frontend/.streamlit/config.toml:
--------------------------------------------------------------------------------
1 | [server]
2 | maxUploadSize = 200
3 | maxMessageSize = 200
4 | [theme]
5 | base = "light"
6 | primaryColor = "#909BF5"
7 | textColor = "#0B0B0B"
8 | font = "sans-serif"
9 |
--------------------------------------------------------------------------------
/frontend/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/datarobot-community/talk-to-my-data-agent/17301f5fad9192aeefa13315f179328b29135334/frontend/__init__.py
--------------------------------------------------------------------------------
/frontend/app.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | import asyncio
16 | from typing import List
17 |
18 | import streamlit as st
19 | from app_settings import PAGE_ICON, apply_custom_css
20 | from datarobot_connect import DataRobotTokenManager
21 | from streamlit.navigation.page import StreamlitPage
22 |
23 | pages: List[StreamlitPage] = [
24 | st.Page("01_connect_and_explore.py", title="Connect & Explore"),
25 | st.Page("02_chat_with_data.py", title="AI Data Analyst"),
26 | ]
27 |
28 | st.set_page_config(
29 | page_icon=PAGE_ICON,
30 | layout="wide",
31 | initial_sidebar_state="expanded",
32 | )
33 | apply_custom_css()
34 |
35 | if "datarobot_connect" not in st.session_state:
36 | datarobot_connect = DataRobotTokenManager()
37 | st.session_state.datarobot_connect = datarobot_connect
38 |
39 |
40 | asyncio.run(
41 | st.session_state.datarobot_connect.display_info(
42 | st.sidebar.container(key="user_info")
43 | )
44 | )
45 |
46 | pg = st.navigation(pages)
47 | pg.run()
48 |
--------------------------------------------------------------------------------
/frontend/app_settings.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 |
16 | import sys
17 |
18 | import streamlit as st
19 | from streamlit_theme import st_theme
20 |
21 | sys.path.append("..")
22 | from utils.schema import AppInfra
23 |
24 | PAGE_ICON = "./datarobot_favicon.png"
25 |
26 |
27 | def display_page_logo() -> None:
28 | theme = st_theme()
29 | # logo placeholder used for initial load
30 | logo = ''
31 | if theme:
32 | if theme.get("base") == "light":
33 | logo = "./DataRobot_black.svg"
34 | else:
35 | logo = "./DataRobot_white.svg"
36 | with st.container(key="datarobot-logo"):
37 | st.image(logo, width=200)
38 |
39 |
40 | def get_database_logo(app_infra: AppInfra) -> None:
41 | if app_infra.database == "snowflake":
42 | st.image("./Snowflake.svg", width=100)
43 | elif app_infra.database == "bigquery":
44 | st.image("./Google_Cloud.svg", width=100)
45 | elif app_infra.database == "sap":
46 | st.image("./sap.svg", width=100)
47 | return None
48 |
49 |
50 | def get_database_loader_message(app_infra: AppInfra) -> str:
51 | if app_infra.database == "snowflake":
52 | return "Load Datasets from Snowflake"
53 | elif app_infra.database == "bigquery":
54 | return "Load Datasets from BigQuery"
55 | elif app_infra.database == "sap":
56 | return "Load Datasets from SAP"
57 | return "No database available"
58 |
59 |
60 | def apply_custom_css() -> None:
61 | with open("./style.css") as f:
62 | css = f.read()
63 | st.markdown(f"", unsafe_allow_html=True)
64 |
--------------------------------------------------------------------------------
/frontend/bot.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/datarobot-community/talk-to-my-data-agent/17301f5fad9192aeefa13315f179328b29135334/frontend/bot.jpg
--------------------------------------------------------------------------------
/frontend/datarobot_favicon.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/datarobot-community/talk-to-my-data-agent/17301f5fad9192aeefa13315f179328b29135334/frontend/datarobot_favicon.png
--------------------------------------------------------------------------------
/frontend/helpers.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | import json
16 | import logging
17 | import sys
18 | import traceback
19 | import uuid
20 | from datetime import datetime
21 | from pathlib import Path
22 | from typing import (
23 | Any,
24 | )
25 |
26 | import streamlit as st
27 |
28 | sys.path.append("..")
29 | from utils.analyst_db import AnalystDB, DataSourceType
30 |
31 | logger = logging.getLogger("DataAnalyst")
32 |
33 |
34 | # Add enhanced error logging function
35 | def log_error_details(error: BaseException, context: dict[str, Any]) -> None:
36 | """Log detailed error information with context
37 |
38 | Args:
39 | error: The exception that occurred
40 | context: Dictionary containing error context
41 | """
42 | error_details = {
43 | "timestamp": datetime.now().isoformat(),
44 | "error_type": type(error).__name__,
45 | "error_message": str(error),
46 | "stack_trace": traceback.format_exc(),
47 | **context,
48 | }
49 |
50 | logger.error(
51 | f"\nERROR DETAILS\n=============\n{json.dumps(error_details, indent=2, default=str)}"
52 | )
53 |
54 |
55 | empty_session_state = {
56 | "initialized": True,
57 | "datasets_names": [],
58 | "cleansed_data_names": [],
59 | "selected_registry_datasets": [],
60 | "data_source": DataSourceType.FILE,
61 | "file_uploader_key": 0,
62 | "processed_file_ids": [],
63 | "chat_messages": [],
64 | "chat_input_key": 0,
65 | "debug_mode": True,
66 | }
67 |
68 |
69 | def state_empty() -> None:
70 | for key, value in empty_session_state.items():
71 | st.session_state[key] = value
72 | logger.info("Session state has been reset to its initial empty state.")
73 |
74 |
75 | def generate_user_id() -> str | None:
76 | email_header = st.context.headers.get("x-user-email")
77 | if email_header:
78 | new_user_id = str(uuid.uuid5(uuid.NAMESPACE_OID, email_header))[:36]
79 | return new_user_id
80 | else:
81 | return None
82 |
83 |
84 | async def state_init() -> None:
85 | if "initialized" not in st.session_state:
86 | state_empty()
87 | user_id = None
88 | if "datarobot_uid" not in st.session_state:
89 | user_id = generate_user_id()
90 | else:
91 | user_id = st.session_state.datarobot_uid
92 | if user_id:
93 | analyst_db = await AnalystDB.create(
94 | user_id=user_id,
95 | db_path=Path("/tmp"),
96 | dataset_db_name="datasets.db",
97 | chat_db_name="chat.db",
98 | )
99 |
100 | st.session_state.analyst_db = analyst_db
101 | else:
102 | logger.warning("datarobot-connect not initialised")
103 | pass
104 |
--------------------------------------------------------------------------------
/frontend/metadata.yaml.jinja:
--------------------------------------------------------------------------------
1 | name: runtime-params
2 |
3 | runtimeParameterDefinitions:
4 | {{ additional_params }}
5 |
--------------------------------------------------------------------------------
/frontend/requirements.txt:
--------------------------------------------------------------------------------
1 | # python sandbox libraries
2 | pandas>=2.2.3,<3.0
3 | numpy>=2.1.3,<3.0
4 | scipy>=1.15.1,<2.0
5 | statsmodels>=0.14.4,<1.0
6 | scikit-learn>=1.6.1,<2.0
7 | lightgbm>=4.5.0,<5.0
8 | tslearn>=0.6.3,<1.0
9 | spacy>=3.8.3,<4.0
10 | pyarrow<19.0.0
11 | polars>=1.22.0,<2.0
12 |
13 |
14 | # plotting
15 | kaleido>=0.2.1,<1.0
16 | plotly>=5.24.1,<6.0
17 | textblob>=0.19.0,<1.0
18 |
19 | # data
20 | openpyxl>=3.1.5,<4.0
21 | snowflake-connector-python>=3.12.4,<4.0
22 | google-cloud-bigquery>=3.27.0,<4.0
23 | google-auth>=2.37.0,<3.0
24 | snowflake-sqlalchemy>=1.7.3,<2.0
25 | sqlalchemy>=2.0.37,<3.0
26 | cryptography>=44.0.0,<45.0
27 | hdbcli>=2.23.27,<3.0
28 |
29 | # genai
30 | openai>=1.59.9,<2
31 | instructor>=1.3.4,<2.0
32 |
33 | # frontend
34 | streamlit==1.44.1
35 | st-theme>=1.2.3,<2.0
36 | streamlit-javascript>=0.1.5,<1.0
37 |
38 | # backend
39 | datarobot>=3.6.0,<4.0
40 | fastapi>=0.115.6,<1.0
41 | python-multipart>=0.0.20,<1.0
42 | psutil>=6.1.1,<7.0
43 | pydantic==2.7.4,<3.0
44 | pydantic-settings==2.4.0,<3.0
45 | joblib>=1.4.2,<2.0
46 | duckdb>=1.2.0,<1.3
47 | fastexcel>=0.12.1,<1.0
48 | aiofiles==24.1.0
49 |
50 | # dev & compatibility
51 | eval_type_backport>=0.2.2,<1.0
52 | db-dtypes>=1.3.1,<2.0
53 | typing-extensions>=4.12.2,<5.0
54 | numba>=0.61.0,<1.0
55 |
--------------------------------------------------------------------------------
/frontend/sap.svg:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/frontend/start-app.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 | #
3 | # Copyright 2023 DataRobot, Inc. and its affiliates.
4 | #
5 | # All rights reserved.
6 | # This is proprietary source code of DataRobot, Inc. and its affiliates.
7 | # Released under the terms of DataRobot Tool and Utility Agreement.
8 | #
9 |
10 | echo "Starting App"
11 | streamlit run 'app.py' --server.maxUploadSize 200
12 |
--------------------------------------------------------------------------------
/frontend/style.css:
--------------------------------------------------------------------------------
1 | @import url("https://fonts.googleapis.com/css2?family=Roboto");
2 |
3 | html,
4 | body,
5 | [class*="css"] {
6 | font-family: "Roboto", sans-serif;
7 | }
8 |
9 | .main {
10 | padding: 0rem 1rem;
11 | }
12 |
13 | .stProgress > div > div > div > div {
14 | background-color: #1c83e1;
15 | }
16 |
17 | .stDownloadButton button {
18 | width: 100%;
19 | }
20 |
21 | /* Add rounded corners to images */
22 | img {
23 | border-radius: 8px;
24 | }
25 | /* Remove rounded corners from DataRobot logo */
26 | .st-key-datarobot-logo img {
27 | border-radius: 0;
28 | }
29 | /* Add rounded corners to chat input */
30 | .stChatInput {
31 | border-radius: 12px;
32 | }
33 |
34 | /* Add rounded corners to chat messages */
35 | .stChatMessage {
36 | border-radius: 12px;
37 | padding-right: 0;
38 | }
39 |
40 | /* Add rounded corners to buttons */
41 | button {
42 | border-radius: 8px !important;
43 | }
44 |
45 | /* Add rounded corners to input fields */
46 | .stTextInput input {
47 | border-radius: 8px;
48 | }
49 |
50 | div[class*="st-key-delete_msg_"] {
51 | height: 40px;
52 | width: 30px;
53 | align-self: end;
54 | margin-right: 16px;
55 | }
56 |
--------------------------------------------------------------------------------
/frontend/you.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/datarobot-community/talk-to-my-data-agent/17301f5fad9192aeefa13315f179328b29135334/frontend/you.jpg
--------------------------------------------------------------------------------
/infra/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/datarobot-community/talk-to-my-data-agent/17301f5fad9192aeefa13315f179328b29135334/infra/__init__.py
--------------------------------------------------------------------------------
/infra/components/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/datarobot-community/talk-to-my-data-agent/17301f5fad9192aeefa13315f179328b29135334/infra/components/__init__.py
--------------------------------------------------------------------------------
/infra/feature_flag_requirements.yaml:
--------------------------------------------------------------------------------
1 | ---
2 | ENABLE_MLOPS: true
3 | ENABLE_CUSTOM_INFERENCE_MODEL: true
4 | ENABLE_PUBLIC_NETWORK_ACCESS_FOR_ALL_CUSTOM_MODELS: true
5 | ENABLE_MLOPS_TEXT_GENERATION_TARGET_TYPE: true
6 | ENABLE_MLOPS_RESOURCE_REQUEST_BUNDLES: true
7 |
--------------------------------------------------------------------------------
/infra/settings_database.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 | import os
15 |
16 | from utils.schema import DatabaseConnectionType
17 |
18 | # Valid values are: "snowflake", "bigquery", "sap" or "no_database"
19 | DATABASE_CONNECTION_TYPE: DatabaseConnectionType = os.getenv( # type: ignore[assignment]
20 | "DATABASE_CONNECTION_TYPE", "no_database"
21 | )
22 |
--------------------------------------------------------------------------------
/infra/settings_generative.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 | from __future__ import annotations
15 |
16 | import datarobot as dr
17 | import pulumi
18 | import pulumi_datarobot as datarobot
19 | from datarobot_pulumi_utils.pulumi.stack import PROJECT_NAME
20 | from datarobot_pulumi_utils.schema.custom_models import (
21 | CustomModelArgs,
22 | DeploymentArgs,
23 | RegisteredModelArgs,
24 | )
25 | from datarobot_pulumi_utils.schema.exec_envs import RuntimeEnvironments
26 | from datarobot_pulumi_utils.schema.llms import (
27 | LLMBlueprintArgs,
28 | LLMs,
29 | LLMSettings,
30 | PlaygroundArgs,
31 | )
32 |
33 | from utils.schema import LLMDeploymentSettings
34 |
35 | LLM = LLMs.AZURE_OPENAI_GPT_4_O
36 |
37 | custom_model_args = CustomModelArgs(
38 | resource_name=f"Generative Analyst Custom Model [{PROJECT_NAME}]",
39 | name="Generative Analyst Assistant", # built-in QA app uses this as the AI's name
40 | target_name=LLMDeploymentSettings().target_feature_name,
41 | target_type=dr.enums.TARGET_TYPE.TEXT_GENERATION,
42 | replicas=2,
43 | base_environment_id=RuntimeEnvironments.PYTHON_312_MODERATIONS.value.id,
44 | opts=pulumi.ResourceOptions(delete_before_replace=True),
45 | )
46 |
47 | registered_model_args = RegisteredModelArgs(
48 | resource_name=f"Generative Analyst Registered Model [{PROJECT_NAME}]",
49 | )
50 |
51 |
52 | deployment_args = DeploymentArgs(
53 | resource_name=f"Generative Analyst Deployment [{PROJECT_NAME}]",
54 | label=f"Generative Analyst Deployment [{PROJECT_NAME}]",
55 | association_id_settings=datarobot.DeploymentAssociationIdSettingsArgs(
56 | column_names=["association_id"],
57 | auto_generate_id=False,
58 | required_in_prediction_requests=True,
59 | ),
60 | predictions_data_collection_settings=datarobot.DeploymentPredictionsDataCollectionSettingsArgs(
61 | enabled=True,
62 | ),
63 | predictions_settings=(
64 | datarobot.DeploymentPredictionsSettingsArgs(min_computes=0, max_computes=2)
65 | ),
66 | )
67 |
68 | playground_args = PlaygroundArgs(
69 | resource_name=f"Generative Analyst Playground [{PROJECT_NAME}]",
70 | )
71 |
72 | llm_blueprint_args = LLMBlueprintArgs(
73 | resource_name=f"Generative Analyst LLM Blueprint [{PROJECT_NAME}]",
74 | llm_id=LLM.name,
75 | llm_settings=LLMSettings(
76 | max_completion_length=2048,
77 | temperature=0.1,
78 | ),
79 | )
80 |
--------------------------------------------------------------------------------
/infra/settings_main.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 | from pathlib import Path
15 |
16 | PROJECT_ROOT = Path(__file__).resolve().parent.parent.absolute()
17 |
--------------------------------------------------------------------------------
/infra/settings_proxy_llm.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 | from __future__ import annotations
15 |
16 | import os
17 |
18 | import datarobot as dr
19 | import pulumi
20 | from datarobot_pulumi_utils.pulumi.stack import PROJECT_NAME
21 | from datarobot_pulumi_utils.schema.custom_models import (
22 | CustomModelArgs,
23 | DeploymentArgs,
24 | RegisteredModelArgs,
25 | )
26 | from datarobot_pulumi_utils.schema.exec_envs import RuntimeEnvironments
27 |
28 | from utils.schema import LLMDeploymentSettings
29 |
30 | CHAT_MODEL_NAME = os.getenv("CHAT_MODEL_NAME")
31 |
32 | custom_model_args = CustomModelArgs(
33 | resource_name=f"Data Analyst Proxy LLM Custom Model [{PROJECT_NAME}]",
34 | name=f"Data Analyst Proxy LLM Custom Model [{PROJECT_NAME}]",
35 | target_name=LLMDeploymentSettings().target_feature_name,
36 | target_type=dr.enums.TARGET_TYPE.TEXT_GENERATION,
37 | replicas=2,
38 | base_environment_id=RuntimeEnvironments.PYTHON_312_MODERATIONS.value.id,
39 | opts=pulumi.ResourceOptions(delete_before_replace=True),
40 | )
41 |
42 | registered_model_args = RegisteredModelArgs(
43 | resource_name=f"Data Analyst Proxy LLM Registered Model [{PROJECT_NAME}]",
44 | )
45 |
46 | deployment_args = DeploymentArgs(
47 | resource_name=f"Data Analyst Proxy LLM Deployment [{PROJECT_NAME}]",
48 | label=f"Data Analyst Proxy LLM Deployment [{PROJECT_NAME}]",
49 | )
50 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | # infra
2 | pulumi==3.153.0
3 | pulumi-datarobot>=0.8.15,<0.9
4 | datarobot-pulumi-utils>=0.0.2.post8,<0.1
5 | jinja2>=3.1.5,<4.0
6 | PyYAML>=6.0.2,<7.0
7 |
8 | # dev
9 | mypy>=1.14.1,<2.0
10 | pandas-stubs>=2.2.3.241126,<3.0
11 | types-requests>=2.32.0.20241016,<3.0
12 | types-pyyaml>=6.0.12.20241230,<7.0
13 | ruff>=0.9,<1.0
14 | pytest>=8.3.4,<9.0
15 | pytest-cov>=6.0.0,<7.0
16 | pytest-mock>=3.14.0,<4.0
17 | pytest-asyncio>=0.25.2,<1.0
18 | pytest-httpx>=0.30.0,<0.31.0
19 |
20 | # genai testing
21 | rich>=13.9.4,<14.0
22 | inquirer>=3.4.0,<4.0
23 | nbformat>=5.10.0,<6.0
24 |
25 | # ======= frontend/requirements.txt ======= #
26 |
27 | # python sandbox libraries
28 | pandas>=2.2.3,<3.0
29 | numpy>=2.1.3,<3.0
30 | scipy>=1.15.1,<2.0
31 | statsmodels>=0.14.4,<1.0
32 | scikit-learn>=1.6.1,<2.0
33 | lightgbm>=4.5.0,<5.0
34 | tslearn>=0.6.3,<1.0
35 | spacy>=3.8.3,<4.0
36 | pyarrow<19.0.0
37 | polars>=1.22.0,<2.0
38 |
39 | # plotting
40 | kaleido==0.2.0
41 | plotly>=5.24.1,<6.0
42 | textblob>=0.19.0,<1.0
43 |
44 | # data
45 | openpyxl>=3.1.5,<4.0
46 | snowflake-connector-python>=3.12.4,<4.0
47 | google-cloud-bigquery>=3.27.0,<4.0
48 | google-auth>=2.37.0,<3.0
49 | snowflake-sqlalchemy>=1.7.3,<2.0
50 | sqlalchemy>=2.0.37,<3.0
51 | cryptography>=44.0.0,<45.0
52 | hdbcli>=2.23.27,<3.0
53 |
54 | # genai
55 | openai>=1.59.9,<2
56 | instructor>=1.3.4,<2.0
57 | boto3>=1.36.2,<2.0
58 |
59 | # frontend
60 | streamlit==1.44.1
61 | st-theme>=1.2.3,<2.0
62 | streamlit-javascript>=0.1.5,<1.0
63 |
64 | # backend
65 | datarobot>=3.6.0,<4.0
66 | fastapi>=0.115.6,<1.0
67 | python-multipart>=0.0.20,<1.0
68 | uvicorn==0.34.2,<1.0
69 | psutil>=6.1.1,<7.0
70 | pydantic==2.7.4,<3.0
71 | pydantic-settings==2.4.0,<3.0
72 | joblib>=1.4.2,<2.0
73 | duckdb>=1.2.0,<1.3
74 | fastexcel>=0.12.1,<1.0
75 | aiofiles==24.1.0
76 | types-aiofiles==24.1.0.20241221
77 | httpx>=0.23.0,<1.0,<1.0
78 | pillow==11.2.1
79 |
80 | # dev & compatibility
81 | eval_type_backport>=0.2.2,<1.0
82 | db-dtypes>=1.3.1,<2.0
83 | typing-extensions>=4.12.2,<5.0
84 | numba>=0.61.0,<1.0
85 |
--------------------------------------------------------------------------------
/set_env.bat:
--------------------------------------------------------------------------------
1 | @echo off
2 | setlocal enabledelayedexpansion
3 | :: Define LF (Line Feed) at the start of your script
4 | set ^"LF=^
5 |
6 | ^" The empty line above is critical - do not remove
7 |
8 | :: Function to load .env file
9 | call :load_env_file || exit /b
10 |
11 | endlocal
12 |
13 | :: Function to activate the virtual environment if it exists
14 | call :activate_virtual_environment
15 | call :set_env_vars
16 |
17 | exit /b
18 |
19 |
20 | :load_env_file
21 | if not exist .env (
22 | echo Error: .env file not found.
23 | echo Please create a .env file with VAR_NAME=value pairs.
24 | exit /b 1
25 | )
26 | @echo off
27 | python -c "from quickstart import load_dotenv; env_vars = load_dotenv(); f = open('set_env_vars.bat', 'w'); [f.write(f'set \"{key}={value.replace(\"\n\", \" \").replace(\"\\r\", \" \")}\"\n') for key, value in env_vars.items()]; f.close()"
28 | exit /b 0
29 |
30 | :activate_virtual_environment
31 | if exist .venv\Scripts\activate.bat (
32 | echo Activated virtual environment found at .venv\Scripts\activate.bat
33 | call .venv\Scripts\activate.bat
34 | )
35 | exit /b 0
36 |
37 | :set_env_vars
38 | endlocal
39 | call set_env_vars.bat
40 | del set_env_vars.bat
41 | echo Environment variables have been set.
42 | exit /b 0
43 |
44 |
--------------------------------------------------------------------------------
/set_env.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | PULUMI_MANAGED_VENV="runtime: name: python options: toolchain: pip virtualenv: .venv"
4 |
5 | MANUAL_DEP_INSTRUCTIONS="Please comment out the following three lines 'Pulumi.yaml' to enable manual dependency management with Pulumi:
6 | # options:
7 | # toolchain: pip
8 | # virtualenv: .venv
9 |
10 | Then install dependencies manually (make sure to address any conflicts pip identifies):
11 | pip install -r requirements.txt"
12 |
13 | # Function to check if script is sourced
14 | check_if_sourced() {
15 | if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
16 | echo "Please run this script using 'source' or '.' like this:"
17 | echo " source ${0}"
18 | echo " . ${0}"
19 | return 1
20 | fi
21 | }
22 |
23 | # Function to check for DataRobot Codespaces environment
24 | check_datarobot_environment() {
25 | if [[ -n "${DATAROBOT_NOTEBOOK_IMAGE}" ]]; then
26 | if tr -d '\n' < Pulumi.yaml | grep -Fq "$PULUMI_MANAGED_VENV"; then
27 | echo "DR Codespaces requires manual management of dependencies."
28 | echo "$MANUAL_DEP_INSTRUCTIONS"
29 | return 1
30 | fi
31 | fi
32 | }
33 |
34 | # Function to check for active conda environment
35 | check_active_conda_env() {
36 | if [[ -n "${CONDA_DEFAULT_ENV}" ]]; then
37 | if tr -d '\n' < Pulumi.yaml | grep -Fq "$PULUMI_MANAGED_VENV"; then
38 | echo "Using Pulumi with conda requires manual management of dependencies."
39 | echo "$MANUAL_DEP_INSTRUCTIONS"
40 | return 1
41 | fi
42 | fi
43 | }
44 |
45 |
46 | # Function to load .env file
47 | load_env_file() {
48 | if [ ! -f .env ]; then
49 | echo "Error: .env file not found."
50 | echo "Please create a .env file with VAR_NAME=value pairs."
51 | return 1
52 | fi
53 | set -a
54 | source .env
55 | set +a
56 | echo "Environment variables from .env have been set."
57 | }
58 |
59 | # Function to activate the virtual environment if it exists
60 | activate_virtual_environment() {
61 | if [ -f .venv/bin/activate ]; then
62 | echo "Activated virtual environment found at .venv/bin/activate"
63 | source .venv/bin/activate
64 | fi
65 | }
66 |
67 | # Main execution
68 | if check_if_sourced ; then
69 | load_env_file
70 | activate_virtual_environment
71 | check_datarobot_environment
72 | check_active_conda_env
73 | fi
74 |
--------------------------------------------------------------------------------
/test_pydantic.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 | from typing import Any, Hashable
16 |
17 | import pandas as pd
18 | from pandas.testing import assert_frame_equal
19 | from pydantic import BaseModel, ConfigDict, Field, computed_field, field_validator
20 |
21 |
22 | class AnalystDataset(BaseModel):
23 | name: str
24 | data: pd.DataFrame = Field(default_factory=pd.DataFrame) # Removed exclude=True
25 |
26 | model_config = ConfigDict(
27 | arbitrary_types_allowed=True,
28 | json_encoders={pd.DataFrame: lambda df: df.to_dict(orient="records")},
29 | )
30 |
31 | @field_validator("data", mode="before")
32 | @classmethod
33 | def validate_dataframe(cls, v: Any) -> pd.DataFrame:
34 | if isinstance(v, pd.DataFrame):
35 | return v
36 | elif isinstance(v, list):
37 | try:
38 | return pd.DataFrame.from_records(v)
39 | except Exception as e:
40 | raise ValueError("Invalid data format") from e
41 | else:
42 | raise ValueError("data has to be either list of records or pd.DataFrame")
43 |
44 | def to_df(self) -> pd.DataFrame:
45 | return self.data
46 |
47 | @property
48 | def columns(self) -> list[str]:
49 | return self.data.columns.tolist()
50 |
51 | @computed_field
52 | def data_records(self) -> list[dict[Hashable, Any]]:
53 | return self.data.to_dict(orient="records")
54 |
55 |
56 | # Example usage:
57 | if __name__ == "__main__":
58 | # Create a DataFrame.
59 | df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
60 | data = {"df": df}
61 |
62 | # Initialize the model.
63 | model = AnalystDataset(name="test", data=data["df"])
64 |
65 | # Serialize the model to JSON.
66 | serialized = model.model_dump_json()
67 | print(serialized)
68 |
69 | # Deserialize back into an AnalystDataset instance.
70 | deserialized = AnalystDataset.model_validate_json(serialized)
71 |
72 | # Verify that the deserialized DataFrame matches the original.
73 | assert_frame_equal(deserialized.to_df(), data["df"])
74 | print("DataFrame round-trip successful!")
75 |
--------------------------------------------------------------------------------
/trivy-ignore.rego:
--------------------------------------------------------------------------------
1 | package trivy
2 |
3 | import data.lib.trivy
4 |
5 | default ignore := false
6 |
--------------------------------------------------------------------------------
/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/datarobot-community/talk-to-my-data-agent/17301f5fad9192aeefa13315f179328b29135334/utils/__init__.py
--------------------------------------------------------------------------------
/utils/resources.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 | from __future__ import annotations
15 |
16 | import json
17 | import subprocess
18 | from typing import Any, Dict, Mapping, Tuple, Type, Union
19 |
20 | from pydantic import AliasChoices, Field
21 | from pydantic_settings import (
22 | BaseSettings,
23 | EnvSettingsSource,
24 | PydanticBaseSettingsSource,
25 | SettingsConfigDict,
26 | )
27 | from pydantic_settings.sources import parse_env_vars
28 |
29 |
30 | class PulumiSettingsSource(EnvSettingsSource):
31 | """Pulumi stack outputs as a pydantic settings source."""
32 |
33 | _PULUMI_OUTPUTS: Dict[str, str] = {}
34 | _PULUMI_CALLED: bool = False
35 |
36 | def __init__(self, *args: Any, **kwargs: Any) -> None:
37 | self.read_pulumi_outputs()
38 | super().__init__(*args, **kwargs)
39 |
40 | def read_pulumi_outputs(self) -> None:
41 | try:
42 | raw_outputs = json.loads(
43 | subprocess.check_output(
44 | ["pulumi", "stack", "output", "-j"],
45 | text=True,
46 | ).strip()
47 | )
48 | self._PULUMI_OUTPUTS = {
49 | k: v if isinstance(v, str) else json.dumps(v)
50 | for k, v in raw_outputs.items()
51 | }
52 | except BaseException:
53 | self._PULUMI_OUTPUTS = {}
54 |
55 | def _load_env_vars(self) -> Mapping[str, Union[str, None]]:
56 | return parse_env_vars(
57 | self._PULUMI_OUTPUTS,
58 | self.case_sensitive,
59 | self.env_ignore_empty,
60 | self.env_parse_none_str,
61 | )
62 |
63 |
64 | class DynamicSettings(BaseSettings):
65 | """Settings that come from pulumi stack outputs or DR runtime parameters"""
66 |
67 | model_config = SettingsConfigDict(extra="ignore", populate_by_name=True)
68 |
69 | @classmethod
70 | def settings_customise_sources(
71 | cls,
72 | settings_cls: Type[BaseSettings],
73 | init_settings: PydanticBaseSettingsSource,
74 | env_settings: PydanticBaseSettingsSource,
75 | dotenv_settings: PydanticBaseSettingsSource,
76 | file_secret_settings: PydanticBaseSettingsSource,
77 | ) -> Tuple[PydanticBaseSettingsSource, ...]:
78 | return (
79 | init_settings,
80 | PulumiSettingsSource(settings_cls),
81 | env_settings,
82 | )
83 |
84 |
85 | app_settings_env_name: str = "ANALYST_APP_SETTINGS"
86 |
87 | app_env_name: str = "DATAROBOT_APPLICATION_ID"
88 |
89 | llm_deployment_env_name: str = "LLM_DEPLOYMENT_ID"
90 |
91 |
92 | class LLMDeployment(DynamicSettings):
93 | id: str = Field(
94 | validation_alias=AliasChoices(
95 | "MLOPS_RUNTIME_PARAM_" + llm_deployment_env_name,
96 | llm_deployment_env_name,
97 | )
98 | )
99 |
--------------------------------------------------------------------------------
/utils/tools.py:
--------------------------------------------------------------------------------
1 | # Copyright 2024 DataRobot, Inc.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 |
16 | # Add additional tools that can be used by the analysis code execution. Remember to include the necessary imports and provide a docstring and signature for each function.
17 | # signature and docstring will be provided to the LLM in the prompt.
18 | # Uncomment the examples below to get started.
19 |
20 |
21 | # import datarobot as dr
22 | # import pandas as pd
23 | # from datarobot_predict.deployment import predict
24 |
25 | # def calculate_summary(df: pd.DataFrame) -> pd.DataFrame:
26 | # """
27 | # Calculate summary statistics for a DataFrame.
28 |
29 | # Args:
30 | # df (pd.DataFrame): Input DataFrame.
31 |
32 | # Returns:
33 | # pd.DataFrame: Summary statistics including count, mean, std, min, max, and percentiles.
34 | # """
35 | # description = df.describe(percentiles=[0.2, 0.4, 0.6, 0.8])
36 | # return description
37 |
38 |
39 | # def filter_data(df: pd.DataFrame, column: str, value: float) -> pd.DataFrame:
40 | # """
41 | # Filter DataFrame based on a condition.
42 | # Args:
43 | # df (pd.DataFrame): Input DataFrame.
44 | # column (str): Column to apply the filter on.
45 | # value (float): Value to compare against in the filter.
46 | # Returns:
47 | # pd.DataFrame: Filtered DataFrame where the specified column's values are greater than the given value.
48 | # """
49 | # filtered_df = df[df[column] > value]
50 | # return filtered_df
51 |
52 |
53 | # def call_datarobot_deployment(df: pd.DataFrame, deployment_id: str) -> pd.DataFrame:
54 | # """
55 | # Call a DataRobot deployment to get predictions.
56 |
57 | # Args:
58 | # df (pd.DataFrame): Input DataFrame with features for prediction.
59 | # deployment_id (str): ID of the DataRobot deployment to use for predictions.
60 |
61 | # Returns:
62 | # pd.DataFrame: DataFrame containing the predictions from DataRobot. The prediction column is named 'predictions'.
63 | # """
64 | # deployment = dr.Deployment.get(deployment_id) # type: ignore[attr-defined]
65 | # prediction_response: pd.DataFrame = predict(
66 | # deployment=deployment, data_frame=df
67 | # ).dataframe
68 |
69 | # prediction_response.columns = [
70 | # c.replace("_PREDICTION", "")
71 | # for c in prediction_response.columns # type: ignore[assignment]
72 | # ]
73 |
74 | # if deployment.model is not None:
75 | # target_column = deployment.model.get("target_name")
76 | # if target_column:
77 | # prediction_response["predictions"] = prediction_response[target_column]
78 |
79 | # return prediction_response
80 |
--------------------------------------------------------------------------------
| 
navigate(ROUTES.DATA)}
17 | />
18 | 
7 | Loading...
8 |