raif 1.0.0 → 1.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 232a25681c254ff213ce9565506bffc805df041361e145016eee17f0ca3514f3
-  data.tar.gz: 24c54f8c2d5ea95664a55a57b7ed8c3e22d3040f79f28311eb60083aa2d9f47e
+  metadata.gz: 1af0a7a003990a40716f3d07ec87838e5fb485147bb87a9b35ae1fe23f642501
+  data.tar.gz: b82bba67ed23a6e47e1bd02f5c64f9448c7374017c097a012ccdae2e3795be34
 SHA512:
-  metadata.gz: c07aa30bcc9b040f75876124fec474bf48f13222ccbb053b7f0dba2313dbd459965d37fcd2b305d399dd503495cfadc2ca4f2b5c8e48f604528f6555d46e404e
-  data.tar.gz: 9d33cca29d90ed9dda3d7fd93c988b4b48999db4394be62044350648f11083ed03138201f76bb16569fbad6c61c2ff8b917af47d1f4a6e687c527bb9cd32b4b0
+  metadata.gz: d9dd7273eeccb284d7ee720fe586dc1e86b6e3355da378a719310cbe3209c508a57a3ee1cfcaa60762cd2860f76960b4f89aa39d7f5eafa90eb97b78adab6123
+  data.tar.gz: 26013bb1beb60367d878b451163594c50e36ef046b8ae0864a4728c4d0ae862e79ec22ab5ccb6c0acd88991b9f227b3af4b31daedc76e90755e9d87bce45f25c

data/README.md

@@ -6,7 +6,7 @@
 [![Documentation](https://img.shields.io/badge/docs-YARD-blue.svg)](https://cultivatelabs.github.io/raif/)
 
 
-Raif (Ruby AI Framework) is a Rails engine that helps you add AI-powered features to your Rails apps, such as [tasks](#tasks), [conversations](#conversations), and [agents](#agents).  It supports for multiple LLM providers including [OpenAI](#openai), [Anthropic Claude](#anthropic-claude), and [AWS Bedrock](#aws-bedrock).
+Raif (Ruby AI Framework) is a Rails engine that helps you add AI-powered features to your Rails apps, such as [tasks](#tasks), [conversations](#conversations), and [agents](#agents).  It supports for multiple LLM providers including [OpenAI](#openai), [Anthropic Claude](#anthropic-claude), [AWS Bedrock](#aws-bedrock), and [OpenRouter](#openrouter).
 
 Raif is built by [Cultivate Labs](https://www.cultivatelabs.com) and is used to power [ARC](https://www.arcanalysis.ai), an AI-powered research & analysis platform.
 
@@ -15,6 +15,7 @@ Raif is built by [Cultivate Labs](https://www.cultivatelabs.com) and is used to
   - [OpenAI](#openai)
   - [Anthropic Claude](#anthropic-claude)
   - [AWS Bedrock (Claude)](#aws-bedrock-claude)
+  - [OpenRouter](#openrouter)
 - [Chatting with the LLM](#chatting-with-the-llm)
 - [Key Raif Concepts](#key-raif-concepts)
   - [Tasks](#tasks)
@@ -22,12 +23,16 @@ Raif is built by [Cultivate Labs](https://www.cultivatelabs.com) and is used to
     - [Conversation Types](#conversation-types)
   - [Agents](#agents)
   - [Model Tools](#model-tools)
+- [Images/Files/PDF's](#imagesfilespdfs)
+  - [Images/Files/PDF's in Tasks](#imagesfilespdfs-in-tasks)
+- [Embedding Models](#embedding-models)
 - [Web Admin](#web-admin)
 - [Customization](#customization)
   - [Controllers](#controllers)
   - [Models](#models)
   - [Views](#views)
   - [System Prompts](#system-prompts)
+  - [Adding LLM Models](#adding-llm-models)
 - [Testing](#testing)
 - [Demo App](#demo-app)
 - [License](#license)
@@ -122,6 +127,28 @@ Currently supported Bedrock models:
 
 Note: Raif utilizes the [AWS Bedrock gem](https://docs.aws.amazon.com/sdk-for-ruby/v3/api/Aws/BedrockRuntime/Client.html) and AWS credentials should be configured via the AWS SDK (environment variables, IAM role, etc.)
 
+## OpenRouter
+[OpenRouter](https://openrouter.ai/) is a unified API that provides access to multiple AI models from different providers including Anthropic, Meta, Google, and more.
+
+```ruby
+Raif.configure do |config|
+  config.open_router_models_enabled = true
+  config.open_router_api_key = ENV["OPENROUTER_API_KEY"]
+  config.open_router_app_name = "Your App Name" # Optional
+  config.open_router_site_url = "https://yourdomain.com" # Optional
+  config.default_llm_model_key = "open_router_claude_3_7_sonnet"
+end
+```
+
+Currently included OpenRouter models:
+- `open_router_claude_3_7_sonnet`
+- `open_router_llama_3_3_70b_instruct`
+- `open_router_llama_3_1_8b_instruct`
+- `open_router_gemini_2_0_flash`
+- `open_router_deepseek_chat_v3`
+
+See [Adding LLM Models](#adding-llm-models) for more information on adding new OpenRouter models.
+
 # Chatting with the LLM
 
 When using Raif, it's often useful to use one of the [higher level abstractions](#key-raif-concepts) in your application. But when needed, you can utilize `Raif::Llm` to chat with the model directly. All calls to the LLM will create and return a `Raif::ModelCompletion` record, providing you a log of all interactions with the LLM which can be viewed in the [web admin](#web-admin).
@@ -160,7 +187,7 @@ puts model_completion.parsed_response # will strip backticks, parse the JSON, an
 # Key Raif Concepts
 
 ## Tasks
-If you have a single-shot task that you want an LLM to do in your application, you should create a `Raif::Task` subclass (see the end of this section for an example of using the task generator), where you'll define the prompt and response format for the task and call via `Raif::Task.run`. For example, say you have a `Document` model in your app and want to have a summarization task for the LLM:
+If you have a single-shot task that you want an LLM to do in your application, you should create a `Raif::Task` subclass, where you'll define the prompt and response format for the task and call via `Raif::Task.run`. For example, say you have a `Document` model in your app and want to have a summarization task for the LLM:
 
 ```bash
 rails generate raif:task DocumentSummarization --response-format html
@@ -171,7 +198,10 @@ This will create a new task in `app/models/raif/tasks/document_summarization.rb`
 ```ruby
 class Raif::Tasks::DocumentSummarization < Raif::ApplicationTask
   llm_response_format :html # options are :html, :text, :json
-  
+  llm_temperature 0.8 # optional, defaults to 0.7
+  llm_response_allowed_tags %w[p b i div strong] # optional, defaults to Rails::HTML5::SafeListSanitizer.allowed_tags
+  llm_response_allowed_attributes %w[style] # optional, defaults to Rails::HTML5::SafeListSanitizer.allowed_attributes
+
   # Any attr_accessor you define can be included as an argument when calling `run`. 
   # E.g. Raif::Tasks::DocumentSummarization.run(document: document, creator: user)
   attr_accessor :document
@@ -229,20 +259,10 @@ module Raif
 
       attr_accessor :topic
 
-      def self.json_response_schema
-        {
-          type: "object",
-          additionalProperties: false,
-          required: ["queries"],
-          properties: {
-            queries: {
-              type: "array",
-              items: {
-                type: "string"
-              }
-            }
-          }
-        }
+      json_response_schema do
+        array :queries do
+          items type: "string"
+        end
       end
 
       def build_prompt
@@ -271,6 +291,8 @@ You are an assistant with expertise in summarizing detailed articles into clear
 You're collaborating with teammate who speaks Spanish. Please respond in Spanish.
 ```
 
+The current list of valid language keys can be found [here](https://github.com/CultivateLabs/raif/blob/main/lib/raif/languages.rb).
+
 ## Conversations
 
 Raif provides `Raif::Conversation` and `Raif::ConversationEntry` models that you can use to  provide an LLM-powered chat interface. It also provides controllers and views for the conversation interface.
@@ -434,39 +456,31 @@ This will create a new model tool in `app/models/raif/model_tools/google_search.
 ```ruby
 class Raif::ModelTools::GoogleSearch < Raif::ModelTool
   # For example tool implementations, see: 
-  # Wikipedia Search Tool: https://github.com/CultivateLabs/raif/blob/main/app/models/raif/model_tools/wikipedia_search_tool.rb
-  # Fetch URL Tool: https://github.com/CultivateLabs/raif/blob/main/app/models/raif/model_tools/fetch_url_tool.rb
+  # Wikipedia Search Tool: https://github.com/CultivateLabs/raif/blob/main/app/models/raif/model_tools/wikipedia_search.rb
+  # Fetch URL Tool: https://github.com/CultivateLabs/raif/blob/main/app/models/raif/model_tools/fetch_url.rb
+
+  # Define the schema for the arguments that the LLM should use when invoking your tool.
+  # It should be a valid JSON schema. When the model invokes your tool,
+  # the arguments it provides will be validated against this schema using JSON::Validator from the json-schema gem.
+  #
+  # All attributes will be required and additionalProperties will be set to false.
+  #
+  # This schema would expect the model to invoke your tool with an arguments JSON object like:
+  # { "query" : "some query here" }
+  tool_arguments_schema do
+    string :query, description: "The query to search for"
+  end
 
   # An example of how the LLM should invoke your tool. This should return a hash with name and arguments keys.
   # `to_json` will be called on it and provided to the LLM as an example of how to invoke your tool.
-  def self.example_model_invocation
+  example_model_invocation do
     {
       "name": tool_name,
       "arguments": { "query": "example query here" }
     }
   end
 
-  # Define your tool's argument schema here. It should be a valid JSON schema.
-  # When the model invokes your tool, the arguments it provides will be validated
-  # against this schema using JSON::Validator from the json-schema gem.
-  def self.tool_arguments_schema
-    # For example:
-    # {
-    #   type: "object",
-    #   additionalProperties: false,
-    #   required: ["query"],
-    #   properties: {
-    #     query: {
-    #       type: "string",
-    #       description: "The query to search for"
-    #     }
-    #   }
-    # }
-    # Would expect the model to invoke your tool with an arguments JSON object like:
-    # { "query" : "some query here" }
-  end
-
-  def self.tool_description
+  tool_description do
     "Description of your tool that will be provided to the LLM so it knows when to invoke it"
   end
 
@@ -506,6 +520,125 @@ class Raif::ModelTools::GoogleSearch < Raif::ModelTool
 end
 ```
 
+## Images/Files/PDF's
+
+Raif supports images, files, and PDF's in the messages sent to the LLM.
+
+To include an image, file/PDF in a message, you can use the `Raif::ModelImageInput` and `Raif::ModelFileInput`.
+
+To include an image:
+```ruby
+# From a local file
+image = Raif::ModelImageInput.new(input: "path/to/image.png")
+
+# From a URL
+image = Raif::ModelImageInput.new(url: "https://example.com/image.png")
+
+# From an ActiveStorage attachment (assumes you have a User model with an avatar attachment)
+image = Raif::ModelImageInput.new(input: user.avatar)
+
+# Then chat with the LLM
+llm = Raif.llm(:open_ai_gpt_4o)
+model_completion = llm.chat(messages: [
+  { role: "user", content: ["What's in this image?", image]}
+])
+```
+
+To include a file/PDF:
+```ruby
+# From a local file
+file = Raif::ModelFileInput.new(input: "path/to/file.pdf")
+
+# From a URL
+file = Raif::ModelFileInput.new(url: "https://example.com/file.pdf")
+
+# From an ActiveStorage attachment (assumes you have a Document model with a pdf attachment)
+file = Raif::ModelFileInput.new(input: document.pdf)
+
+# Then chat with the LLM
+llm = Raif.llm(:open_ai_gpt_4o)
+model_completion = llm.chat(messages: [
+  { role: "user", content: ["What's in this file?", file]}
+])
+```
+
+### Images/Files/PDF's in Tasks
+
+You can include images and files/PDF's when running a `Raif::Task`:
+
+To include a file/PDF:
+```ruby
+file = Raif::ModelFileInput.new(input: "path/to/file.pdf")
+
+# Assumes you've created a PdfContentExtraction task
+task = Raif::Tasks::PdfContentExtraction.run(
+  creator: current_user,
+  files: [file]
+)
+```
+
+To include an image:
+```ruby
+image = Raif::ModelImageInput.new(input: "path/to/image.png")
+
+# Assumes you've created a ImageDescriptionGeneration task
+task = Raif::Tasks::ImageDescriptionGeneration.run(
+  creator: current_user,
+  images: [image]
+)
+```
+
+
+# Embedding Models
+
+Raif supports generation of vector embeddings. You can enable and configure embedding models in your Raif configuration:
+
+```ruby
+Raif.configure do |config|
+  config.open_ai_embedding_models_enabled = true
+  config.aws_bedrock_titan_embedding_models_enabled = true
+  
+  config.default_embedding_model_key = "open_ai_text_embedding_3_small"
+end
+```
+
+## Supported Embedding Models
+
+Raif currently supports the following embedding models:
+
+### OpenAI
+- `open_ai_text_embedding_3_small`
+- `open_ai_text_embed ding_3_large`
+- `open_ai_text_embedding_ada_002`
+
+### AWS Bedrock
+- `bedrock_titan_embed_text_v2`
+
+## Creating Embeddings
+
+By default, Raif will used `Raif.config.default_embedding_model_key` to create embeddings. To create an embedding for a piece of text:
+
+```ruby
+# Generate an embedding for a piece of text
+embedding = Raif.generate_embedding!("Your text here")
+
+# Generate an embedding for a piece of text with a specific number of dimensions
+embedding = Raif.generate_embedding!("Your text here", dimensions: 1024)
+
+# If you're using an OpenAI embedding model, you can pass an array of strings to embed multiple texts at once
+embeddings = Raif.generate_embedding!([
+  "Your text here",
+  "Your other text here"
+])
+```
+
+Or to generate embeddings for a piece of text with a specific model:
+
+```ruby
+model = Raif.embedding_model(:open_ai_text_embedding_3_small)
+embedding = model.generate_embedding!("Your text here")
+```
+
 # Web Admin
 
 Raif includes a web admin interface for viewing all interactions with the LLM. Assuming you have the engine mounted at `/raif`, you can access the admin interface at `/raif/admin`.
@@ -593,6 +726,32 @@ If you don't want to override the system prompt entirely in your task/conversati
 Raif.configure do |config|
   config.conversation_system_prompt_intro = "You are a helpful assistant who specializes in customer support."
   config.task_system_prompt_intro = "You are a helpful assistant who specializes in data analysis."
+  # or with a lambda
+  config.task_system_prompt_intro = ->(task) { "You are a helpful assistant who specializes in #{task.name}." }
+  config.conversation_system_prompt_intro = ->(conversation) { "You are a helpful assistant talking to #{conversation.creator.email}. Today's date is #{Date.today.strftime('%B %d, %Y')}." }
+end
+```
+
+## Adding LLM Models
+
+You can easily add new LLM models to Raif:
+
+```ruby
+# Register the model in Raif's LLM registry
+Raif.register_llm(Raif::Llms::OpenRouter, {
+  key: :open_router_gemini_flash_1_5_8b, # a unique key for the model
+  api_name: "google/gemini-flash-1.5-8b", # name of the model to be used in API calls - needs to match the provider's API name
+  input_token_cost: 0.038 / 1_000_000, # the cost per input token
+  output_token_cost: 0.15 / 1_000_000, # the cost per output token
+})
+
+# Then use the model
+llm = Raif.llm(:open_router_gemini_flash_1_5_8b)
+llm.chat(message: "Hello, world!")
+
+# Or set it as the default LLM model in your initializer
+Raif.configure do |config|
+  config.default_llm_model_key = "open_router_gemini_flash_1_5_8b"
 end
 ```
 

data/app/assets/stylesheets/raif/admin/stats.scss

@@ -0,0 +1,12 @@
+.stats-icon {
+  min-width: 46px;
+  text-align: center;
+}
+
+.stats-card {
+  transition: transform 0.2s;
+}
+
+.stats-card:hover {
+  transform: translateY(-5px);
+}

data/app/controllers/raif/admin/application_controller.rb

@@ -15,6 +15,20 @@ module Raif
         end
       end
 
+      def get_time_range(period)
+        case period
+        when "day"
+          24.hours.ago..Time.current
+        when "week"
+          1.week.ago..Time.current
+        when "month"
+          1.month.ago..Time.current
+        when "all"
+          Time.at(0)..Time.current
+        else
+          24.hours.ago..Time.current
+        end
+      end
     end
   end
 end

data/app/controllers/raif/admin/stats/tasks_controller.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Raif
+  module Admin
+    module Stats
+      class TasksController < Raif::Admin::ApplicationController
+        def index
+          @selected_period = params[:period] || "day"
+          @time_range = get_time_range(@selected_period)
+
+          @task_count = Raif::Task.where(created_at: @time_range).count
+
+          # Get task counts by type
+          @task_counts_by_type = Raif::Task.where(created_at: @time_range).group(:type).count
+
+          # Get costs by task type
+          @task_costs_by_type = Raif::Task.joins(:raif_model_completion)
+            .where(created_at: @time_range)
+            .group(:type)
+            .sum("raif_model_completions.total_cost")
+        end
+      end
+    end
+  end
+end

data/app/controllers/raif/admin/stats_controller.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Raif
+  module Admin
+    class StatsController < Raif::Admin::ApplicationController
+      def index
+        @selected_period = params[:period] || "day"
+        @time_range = get_time_range(@selected_period)
+
+        @model_completion_count = Raif::ModelCompletion.where(created_at: @time_range).count
+        @model_completion_total_cost = Raif::ModelCompletion.where(created_at: @time_range).sum(:total_cost)
+        @task_count = Raif::Task.where(created_at: @time_range).count
+        @conversation_count = Raif::Conversation.where(created_at: @time_range).count
+        @conversation_entry_count = Raif::ConversationEntry.where(created_at: @time_range).count
+        @agent_count = Raif::Agent.where(created_at: @time_range).count
+      end
+    end
+  end
+end

data/app/controllers/raif/admin/tasks_controller.rb

@@ -7,10 +7,26 @@ module Raif
 
       def index
         @task_types = Raif::Task.distinct.pluck(:type)
-        @selected_types = params[:task_types] || []
+        @selected_type = params[:task_types].present? ? params[:task_types] : "all"
+
+        @task_statuses = [:all, :completed, :failed, :in_progress, :pending]
+        @selected_statuses = params[:task_statuses].present? ? params[:task_statuses].to_sym : :all
 
         tasks = Raif::Task.order(created_at: :desc)
-        tasks = tasks.where(type: @selected_types) if @selected_types.present?
+        tasks = tasks.where(type: @selected_type) if @selected_type.present? && @selected_type != "all"
+
+        if @selected_statuses.present? && @selected_statuses != :all
+          case @selected_statuses
+          when :completed
+            tasks = tasks.completed
+          when :failed
+            tasks = tasks.failed
+          when :in_progress
+            tasks = tasks.in_progress
+          when :pending
+            tasks = tasks.pending
+          end
+        end
 
         @pagy, @tasks = pagy(tasks)
       end

data/app/controllers/raif/conversations_controller.rb

@@ -34,7 +34,11 @@ private
   end
 
   def validate_conversation_type
-    head :bad_request unless Raif.config.conversation_types.include?(conversation_type_param)
+    unless Raif.config.conversation_types.include?(conversation_type_param)
+      logger.error("Invalid Raif conversation type - not in Raif.config.conversation_types: #{conversation_type_param}")
+      logger.debug("\n\n\e[33m!!! Make sure to add the conversation type in Raif.config.conversation_types\e[0m\n")
+      head :bad_request
+    end
   end
 
   def raif_conversation_type

data/app/models/raif/agent.rb

@@ -22,14 +22,12 @@ module Raif
     validates :task, presence: true
     validates :system_prompt, presence: true
     validates :max_iterations, presence: true, numericality: { greater_than: 0 }
-    validates :available_model_tools, length: {
-      minimum: 1,
-      message: ->(_object, _data) {
-        I18n.t("raif.agents.errors.available_model_tools.too_short")
-      }
-    }
 
-    before_validation -> { self.system_prompt ||= build_system_prompt }, on: :create
+    before_validation -> {
+      populate_default_model_tools
+      self.system_prompt ||= build_system_prompt
+    },
+      on: :create
 
     attr_accessor :on_conversation_history_entry
 
@@ -59,7 +57,6 @@ module Raif
     def run!(&block)
       self.on_conversation_history_entry = block_given? ? block : nil
       self.started_at = Time.current
-      self.available_model_tools += ["Raif::ModelTools::AgentFinalAnswer"] unless available_model_tools.include?("Raif::ModelTools::AgentFinalAnswer")
       save!
 
       logger.debug <<~DEBUG
@@ -111,8 +108,12 @@ module Raif
 
   private
 
+    def populate_default_model_tools
+      # no-op by default. Can be overridden by subclasses to add default model tools
+    end
+
     def process_iteration_model_completion(model_completion)
-      raise NotImplementedError, "#{self.class.name} must implement execute_agent_iteration"
+      raise NotImplementedError, "#{self.class.name} must implement process_iteration_model_completion"
     end
 
     def native_model_tools
@@ -122,6 +123,7 @@ module Raif
     def add_conversation_history_entry(entry)
       entry_stringified = entry.stringify_keys
       conversation_history << entry_stringified
+      save!
       on_conversation_history_entry.call(entry_stringified) if on_conversation_history_entry.present?
     end
 

data/app/models/raif/agents/native_tool_calling_agent.rb

@@ -4,6 +4,16 @@ module Raif
   module Agents
     class NativeToolCallingAgent < Raif::Agent
       validate :ensure_llm_supports_native_tool_use
+      validates :available_model_tools, length: {
+        minimum: 2,
+        message: ->(_object, _data) {
+          I18n.t("raif.agents.native_tool_calling_agent.errors.available_model_tools.too_short")
+        }
+      }
+
+      before_validation -> {
+        available_model_tools << "Raif::ModelTools::AgentFinalAnswer" unless available_model_tools.include?("Raif::ModelTools::AgentFinalAnswer")
+      }
 
       def build_system_prompt
         <<~PROMPT.strip
@@ -49,7 +59,7 @@ module Raif
         if model_completion.response_tool_calls.blank?
           add_conversation_history_entry({
             role: "assistant",
-            content: "<observation>Error: No tool call found. I need make a tool call at each step. Available tools: #{available_model_tools_map.keys.join(", ")}</observation>" # rubocop:disable Layout/LineLength
+            content: "<observation>Error: No tool call found. I need to make a tool call at each step. Available tools: #{available_model_tools_map.keys.join(", ")}</observation>" # rubocop:disable Layout/LineLength
           })
           return
         end

data/app/models/raif/agents/re_act_agent.rb

@@ -3,6 +3,12 @@
 module Raif
   module Agents
     class ReActAgent < Raif::Agent
+      validates :available_model_tools, length: {
+        minimum: 1,
+        message: ->(_object, _data) {
+          I18n.t("raif.agents.re_act_agent.errors.available_model_tools.too_short")
+        }
+      }
 
       def build_system_prompt
         <<~PROMPT.strip

data/app/models/raif/concerns/has_available_model_tools.rb

@@ -4,7 +4,7 @@ module Raif::Concerns::HasAvailableModelTools
   extend ActiveSupport::Concern
 
   def available_model_tools_map
-    @available_model_tools_map ||= available_model_tools&.map do |tool_name|
+    available_model_tools&.map do |tool_name|
       tool_klass = tool_name.is_a?(String) ? tool_name.constantize : tool_name
       [tool_klass.tool_name, tool_klass]
     end.to_h

data/app/models/raif/concerns/json_schema_definition.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Raif
+  module Concerns
+    module JsonSchemaDefinition
+      extend ActiveSupport::Concern
+
+      class_methods do
+        def json_schema_definition(schema_name, &block)
+          raise ArgumentError, "A block must be provided to define the JSON schema" unless block_given?
+
+          @schemas ||= {}
+          @schemas[schema_name] = Raif::JsonSchemaBuilder.new
+          @schemas[schema_name].instance_eval(&block)
+          @schemas[schema_name]
+        end
+
+        def schema_defined?(schema_name)
+          @schemas&.dig(schema_name).present?
+        end
+
+        def schema_for(schema_name)
+          @schemas[schema_name].to_schema
+        end
+      end
+    end
+  end
+end

data/app/models/raif/concerns/llm_response_parsing.rb

@@ -9,6 +9,25 @@ module Raif::Concerns::LlmResponseParsing
     enum :response_format, Raif::Llm.valid_response_formats, prefix: true
 
     validates :response_format, presence: true, inclusion: { in: response_formats.keys }
+
+    class_attribute :allowed_tags
+    class_attribute :allowed_attributes
+  end
+
+  class_methods do
+    def llm_response_format(format)
+      raise ArgumentError, "response_format must be one of: #{response_formats.keys.join(", ")}" unless response_formats.keys.include?(format.to_s)
+
+      after_initialize -> { self.response_format = format }, if: :new_record?
+    end
+
+    def llm_response_allowed_tags(tags)
+      self.allowed_tags = tags
+    end
+
+    def llm_response_allowed_attributes(attributes)
+      self.allowed_attributes = attributes
+    end
   end
 
   # Parses the response from the LLM into a structured format, based on the response_format.
@@ -39,6 +58,9 @@ module Raif::Concerns::LlmResponseParsing
       end
     end
 
-    ActionController::Base.helpers.sanitize(fragment.to_html).strip
+    allowed_tags = self.class.allowed_tags || Rails::HTML5::SafeListSanitizer.allowed_tags
+    allowed_attributes = self.class.allowed_attributes || Rails::HTML5::SafeListSanitizer.allowed_attributes
+
+    ActionController::Base.helpers.sanitize(fragment.to_html, tags: allowed_tags, attributes: allowed_attributes).strip
   end
 end

data/app/models/raif/concerns/llm_temperature.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::LlmTemperature
+  extend ActiveSupport::Concern
+
+  included do
+    class_attribute :temperature, instance_writer: false
+  end
+
+  class_methods do
+    def llm_temperature(temperature)
+      raise ArgumentError, "temperature must be a number between 0 and 1" unless temperature.is_a?(Numeric) && temperature.between?(0, 1)
+
+      self.temperature = temperature
+    end
+  end
+end