Using GraphQL with LLMs in Ruby

GraphQL is an approach to building APIs designed to be used by a variety of clients. With GraphQL, a client typically uses a provided schema to either query or mutate the data needed. Once defined, these APIs are usually exposed to the web via a framework like Ruby on Rails. Interestingly, the defined schema of GraphQL APIs make them perfect for use with LLMs. This guide walks through the steps needed to build a basic GraphQL API and link it to an LLM using Ruby.

An Introduction to GraphQL

A GraphQL API can be expressed as a schema. A schema simply defines the queries and mutations that are possible within the API. Let's look at a basic example:

type Contact {
  name: String!
}

type SaveContactMutationPayload {
  contact: Contact!
}

type Mutation {
  saveContact(name: String!): SaveContactMutationPayload!
}

type Query {
  contacts: [Contact!]!
}

Queries

A query is typically used to retrieve data. For the above example schema, a query to fetch contacts looks like the following:

query Contacts {
  contacts {
    name
  }
}

The "Contacts" query returns back the following payload:

{
  "data": [
    { "name": "John" },
    { "name": "Paul" },
    { "name": "George" },
    { "name": "Ringo" }
  ]
}

Mutations

A mutation is typically used to modify data. For the above example schema, a mutation to save a contact looks like the following:

mutation SaveContact($name: String!) {
  saveContact(name: $name) {
    contact {
      name
    }
  }
}

The "SaveContact" mutation requires some variables:

{ "name": "Adele" }

When run in combination it also returns back a payload:

{
  "data": {
    "saveContact": {
      "contact": {
        "name": "Adele"
      }
    }
  }
}

Using GraphQL with Ruby

The GraphQL Ruby library makes building a GraphQL API in Ruby simple. The earlier schema can be defined and implemented in just a handful of lines:

require 'bundler/inline'

gemfile do
  source 'https://rubygems.org'
  gem 'graphql'
  gem 'json'
end

Contact = Data.define(:name)

CONTACTS = [
  Contact.new(name: "John"),
  Contact.new(name: "Paul"),
  Contact.new(name: "George"),
  Contact.new(name: "Ringo"),
]

class ContactType < GraphQL::Schema::Object
  field :name, String, null: false
end

class QueryType < GraphQL::Schema::Object
  field :contacts, [ContactType], null: false

  def contacts
    CONTACTS
  end
end

class SaveContactMutation < GraphQL::Schema::Mutation
  null false
  argument :name, String
  field :contact, ContactType

  def resolve(name:)
    contact = Contact.new(name: name)
    CONTACTS << contact
    { contact: }
  end
end

class MutationType < GraphQL::Schema::Object
  field :save_contact, mutation: SaveContactMutation
end

class Schema < GraphQL::Schema
  query QueryType
  mutation MutationType
end

The above Schema can then be supplied with some gql and variables to resolve a query or mutation:

def execute(gql:, variables: {})
  puts JSON.generate(Schema.execute(gql, variables:))
end

Then the schema can be tested against a handfull of queries and mutations:

execute(gql: <<~GQL)
  query Contacts {
    contacts {
      name
    }
  }
GQL

{
  "data": {
    "contacts": [
      { "name": "John" },
      { "name": "Paul" },
      { "name": "George" },
      { "name": "Ringo" }
    ]
  }
}

execute(gql: <<~GQL, variables: { name: "Ringo" })
  mutation SaveContact($name: String!) {
    saveContact(name: $name) {
      contact {
        name
      }
    }
  }
GQL

{
  "data": {
    "saveContact": {
      "contact": {
        "name": "Adele"
      }
    }
  }
}

execute(gql: <<~GQL)
  query Contacts {
    contacts {
      name
    }
  }
GQL

{
  "data": {
    "contacts": [
      { "name": "John" },
      { "name": "Paul" },
      { "name": "George" },
      { "name": "Ringo" },
      { "name": "Adele" }
    ]
  }
}

Using LLMs with Ruby

To integrate with an LLM, OmniAI can be used. It provides a consistent Ruby API for working with a variety of LLMs (Anthropic, Google, Mistral, OpenAI, etc). This example uses Anthropic, but any of the earlier providers can be swapped in. A basic input / output chat loop can be run with the following:

require 'bundler/inline'

gemfile do
  source 'https://rubygems.org'
  gem 'omniai'
  gem 'omniai-anthropic'
end

Message = Data.define(:text, :role)

def listen!
  client = OmniAI::Anthropic::Client.new
  messages = []

  loop do
    print('# ')
    text = gets&.chomp

    break if text.nil? || text.match?(/\A(exit|quit)\z/i)

    messages << Message.new(text: text, role: 'user')

    completion = client.chat do |chat|
      messages.each do |message|
        chat.message(message.text, role: message.role)
      end
    end

    puts(completion.text)

    messages << Message.new(text: completion.text, role: 'assistant')
  rescue Interrupt
    break
  end
end

listen!

This loop keeps track of history (via messages), but currently has no insights into our API. Let’s run a few tests to make sure everything is working as expected:

# Tell me a joke.
Why don't scientists trust atoms? Because they make up everything!

# Tell another with a similar structure.
Why can't you trust a math teacher? Because they have too many problems!

Linking the LLM with GraphQL via Tools

With these basic building blocks in place the glue between the two is an LLM tool. Tools offer an interface between an LLM and an application. OmniAI automatically handles the calling of tools as required. For our case, a GraphQL tool is built using the following:

GRAPHQL_TOOL = OmniAI::Tool.new(
  proc { |gql:, variables: nil|
    variables = JSON.parse(variables) if variables
    JSON.generate(Schema.execute(gql, variables:))
  },
  name: 'graphql_execute',
  description: <<~TEXT,
    Runs a GQL query or mutation on the following schema and return the result:

    <schema>
    #{Schema.to_definition}
    </schema>
  TEXT
  parameters: OmniAI::Tool::Parameters.new(
    properties: {
      gql: OmniAI::Tool::Property.string(description: 'A GQL query or mutation.'),
      variables: OmniAI::Tool::Property.string(description: 'Optional JSON variables.'),
    },
    required: %i[gql]
  )
)

This tool takes in two parameters: gql and variables. It documents the queries and mutations available in the schema (via Schema.to_definition) for the LLM in the description. Lastly, it provides a proc that runs the gql query or mutation on our schema. To allow the LLM to use the tool we simply provide it in the client.chat call.

completion = client.chat(tools: [GRAPHQL_TOOL]) do |chat|
  # ...
end

Summary

That's it! The combined code is as follows:

require 'bundler/inline'

gemfile do
  source 'https://rubygems.org'
  gem 'graphql'
  gem 'json'
  gem 'logger'
  gem 'omniai'
  gem 'omniai-anthropic'
end

Contact = Data.define(:name)
Message = Data.define(:text, :role)

CONTACTS = [
  Contact.new(name: "John"),
  Contact.new(name: "Paul"),
  Contact.new(name: "George")
]

class ContactType < GraphQL::Schema::Object
  field :name, String, null: false
end

class QueryType < GraphQL::Schema::Object
  field :contacts, [ContactType], null: false

  def contacts
    CONTACTS
  end
end

class SaveContactMutation < GraphQL::Schema::Mutation
  argument :name, String
  field :contact, ContactType

  def resolve(name:)
    contact = Contact.new(name: name)
    CONTACTS << contact
    { contact: }
  end
end

class MutationType < GraphQL::Schema::Object
  field :save_contact, mutation: SaveContactMutation
end

class Schema < GraphQL::Schema
  query QueryType
  mutation MutationType
end

GRAPHQL_TOOL = OmniAI::Tool.new(
  proc { |gql:, variables: nil|
    variables = JSON.parse(variables) if variables
    JSON.generate(Schema.execute(gql, variables:))
  },
  name: 'graphql_execute',
  description: <<~TEXT,
    Runs a GQL query or mutation on the following schema and return the result:

    <schema>
    #{Schema.to_definition}
    </schema>
  TEXT
  parameters: OmniAI::Tool::Parameters.new(
    properties: {
      gql: OmniAI::Tool::Property.string(description: 'A GQL query or mutation.'),
      variables: OmniAI::Tool::Property.string(description: 'Optional JSON variables.'),
    },
    required: %i[gql]
  )
)

def listen!
  client = OmniAI::Anthropic::Client.new
  messages = []

  loop do
    print('# ')
    text = gets&.chomp

    break if text.nil? || text.match?(/\A(exit|quit)\z/i)

    messages << Message.new(text: text, role: 'user')

    completion = client.chat(tools: [GRAPHQL_TOOL]) do |chat|
      messages.each do |message|
        chat.message(message.text, role: message.role)
      end
    end

    puts(completion.text)

    messages << Message.new(text: completion.text, role: 'assistant')
  rescue Interrupt
    break
  end
end

listen!

Let’s run a few tests using some natural language to interact with our API:

# Who are my contacts?
Your contacts are:
1. John
2. Paul
3. George

# Can you make sure that Adele and Aretha are contacts?
Both Adele and Aretha have been successfully added to the contacts list!

# Can you include George in my contacts if not already a contact?
I can see that George is already in your contacts list, so there's no need to add them again. Your current contacts are:
1. John
2. Paul
3. George
4. Adele
5. Aretha

If you are curious about other use cases for OmniAI please checkout the documentation for further examples of using tools!

This article originally appeared on https://workflow.ing/blog/articles/using-graphql-with-llms-in-ruby.