Adds AssistantMessage class.

roseflow-ai · Feb 18, 2024 · b18452f · b18452f
1 parent ae9b867
commit b18452f
Showing 1 changed file with 13 additions and 0 deletions.
diff --git 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