A little ollama client for pre-defined prompts.

llme.rb

#!/usr/bin/env ruby

require 'thor'
require 'uri'
require 'net/http'
require 'json'

# > llme joke
# Sure! Here's one:
# 
# Why don't scientists trust atoms?
# 
# Because they make up everything!%
class Llme < Thor
  OLLAMA_HOST = ENV.fetch('OLLAMA_HOST', 'localhost')
  OLLAMA_PORT = ENV.fetch('OLLAMA_PORT', '11434')

  desc 'joke', 'Tell me a joke'
  def joke
    run 'Tell me a joke'
  end

  desc 'architecture', 'llm prompt for generating an architecture.md documentation document for a ruby on rails app'
  def architecture
    context = <<~CONTEXT
      - This is a single ruby on rails application.
      - These are the class comments on all the models.

    CONTEXT
    context << `ag -B 10 '^class' app/models/*.rb`

    prompt = <<~PROMPT
      Please write a document that outlines the architecture of a Ruby on Rails
      application in with the most appropriate language and style that an
      expert techincal writer might use. Cover at least the following sections:
      Overview, Components. And format the document using markdown.
    PROMPT

    warn prompt
    warn "\n"
    run(prompt, model: 'llama2', context: context)
  end


  private

  def run(prompt, model: 'llama2', context: nil)
    uri = URI("http://#{OLLAMA_HOST}:#{OLLAMA_PORT}/v1/chat/completions")

    messages = []
    messages << { role: 'system', content: context } unless context.nil?
    messages << { role: 'user', content: prompt }

    request = Net::HTTP::Post.new(uri, 'Content-Type' => 'application/json')
    request.body = {
      messages: messages,
      model: model,
      temperature: 0.7,
      max_tokens: -1,
      stream: true
    }.to_json

    Net::HTTP.start(uri.hostname, uri.port) do |http|
      http.request(request) do |response|
        response.read_body do |chunk|
          chunk = chunk.split(/\n/).reject{|subhunk| subhunk == 'data: [DONE]'}.join("\n")
          message = JSON.parse(chunk[6..-1]).dig('choices', 0, 'delta', 'content')
          print message
          $stdout.flush
        end
      end
    end
  end
end

Llme.start(ARGV)