feat: add auth to chatui (#1065)

letta-ai · Feb 28, 2024 · 397ddb8 · 397ddb8
1 parent 7d61d32
commit 397ddb8
Show file tree

Hide file tree

Showing 27 changed files with 425 additions and 174 deletions.
diff --git a/chatui/src/app/auth.tsx b/chatui/src/app/auth.tsx
@@ -1,13 +1,15 @@
 import { PropsWithChildren } from 'react';
+import { useLocation, useNavigate } from 'react-router-dom';
 import { useAuthStoreActions, useAuthStoreState } from './libs/auth/auth.store';
-import { useAuthQuery } from './libs/auth/use-auth.query';
 
 const Auth = (props: PropsWithChildren) => {
-	const result = useAuthQuery();
-	const { uuid } = useAuthStoreState();
-	const { setAsAuthenticated } = useAuthStoreActions();
-	if (result.isSuccess && uuid !== result.data.uuid) {
-		setAsAuthenticated(result.data.uuid);
+	const { loggedIn } = useAuthStoreState();
+	const { logout } = useAuthStoreActions();
+	const location = useLocation();
+	const navigate = useNavigate();
+	if (!loggedIn && location.pathname !== '/login') {
+		logout();
+		navigate('/login');
 	}
 	return props.children;
 };

diff --git a/chatui/src/app/libs/agents/agent.ts b/chatui/src/app/libs/agents/agent.ts
@@ -6,6 +6,16 @@ export const AgentSchema = z.object({
 	human: z.string(),
 	persona: z.string(),
 	created_at: z.string(),
+	// TODO: Remove optional once API response returns necessary data
+	memories: z.number().optional(),
+	data_sources: z.number().optional(),
+	last_run: z.string().optional(),
+	tools: z
+		.object({
+			core: z.number(),
+			user_defined: z.number(),
+		})
+		.optional(),
 });
 
 export type Agent = z.infer<typeof AgentSchema>;
diff --git a/chatui/src/app/libs/agents/use-agent-memory.mutation.ts b/chatui/src/app/libs/agents/use-agent-memory.mutation.ts
@@ -1,17 +1,20 @@
 import { useMutation, useQueryClient } from '@tanstack/react-query';
+import { useAuthBearerToken } from '../auth/auth.store';
 import { API_BASE_URL } from '../constants';
 import { AgentMemoryUpdate } from './agent-memory-update';
 
 export const useAgentMemoryUpdateMutation = (userId: string | null | undefined) => {
 	const queryClient = useQueryClient();
+	const bearerToken = useAuthBearerToken();
+
 	return useMutation({
 		mutationFn: async (params: AgentMemoryUpdate) =>
 			await fetch(API_BASE_URL + `/agents/memory?${userId}`, {
 				method: 'POST',
-				headers: { 'Content-Type': ' application/json' },
+				headers: { 'Content-Type': ' application/json', Authorization: bearerToken },
 				body: JSON.stringify(params),
 			}).then((res) => res.json()),
-		onSuccess: (res, { agent_id }) =>
+		onSuccess: (_, { agent_id }) =>
 			queryClient.invalidateQueries({ queryKey: [userId, 'agents', 'entry', agent_id, 'memory'] }),
 	});
 };
diff --git a/chatui/src/app/libs/agents/use-agent-memory.query.ts b/chatui/src/app/libs/agents/use-agent-memory.query.ts
@@ -1,13 +1,18 @@
 import { useQuery } from '@tanstack/react-query';
+import { useAuthBearerToken } from '../auth/auth.store';
 import { API_BASE_URL } from '../constants';
 import { AgentMemory } from './agent-memory';
 
-export const useAgentMemoryQuery = (userId: string | null | undefined, agentId: string | null | undefined) =>
-	useQuery({
+export const useAgentMemoryQuery = (userId: string | null | undefined, agentId: string | null | undefined) => {
+	const bearerToken = useAuthBearerToken();
+	return useQuery({
 		queryKey: [userId, 'agents', 'entry', agentId, 'memory'],
 		queryFn: async () =>
-			(await fetch(API_BASE_URL + `/agents/memory?agent_id=${agentId}&user_id=${userId}`).then((res) =>
-				res.json()
-			)) as Promise<AgentMemory>,
+			(await fetch(API_BASE_URL + `/agents/memory?agent_id=${agentId}&user_id=${userId}`, {
+				headers: {
+					Authorization: bearerToken,
+				},
+			}).then((res) => res.json())) as Promise<AgentMemory>,
 		enabled: !!userId && !!agentId,
 	});
+};
diff --git a/chatui/src/app/libs/agents/use-agents.mutation.ts b/chatui/src/app/libs/agents/use-agents.mutation.ts
@@ -1,14 +1,16 @@
 import { useMutation, useQueryClient } from '@tanstack/react-query';
+import { useAuthBearerToken } from '../auth/auth.store';
 import { API_BASE_URL } from '../constants';
 import { Agent } from './agent';
 
 export const useAgentsCreateMutation = (userId: string | null | undefined) => {
 	const queryClient = useQueryClient();
+	const bearerToken = useAuthBearerToken();
 	return useMutation({
-		mutationFn: async (params: { name: string; human: string; persona: string; model: string }) => {
+		mutationFn: async (params: { name: string; human: string; persona: string; model: string }): Promise<Agent> => {
 			const response = await fetch(API_BASE_URL + '/agents', {
 				method: 'POST',
-				headers: { 'Content-Type': ' application/json' },
+				headers: { 'Content-Type': ' application/json', Authorization: bearerToken },
 				body: JSON.stringify({ config: params, user_id: userId }),
 			});
 
@@ -18,7 +20,7 @@ export const useAgentsCreateMutation = (userId: string | null | undefined) => {
 				throw new Error(errorBody || 'Error creating agent');
 			}
 
-			return response.json() as Promise<Agent>;
+			return await response.json();
 		},
 		onSuccess: () => queryClient.invalidateQueries({ queryKey: [userId, 'agents', 'list'] }),
 	});

diff --git a/chatui/src/app/libs/agents/use-agents.query.ts b/chatui/src/app/libs/agents/use-agents.query.ts
@@ -1,14 +1,21 @@
 import { useQuery } from '@tanstack/react-query';
+import { useAuthBearerToken } from '../auth/auth.store';
 import { API_BASE_URL } from '../constants';
 import { Agent } from './agent';
 
-export const useAgentsQuery = (userId: string | null | undefined) =>
-	useQuery({
+export const useAgentsQuery = (userId: string | null | undefined) => {
+	const bearerToken = useAuthBearerToken();
+	return useQuery({
 		queryKey: [userId, 'agents', 'list'],
 		enabled: !!userId,
 		queryFn: async () =>
-			(await fetch(API_BASE_URL + `/agents?user_id=${userId}`).then((res) => res.json())) as Promise<{
+			(await fetch(API_BASE_URL + `/agents?user_id=${userId}`, {
+				headers: {
+					Authorization: bearerToken,
+				},
+			}).then((res) => res.json())) as Promise<{
 				num_agents: number;
 				agents: Agent[];
 			}>,
 	});
+};
diff --git a/chatui/src/app/libs/auth/auth.store.ts b/chatui/src/app/libs/auth/auth.store.ts
@@ -3,19 +3,44 @@ import { createJSONStorage, persist } from 'zustand/middleware';
 
 export type AuthState = {
 	uuid: string | null;
+	token: string | null;
+	loggedIn: boolean;
+};
+export type AuthActions = {
+	setAsAuthenticated: (uuid: string, token?: string) => void;
+	setToken: (token: string) => void;
+	logout: () => void;
 };
-export type AuthActions = { setAsAuthenticated: (uuid: string) => void };
 
 const useAuthStore = create(
 	persist<{ auth: AuthState; actions: AuthActions }>(
 		(set, get) => ({
-			auth: { uuid: null },
+			auth: { uuid: null, token: null, loggedIn: false },
 			actions: {
-				setAsAuthenticated: (uuid: string) =>
+				setToken: (token: string) =>
+					set((prev) => ({
+						...prev,
+						auth: {
+							...prev.auth,
+							token,
+						},
+					})),
+				setAsAuthenticated: (uuid: string, token?: string) =>
 					set((prev) => ({
 						...prev,
 						auth: {
+							token: token ?? prev.auth.token,
 							uuid,
+							loggedIn: true,
+						},
+					})),
+				logout: () =>
+					set((prev) => ({
+						...prev,
+						auth: {
+							token: null,
+							uuid: null,
+							loggedIn: false,
 						},
 					})),
 			},
@@ -30,3 +55,7 @@ const useAuthStore = create(
 
 export const useAuthStoreState = () => useAuthStore().auth;
 export const useAuthStoreActions = () => useAuthStore().actions;
+export const useAuthBearerToken = () => {
+	const { auth } = useAuthStore();
+	return auth.token ? `Bearer ${auth.token}` : '';
+};
diff --git a/chatui/src/app/libs/auth/use-auth.mutation.ts b/chatui/src/app/libs/auth/use-auth.mutation.ts
@@ -0,0 +1,23 @@
+import { useMutation } from '@tanstack/react-query';
+import { API_BASE_URL } from '../constants';
+import { useAuthStoreActions } from './auth.store';
+
+export type AuthResponse = { uuid: string };
+export const useAuthMutation = () => {
+	const { setAsAuthenticated } = useAuthStoreActions();
+	return useMutation({
+		mutationKey: ['auth'],
+		mutationFn: (password: string) =>
+			fetch(API_BASE_URL + `/auth`, {
+				method: 'POST',
+				headers: { 'Content-Type': ' application/json' },
+				body: JSON.stringify({ password }),
+			}).then((res) => {
+				if (!res.ok) {
+					throw new Error('Network response was not ok');
+				}
+				return res.json();
+			}) as Promise<AuthResponse>,
+		onSuccess: (data, password) => setAsAuthenticated(data.uuid, password),
+	});
+};
diff --git a/chatui/src/app/libs/humans/use-humans.mutation.ts b/chatui/src/app/libs/humans/use-humans.mutation.ts
@@ -0,0 +1,27 @@
+import { useMutation, useQueryClient } from '@tanstack/react-query';
+import { useAuthBearerToken } from '../auth/auth.store';
+import { API_BASE_URL } from '../constants';
+import { Human } from './human';
+
+export const useHumansCreateMutation = (userId: string | null | undefined) => {
+	const queryClient = useQueryClient();
+	const bearerToken = useAuthBearerToken();
+	return useMutation({
+		mutationFn: async (params: { name: string; text: string }): Promise<Human> => {
+			const response = await fetch(API_BASE_URL + '/agents', {
+				method: 'POST',
+				headers: { 'Content-Type': ' application/json', Authorization: bearerToken },
+				body: JSON.stringify({ config: params, user_id: userId }),
+			});
+
+			if (!response.ok) {
+				// Throw an error if the response is not OK
+				const errorBody = await response.text();
+				throw new Error(errorBody || 'Error creating human');
+			}
+
+			return await response.json();
+		},
+		onSuccess: () => queryClient.invalidateQueries({ queryKey: [userId, 'humans', 'list'] }),
+	});
+};
diff --git a/chatui/src/app/libs/humans/use-humans.query.ts b/chatui/src/app/libs/humans/use-humans.query.ts
@@ -1,13 +1,19 @@
 import { useQuery } from '@tanstack/react-query';
+import { useAuthBearerToken } from '../auth/auth.store';
 import { API_BASE_URL } from '../constants';
 import { Human } from './human';
 
-export const useHumansQuery = (userId: string | null | undefined) =>
-	useQuery({
+export const useHumansQuery = (userId: string | null | undefined) => {
+	const bearerToken = useAuthBearerToken();
+	return useQuery({
 		queryKey: [userId, 'humans', 'list'],
 		enabled: !!userId,
 		queryFn: async () => {
-			const response = await fetch(`${API_BASE_URL}/humans?user_id=${encodeURIComponent(userId || '')}`);
+			const response = await fetch(`${API_BASE_URL}/humans?user_id=${encodeURIComponent(userId || '')}`, {
+				headers: {
+					Authorization: bearerToken,
+				},
+			});
 			if (!response.ok) {
 				throw new Error('Network response was not ok for fetching humans');
 			}
@@ -16,3 +22,4 @@ export const useHumansQuery = (userId: string | null | undefined) =>
 			}>;
 		},
 	});
+};
diff --git a/chatui/src/app/libs/messages/message-history.store.ts b/chatui/src/app/libs/messages/message-history.store.ts
@@ -13,7 +13,7 @@ export type MessageHistory = {
 
 const useMessageHistoryStore = create(
 	persist<{ history: MessageHistory; actions: { addMessage: (key: string, message: Message) => void } }>(
-		(set, get) => ({
+		(set) => ({
 			history: {},
 			actions: {
 				addMessage: (key: string, message: Message) =>

diff --git a/chatui/src/app/libs/messages/message-stream.store.ts b/chatui/src/app/libs/messages/message-stream.store.ts
@@ -18,7 +18,7 @@ const useMessageStreamStore = create(
 		{
 			socket: null as EventSource | null,
 			socketURL: null as string | null,
-			readyState: ReadyState.IDLE,
+			readyState: ReadyState.IDLE as ReadyState,
 			abortController: null as AbortController | null,
 			onMessageCallback: ((message: Message) =>
 				console.warn('No message callback set up. Simply logging message', message)) as (message: Message) => void,
@@ -30,10 +30,12 @@ const useMessageStreamStore = create(
 					agentId,
 					message,
 					role,
+					bearerToken,
 				}: {
 					userId: string;
 					agentId: string;
 					message: string;
+					bearerToken: string;
 					role?: 'user' | 'system';
 				}) => {
 					const abortController = new AbortController();
@@ -49,7 +51,7 @@ const useMessageStreamStore = create(
 						});
 					void fetchEventSource(ENDPOINT_URL, {
 						method: 'POST',
-						headers: { 'Content-Type': 'application/json', Accept: 'text/event-stream' },
+						headers: { 'Content-Type': 'application/json', Accept: 'text/event-stream', Authorization: bearerToken },
 						body: JSON.stringify({
 							user_id: userId,
 							agent_id: agentId,

diff --git a/chatui/src/app/libs/messages/use-messages.query.ts b/chatui/src/app/libs/messages/use-messages.query.ts
@@ -1,19 +1,27 @@
 import { useQuery } from '@tanstack/react-query';
+import { useAuthBearerToken } from '../auth/auth.store';
 import { API_BASE_URL } from '../constants';
 
 export const useMessagesQuery = (
 	userId: string | null | undefined,
 	agentId: string | null | undefined,
 	start = 0,
 	count = 10
-) =>
-	useQuery({
+) => {
+	const bearerToken = useAuthBearerToken();
+	return useQuery({
 		queryKey: [userId, 'agents', 'item', agentId, 'messages', 'list', start, count],
 		queryFn: async () =>
 			(await fetch(
-				API_BASE_URL + `/agents/message?agent_id=${agentId}&user_id=${userId}&start=${start}&count=${count}`
+				API_BASE_URL + `/agents/message?agent_id=${agentId}&user_id=${userId}&start=${start}&count=${count}`,
+				{
+					headers: {
+						Authorization: bearerToken,
+					},
+				}
 			).then((res) => res.json())) as Promise<{
 				messages: { role: string; name: string; content: string }[];
 			}>,
 		enabled: !!userId && !!agentId,
 	});
+};
diff --git a/chatui/src/app/libs/models/use-models.query.ts b/chatui/src/app/libs/models/use-models.query.ts
@@ -1,13 +1,19 @@
 import { useQuery } from '@tanstack/react-query';
+import { useAuthBearerToken } from '../auth/auth.store';
 import { API_BASE_URL } from '../constants';
 import { Model } from './model';
 
-export const useModelsQuery = (userId: string | null | undefined) =>
-	useQuery({
+export const useModelsQuery = (userId: string | null | undefined) => {
+	const bearerToken = useAuthBearerToken();
+	return useQuery({
 		queryKey: [userId, 'models', 'list'],
 		enabled: !!userId,
 		queryFn: async () => {
-			const response = await fetch(`${API_BASE_URL}/models?user_id=${encodeURIComponent(userId || '')}`);
+			const response = await fetch(`${API_BASE_URL}/models?user_id=${encodeURIComponent(userId || '')}`, {
+				headers: {
+					Authorization: bearerToken,
+				},
+			});
 			if (!response.ok) {
 				throw new Error('Network response was not ok for fetching models');
 			}
@@ -16,3 +22,4 @@ export const useModelsQuery = (userId: string | null | undefined) =>
 			}>;
 		},
 	});
+};