From b18452f1f25e00922dc10a073658f27411c36789 Mon Sep 17 00:00:00 2001
From: Lauri Jutila <git@laurijutila.com>
Date: Sun, 18 Feb 2024 14:40:51 +0200
Subject: [PATCH] Adds AssistantMessage class.

---
 lib/roseflow/chat/message.rb | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/lib/roseflow/chat/message.rb b/lib/roseflow/chat/message.rb
index a8a66ec..079a6bc 100644
--- a/lib/roseflow/chat/message.rb
+++ b/lib/roseflow/chat/message.rb
@@ -2,6 +2,7 @@
 
 require "dry-struct"
 require "roseflow/types"
+require "roseflow/prompt"
 
 module Roseflow
   module Chat
@@ -51,5 +52,17 @@ def self.from(input)
         )
       end
     end
+
+    class AssistantMessage < Message
+      attribute :role, Types::String.constrained(included_in: %w(assistant))
+
+      def self.from(input)
+        input = input.is_a?(Roseflow::Prompt) ? input.call : input
+        new(
+          role: "assistant",
+          content: input
+        )
+      end
+    end
   end
 end