-
Notifications
You must be signed in to change notification settings - Fork 0
/
openai_client.rb
175 lines (156 loc) · 5.74 KB
/
openai_client.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
require "openai"
##
# The `OpenAIClient` class provides an interface for interacting with the OpenAI API.
#
# Use the `initialize` method to create an instance of the `OpenAIClient` class:
#
# ```ruby
# client = OpenAIClient.new(api_key, model)
# ```
#
# You'll need an API key from OpenAI to use this class. You can sign up for an API key on the OpenAI website.
#
# The `OpenAIClient` class provides the following methods for generating responses from the API:
#
# - {#generate_json_response} - Generates a JSON response from the API.
# - {#generate_code} - Generates Ruby code from a given prompt using the API.
# - {#generate_raw_response} - Generates a raw response from the API.
#
# Example usage:
#
# ```ruby
# client = OpenAIClient.new(api_key, model)
# response = client.generate_code("Sort an array of numbers.")
# ```
class OpenAIClient
##
# Initializes a new OpenAIClient object.
#
# @param [String] api_key Your OpenAI API key.
# @param [String] model The name of the model to use for generating responses.
def initialize(api_key, model)
OpenAI.configure do |config|
config.access_token = api_key
end
@client = OpenAI::Client.new
@model = model
end
##
# Generates a JSON response from the OpenAI API.
#
# @param [Array] messages The messages to send to the API.
# @param [Integer] max_tokens The maximum number of tokens to generate.
# @param [Integer] n The number of responses to generate.
# @param [String, nil] stop A string at which generation is stopped.
# @param [Float] temperature The sampling temperature to use.
# @param [Integer] max_attempts The maximum number of attempts to generate a response.
# @return [String] The generated JSON response.
def generate_json_response(messages, max_tokens = 1000, n = 1, stop = nil, temperature = 0.5, max_attempts = 5)
attempt = 1
while attempt <= max_attempts
response = @client.chat(parameters: {
model: @model,
messages: messages,
max_tokens: max_tokens,
n: n,
stop: stop,
temperature: temperature,
}
)
if response["choices"].nil?
puts response
raise "Error: The 'choices' key is missing or has a nil value in the response."
end
content = response["choices"][0]["message"]["content"]
if ENV['DEBUG'] == 'true'
puts "---- Full Response Message Content ----"
puts content
puts "----"
end
last_line = content.lines.last.chomp
if last_line == "###END###"
content = content.chomp("\n###END###")
return content
else
attempt += 1
puts "Error: The last line must be '###END###', aborted gpt message detected."
sleep(1)
end
end
raise "Unable to obtain a complete or properly-formatted message from chatgpt after #{max_attempts} attempts."
end
##
# Generates Ruby code from a given prompt using the OpenAI API.
#
# @param [String] prompt The prompt to generate code from.
# @param [Integer] max_tokens The maximum number of tokens to generate.
# @param [Integer] n The number of responses to generate.
# @param [String, nil] stop A string at which generation is stopped.
# @param [Float] temperature The sampling temperature to use.
# @param [Integer] max_attempts The maximum number of attempts to generate a response.
# @return [String] The generated Ruby code.
def generate_code(prompt, max_tokens = 150, n = 1, stop = nil, temperature = 0.5, max_attempts = 5)
attempt = 1
while attempt <= max_attempts
response = @client.chat(parameters: {
model: @model,
messages: [
{role: "system", content: "You are a helpful assistant that translates English descriptions into concise Ruby code."},
{role: "user", content: "#{prompt}"}
],
max_tokens: max_tokens,
n: n,
stop: stop,
temperature: temperature
})
content = response["choices"][0]["message"]["content"]
backticks_count = content.scan(/```/).count
if backticks_count % 2 == 0
code_blocks = content.scan(/```ruby\n(.*?)\n```/m).flatten
extracted_code = code_blocks.map(&:strip).join("\n\n")
if ENV['DEBUG'] == 'true'
puts "---- Full Response Message Content ----"
puts content
puts "---- Extracted Code Blocks ----"
puts extracted_code
end
return extracted_code
else
attempt += 1
puts "Unbalanced triple backticks detected. Retrying (attempt #{attempt})..." if ENV['DEBUG'] == 'true'
end
end
raise "Unable to generate code with balanced triple backticks after #{max_attempts} attempts."
end
##
# Generates a raw response from the OpenAI API.
#
# @param [Array] messages The messages to send to the API.
# @param [Integer] max_tokens The maximum number of tokens to generate.
# @param [Integer] n The number of responses to generate.
# @param [String, nil] stop A string at which generation is stopped.
# @param [Float] temperature The sampling temperature to use.
# @return [String] The raw generated response.
def generate_raw_response(messages, max_tokens = 1000, n = 1, stop = nil, temperature = 0.5)
response = @client.chat(parameters: {
model: @model,
messages: messages,
max_tokens: max_tokens,
n: n,
stop: stop,
temperature: temperature,
}
)
if response["choices"].nil?
puts response
raise "Error: The 'choices' key is missing or has a nil value in the response."
end
content = response["choices"][0]["message"]["content"]
if ENV['DEBUG'] == 'true'
puts "---- Full Response Message Content ----"
puts content
puts "----"
end
content
end
end