Class: Langchain::Tool::Database

Inherits:
Object
  • Object
show all
Extended by:
Langchain::ToolDefinition
Includes:
DependencyHelper
Defined in:
lib/langchain/tool/database.rb

Overview

Connects to a SQL database, executes SQL queries, and outputs DB schema for Agents to use

Gem requirements:

gem "sequel", "~> 5.87.0"

Usage:

database = Langchain::Tool::Database.new(connection_string: "postgres://user:password@localhost:5432/db_name")

Instance Attribute Summary collapse

Instance Method Summary collapse

Methods included from Langchain::ToolDefinition

define_function, extended, function_schemas, tool_name

Methods included from DependencyHelper

#depends_on

Constructor Details

#initialize(connection_string:, tables: [], exclude_tables: []) ⇒ Database

Establish a database connection

Parameters:

  • connection_string (String) —

    Database connection info, e.g. ‘postgres://user:password@localhost:5432/db_name’

  • tables (Array<Symbol>) (defaults to: []) —

    The tables to use. Will use all if empty.

  • except_tables (Array<Symbol>) —

    The tables to exclude. Will exclude none if empty.

Raises:

  • (StandardError) 

39
40
41
42
43
44
45
46
47
48
 
# File 'lib/langchain/tool/database.rb', line 39

def initialize(connection_string:, tables: [], exclude_tables: [])
  depends_on "sequel"

  raise StandardError, "connection_string parameter cannot be blank" if connection_string.empty?

  @db = Sequel.connect(connection_string)
  # TODO: This is a bug, these 2 parameters are completely ignored.
  @requested_tables = tables
  @excluded_tables = exclude_tables
end

Instance Attribute Details

#db ⇒ Object (readonly)

Returns the value of attribute db.


31
32
33
 
# File 'lib/langchain/tool/database.rb', line 31

def db
  @db
end

#excluded_tables ⇒ Object (readonly)

Returns the value of attribute excluded_tables.


31
32
33
 
# File 'lib/langchain/tool/database.rb', line 31

def excluded_tables
  @excluded_tables
end

#requested_tables ⇒ Object (readonly)

Returns the value of attribute requested_tables.


31
32
33
 
# File 'lib/langchain/tool/database.rb', line 31

def requested_tables
  @requested_tables
end

Instance Method Details

#describe_tables(tables: []) ⇒ Langchain::Tool::Response

Database Tool: Returns the schema for a list of tables

Parameters:

  • tables (Array<String>) (defaults to: []) —

    The tables to describe.

Returns:

  • (Langchain::Tool::Response) —

    The schema for the tables


61
62
63
64
65
66
67
68
69
70
71
72
73
 
# File 'lib/langchain/tool/database.rb', line 61

def describe_tables(tables: [])
  return "No tables specified" if tables.empty?

  Langchain.logger.debug("#{self.class} - Describing tables: #{tables}")

  result = tables
    .map do |table|
      describe_table(table)
    end
    .join("\n")

  tool_response(content: result)
end

#dump_schema ⇒ Langchain::Tool::Response

Database Tool: Returns the database schema

Returns:

  • (Langchain::Tool::Response) —

    Database schema


78
79
80
81
82
83
84
85
86
 
# File 'lib/langchain/tool/database.rb', line 78

def dump_schema
  Langchain.logger.debug("#{self.class} - Dumping schema tables and keys")

  schemas = db.tables.map do |table|
    describe_table(table)
  end

  tool_response(content: schemas.join("\n"))
end

#execute(input:) ⇒ Langchain::Tool::Response

Database Tool: Executes a SQL query and returns the results

Parameters:

  • input (String) —

    SQL query to be executed

Returns:

  • (Langchain::Tool::Response) —

    Results from the SQL query


92
93
94
95
96
97
98
99
 
# File 'lib/langchain/tool/database.rb', line 92

def execute(input:)
  Langchain.logger.debug("#{self.class} - Executing \"#{input}\"")

  tool_response(content: db[input].to_a)
rescue Sequel::DatabaseError => e
  Langchain.logger.error("#{self.class} - #{e.message}")
  tool_response(content: e.message)
end

#list_tables ⇒ Langchain::Tool::Response

Database Tool: Returns a list of tables in the database

Returns:

  • (Langchain::Tool::Response) —

    List of tables in the database


53
54
55
 
# File 'lib/langchain/tool/database.rb', line 53

def list_tables
  tool_response(content: db.tables)
end