1
0
Fork 0
mirror of https://github.com/maybe-finance/maybe.git synced 2025-08-02 12:05:19 +02:00

improvements(ai): Improve AI streaming UI/UX interactions + better separation of AI provider responsibilities (#2039)

* Start refactor

* Interface updates

* Rework Assistant, Provider, and tests for better domain boundaries

* Consolidate and simplify OpenAI provider and provider concepts

* Clean up assistant streaming

* Improve assistant message orchestration logic

* Clean up "thinking" UI interactions

* Remove stale class

* Regenerate VCR test responses
This commit is contained in:
Zach Gollwitzer 2025-04-01 07:21:54 -04:00 committed by GitHub
parent 6331788b33
commit 5cf758bd03
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
33 changed files with 1179 additions and 624 deletions

View file

@ -1,5 +1,4 @@
require "test_helper"
require "ostruct"
class AssistantTest < ActiveSupport::TestCase
include ProviderTestHelper
@ -8,74 +7,109 @@ class AssistantTest < ActiveSupport::TestCase
@chat = chats(:two)
@message = @chat.messages.create!(
type: "UserMessage",
content: "Help me with my finances",
content: "What is my net worth?",
ai_model: "gpt-4o"
)
@assistant = Assistant.for_chat(@chat)
@provider = mock
@assistant.expects(:get_model_provider).with("gpt-4o").returns(@provider)
end
test "responds to basic prompt" do
text_chunk = OpenStruct.new(type: "output_text", data: "Hello from assistant")
response_chunk = OpenStruct.new(
type: "response",
data: OpenStruct.new(
id: "1",
model: "gpt-4o",
messages: [
OpenStruct.new(
id: "1",
content: "Hello from assistant",
)
],
functions: []
)
)
test "errors get added to chat" do
@assistant.expects(:get_model_provider).with("gpt-4o").returns(@provider)
@provider.expects(:chat_response).with do |message, **options|
options[:streamer].call(text_chunk)
options[:streamer].call(response_chunk)
true
end
error = StandardError.new("test error")
@provider.expects(:chat_response).returns(provider_error_response(error))
assert_difference "AssistantMessage.count", 1 do
@chat.expects(:add_error).with(error).once
assert_no_difference "AssistantMessage.count" do
@assistant.respond_to(@message)
end
end
test "responds with tool function calls" do
function_request_chunk = OpenStruct.new(type: "function_request", data: "get_net_worth")
text_chunk = OpenStruct.new(type: "output_text", data: "Your net worth is $124,200")
response_chunk = OpenStruct.new(
type: "response",
data: OpenStruct.new(
id: "1",
model: "gpt-4o",
messages: [
OpenStruct.new(
id: "1",
content: "Your net worth is $124,200",
)
],
functions: [
OpenStruct.new(
id: "1",
call_id: "1",
name: "get_net_worth",
arguments: "{}",
result: "$124,200"
)
]
)
test "responds to basic prompt" do
@assistant.expects(:get_model_provider).with("gpt-4o").returns(@provider)
text_chunks = [
provider_text_chunk("I do not "),
provider_text_chunk("have the information "),
provider_text_chunk("to answer that question")
]
response_chunk = provider_response_chunk(
id: "1",
model: "gpt-4o",
messages: [ provider_message(id: "1", text: text_chunks.join) ],
function_requests: []
)
response = provider_success_response(response_chunk.data)
@provider.expects(:chat_response).with do |message, **options|
options[:streamer].call(function_request_chunk)
options[:streamer].call(text_chunk)
text_chunks.each do |text_chunk|
options[:streamer].call(text_chunk)
end
options[:streamer].call(response_chunk)
true
end.returns(response)
assert_difference "AssistantMessage.count", 1 do
@assistant.respond_to(@message)
message = @chat.messages.ordered.where(type: "AssistantMessage").last
assert_equal "I do not have the information to answer that question", message.content
assert_equal 0, message.tool_calls.size
end
end
test "responds with tool function calls" do
@assistant.expects(:get_model_provider).with("gpt-4o").returns(@provider).once
# Only first provider call executes function
Assistant::Function::GetAccounts.any_instance.stubs(:call).returns("test value").once
# Call #1: Function requests
call1_response_chunk = provider_response_chunk(
id: "1",
model: "gpt-4o",
messages: [],
function_requests: [
provider_function_request(id: "1", call_id: "1", function_name: "get_accounts", function_args: "{}")
]
)
call1_response = provider_success_response(call1_response_chunk.data)
# Call #2: Text response (that uses function results)
call2_text_chunks = [
provider_text_chunk("Your net worth is "),
provider_text_chunk("$124,200")
]
call2_response_chunk = provider_response_chunk(
id: "2",
model: "gpt-4o",
messages: [ provider_message(id: "1", text: call2_text_chunks.join) ],
function_requests: []
)
call2_response = provider_success_response(call2_response_chunk.data)
sequence = sequence("provider_chat_response")
@provider.expects(:chat_response).with do |message, **options|
call2_text_chunks.each do |text_chunk|
options[:streamer].call(text_chunk)
end
options[:streamer].call(call2_response_chunk)
true
end.returns(call2_response).once.in_sequence(sequence)
@provider.expects(:chat_response).with do |message, **options|
options[:streamer].call(call1_response_chunk)
true
end.returns(call1_response).once.in_sequence(sequence)
assert_difference "AssistantMessage.count", 1 do
@assistant.respond_to(@message)
@ -83,4 +117,34 @@ class AssistantTest < ActiveSupport::TestCase
assert_equal 1, message.tool_calls.size
end
end
private
def provider_function_request(id:, call_id:, function_name:, function_args:)
Provider::LlmConcept::ChatFunctionRequest.new(
id: id,
call_id: call_id,
function_name: function_name,
function_args: function_args
)
end
def provider_message(id:, text:)
Provider::LlmConcept::ChatMessage.new(id: id, output_text: text)
end
def provider_text_chunk(text)
Provider::LlmConcept::ChatStreamChunk.new(type: "output_text", data: text)
end
def provider_response_chunk(id:, model:, messages:, function_requests:)
Provider::LlmConcept::ChatStreamChunk.new(
type: "response",
data: Provider::LlmConcept::ChatResponse.new(
id: id,
model: model,
messages: messages,
function_requests: function_requests
)
)
end
end

View file

@ -6,16 +6,11 @@ class Provider::OpenaiTest < ActiveSupport::TestCase
setup do
@subject = @openai = Provider::Openai.new(ENV.fetch("OPENAI_ACCESS_TOKEN", "test-openai-token"))
@subject_model = "gpt-4o"
@chat = chats(:two)
end
test "openai errors are automatically raised" do
VCR.use_cassette("openai/chat/error") do
response = @openai.chat_response(UserMessage.new(
chat: @chat,
content: "Error test",
ai_model: "invalid-model-that-will-trigger-api-error"
))
response = @openai.chat_response("Test", model: "invalid-model-that-will-trigger-api-error")
assert_not response.success?
assert_kind_of Provider::Openai::Error, response.error
@ -24,113 +19,145 @@ class Provider::OpenaiTest < ActiveSupport::TestCase
test "basic chat response" do
VCR.use_cassette("openai/chat/basic_response") do
message = @chat.messages.create!(
type: "UserMessage",
content: "This is a chat test. If it's working, respond with a single word: Yes",
ai_model: @subject_model
response = @subject.chat_response(
"This is a chat test. If it's working, respond with a single word: Yes",
model: @subject_model
)
response = @subject.chat_response(message)
assert response.success?
assert_equal 1, response.data.messages.size
assert_includes response.data.messages.first.content, "Yes"
assert_includes response.data.messages.first.output_text, "Yes"
end
end
test "streams basic chat response" do
VCR.use_cassette("openai/chat/basic_response") do
VCR.use_cassette("openai/chat/basic_streaming_response") do
collected_chunks = []
mock_streamer = proc do |chunk|
collected_chunks << chunk
end
message = @chat.messages.create!(
type: "UserMessage",
content: "This is a chat test. If it's working, respond with a single word: Yes",
ai_model: @subject_model
response = @subject.chat_response(
"This is a chat test. If it's working, respond with a single word: Yes",
model: @subject_model,
streamer: mock_streamer
)
@subject.chat_response(message, streamer: mock_streamer)
tool_call_chunks = collected_chunks.select { |chunk| chunk.type == "function_request" }
text_chunks = collected_chunks.select { |chunk| chunk.type == "output_text" }
response_chunks = collected_chunks.select { |chunk| chunk.type == "response" }
assert_equal 1, text_chunks.size
assert_equal 1, response_chunks.size
assert_equal 0, tool_call_chunks.size
assert_equal "Yes", text_chunks.first.data
assert_equal "Yes", response_chunks.first.data.messages.first.content
assert_equal "Yes", response_chunks.first.data.messages.first.output_text
assert_equal response_chunks.first.data, response.data
end
end
test "chat response with tool calls" do
VCR.use_cassette("openai/chat/tool_calls") do
response = @subject.chat_response(
tool_call_message,
test "chat response with function calls" do
VCR.use_cassette("openai/chat/function_calls") do
prompt = "What is my net worth?"
functions = [
{
name: "get_net_worth",
description: "Gets a user's net worth",
params_schema: { type: "object", properties: {}, required: [], additionalProperties: false },
strict: true
}
]
first_response = @subject.chat_response(
prompt,
model: @subject_model,
instructions: "Use the tools available to you to answer the user's question.",
available_functions: [ PredictableToolFunction.new(@chat) ]
functions: functions
)
assert response.success?
assert_equal 1, response.data.functions.size
assert_equal 1, response.data.messages.size
assert_includes response.data.messages.first.content, PredictableToolFunction.expected_test_result
assert first_response.success?
function_request = first_response.data.function_requests.first
assert function_request.present?
second_response = @subject.chat_response(
prompt,
model: @subject_model,
function_results: [ {
call_id: function_request.call_id,
output: { amount: 10000, currency: "USD" }.to_json
} ],
previous_response_id: first_response.data.id
)
assert second_response.success?
assert_equal 1, second_response.data.messages.size
assert_includes second_response.data.messages.first.output_text, "$10,000"
end
end
test "streams chat response with tool calls" do
VCR.use_cassette("openai/chat/tool_calls") do
test "streams chat response with function calls" do
VCR.use_cassette("openai/chat/streaming_function_calls") do
collected_chunks = []
mock_streamer = proc do |chunk|
collected_chunks << chunk
end
prompt = "What is my net worth?"
functions = [
{
name: "get_net_worth",
description: "Gets a user's net worth",
params_schema: { type: "object", properties: {}, required: [], additionalProperties: false },
strict: true
}
]
# Call #1: First streaming call, will return a function request
@subject.chat_response(
tool_call_message,
prompt,
model: @subject_model,
instructions: "Use the tools available to you to answer the user's question.",
available_functions: [ PredictableToolFunction.new(@chat) ],
functions: functions,
streamer: mock_streamer
)
text_chunks = collected_chunks.select { |chunk| chunk.type == "output_text" }
text_chunks = collected_chunks.select { |chunk| chunk.type == "output_text" }
tool_call_chunks = collected_chunks.select { |chunk| chunk.type == "function_request" }
response_chunks = collected_chunks.select { |chunk| chunk.type == "response" }
assert_equal 1, tool_call_chunks.count
assert text_chunks.count >= 1
assert_equal 1, response_chunks.count
assert_equal 0, text_chunks.size
assert_equal 1, response_chunks.size
assert_includes response_chunks.first.data.messages.first.content, PredictableToolFunction.expected_test_result
first_response = response_chunks.first.data
function_request = first_response.function_requests.first
# Reset collected chunks for the second call
collected_chunks = []
# Call #2: Second streaming call, will return a function result
@subject.chat_response(
prompt,
model: @subject_model,
function_results: [
{
call_id: function_request.call_id,
output: { amount: 10000, currency: "USD" }
}
],
previous_response_id: first_response.id,
streamer: mock_streamer
)
text_chunks = collected_chunks.select { |chunk| chunk.type == "output_text" }
response_chunks = collected_chunks.select { |chunk| chunk.type == "response" }
assert text_chunks.size >= 1
assert_equal 1, response_chunks.size
assert_includes response_chunks.first.data.messages.first.output_text, "$10,000"
end
end
private
def tool_call_message
UserMessage.new(chat: @chat, content: "What is my net worth?", ai_model: @subject_model)
end
class PredictableToolFunction < Assistant::Function
class << self
def expected_test_result
"$124,200"
end
def name
"get_net_worth"
end
def description
"Gets user net worth data"
end
end
def call(params = {})
self.class.expected_test_result
end
end
end