bot.rb

require "bundler/setup"
require "dotenv/load"
require "telegram/bot"
require "ruby_llm"
require "net/http"
require "json"

TELEGRAM_TOKEN = ENV.fetch("TELEGRAM_BOT_TOKEN")
BASE_URL = "https://api.telegram.org/bot#{TELEGRAM_TOKEN}"

RubyLLM.configure do |config|
  config.openai_api_key = ENV["OPENAI_API_KEY"] if ENV["OPENAI_API_KEY"]
  config.anthropic_api_key = ENV["ANTHROPIC_API_KEY"] if ENV["ANTHROPIC_API_KEY"]

  # Vertex AI — uses Application Default Credentials (ADC) when no service account key is set
  config.vertexai_project_id = ENV["VERTEXAI_PROJECT_ID"] if ENV["VERTEXAI_PROJECT_ID"]
  config.vertexai_location = ENV["VERTEXAI_LOCATION"] if ENV["VERTEXAI_LOCATION"]
end

MODEL = ENV.fetch("LLM_MODEL", "gemini-2.0-flash-001")
PROVIDER = ENV["LLM_PROVIDER"]&.to_sym || :vertexai
DRAFT_THROTTLE_MS = 500

CHATS = {} # chat_id => RubyLLM::Chat

def telegram_api(method_name, params = {})
  uri = URI("#{BASE_URL}/#{method_name}")
  http = Net::HTTP.new(uri.host, uri.port)
  http.use_ssl = true

  request = Net::HTTP::Post.new(uri.path)
  request["Content-Type"] = "application/json"
  request.body = params.to_json

  response = http.request(request)
  JSON.parse(response.body)
end

def send_message_draft(chat_id:, draft_id:, text:)
  telegram_api("sendMessageDraft", { chat_id: chat_id, draft_id: draft_id, text: text })
end

def send_message(chat_id:, text:)
  telegram_api("sendMessage", { chat_id: chat_id, text: text })
end

MAX_MESSAGE_LENGTH = 4096

# Streams arbitrary text to a Telegram chat using sendMessageDraft for live updates.
# Yields an `emit` callable — call emit.call(chunk) to feed text chunks.
# Handles pagination automatically when content exceeds MAX_MESSAGE_LENGTH.
def stream_text(chat_id:)
  draft_id = rand(1..2_147_483_647)
  accumulated = ""
  page_offset = 0
  last_draft_at = 0

  threads = []

  emit = lambda do |chunk|
    accumulated += chunk

    if accumulated.length - page_offset >= MAX_MESSAGE_LENGTH
      send_message(chat_id: chat_id, text: accumulated[page_offset, MAX_MESSAGE_LENGTH])
      page_offset += MAX_MESSAGE_LENGTH
      draft_id = rand(1..2_147_483_647)
    end

    now = Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond)
    if now - last_draft_at >= DRAFT_THROTTLE_MS
      current_draft_id = draft_id
      text = accumulated[page_offset, MAX_MESSAGE_LENGTH - 2].dup
      threads << Thread.new { send_message_draft(chat_id: chat_id, draft_id: current_draft_id, text: text + " ▌") }
      last_draft_at = now
    end
  end

  yield emit

  remaining = accumulated[page_offset..]

  # Wait for all threads to complete before sending the final message
  threads.each(&:join)

  send_message(chat_id: chat_id, text: remaining) unless remaining.empty?
end

def stream_llm_response(chat_id:, user_message:)
  chat = CHATS[chat_id] ||= RubyLLM.chat(model: MODEL, provider: PROVIDER)
  final_message = nil

  stream_text(chat_id: chat_id) do |emit|
    final_message = chat.ask(user_message) do |chunk|
      emit.call(chunk.content) if chunk.content
    end
  end
rescue RubyLLM::Error => e
  send_message(chat_id: chat_id, text: "Error: #{e.message}")
rescue => e
  send_message(chat_id: chat_id, text: "Unexpected error: #{e.message}")
ensure
  puts "#{MODEL} | #{(final_message&.input_tokens || 0) + (final_message&.output_tokens || 0)} tokens"
end

puts "Bot starting with model: #{MODEL} (#{PROVIDER})"

Telegram::Bot::Client.run(TELEGRAM_TOKEN) do |bot|
  bot.listen do |message|
    next unless message.respond_to?(:text) && message.text

    case message.text.split(" ")[0]
    when "/start"
      bot.api.send_message(chat_id: message.chat.id, text: "Hello! Send me any message and I'll respond using AI.")
    when "/reset"
      CHATS.delete(message.chat.id)
      bot.api.send_message(chat_id: message.chat.id, text: "Conversation cleared.")
    when "/long"
      text = "lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum." * 20
      stream_text(chat_id: message.chat.id) do |emit|
        text.chars.each do |char|
          emit.call(char)
          sleep rand(0.001..0.002)
        end
      end
    when "/rick"
      stream_text(chat_id: message.chat.id) do |emit|
        ["Never gonna give you up", "Never gonna let you down", "Never gonna run around and desert you", "Never gonna make you cry", "Never gonna say goodbye", "Never gonna tell a lie and hurt you", "🕺"].each do |line|
          emit.call("\n" + line)
          sleep 1.5
        end
      end
    else
      stream_llm_response(chat_id: message.chat.id, user_message: message.text)
    end
  end
end