diff --git a/src/lib/types.ts b/src/lib/types.ts
index 05cf753..dfc3b72 100644
--- a/src/lib/types.ts
+++ b/src/lib/types.ts
@@ -1,5 +1,18 @@
-export type ChatTurnMessage = {
-  role: 'user' | 'assistant' | 'system';
+import { ToolCall } from './models/types';
+
+export type SystemMessage = {
+  role: 'system';
+  content: string;
+};
+
+export type AssistantMessage = {
+  role: 'assistant';
+  content: string;
+  tool_calls?: ToolCall[];
+};
+
+export type UserMessage = {
+  role: 'user';
   content: string;
 };
 
@@ -10,7 +23,13 @@ export type ToolMessage = {
   content: string;
 };
 
-export type Message = ChatTurnMessage | ToolMessage;
+export type ChatTurnMessage = UserMessage | AssistantMessage;
+
+export type Message =
+  | UserMessage
+  | AssistantMessage
+  | SystemMessage
+  | ToolMessage;
 
 export type Chunk = {
   content: string;