mirror of
https://github.com/discourse/discourse-ai.git
synced 2025-07-08 23:32:45 +00:00
ff2e18f9ca
This change fixes two bugs and adds a safeguard. The first issue is that the schema Gemini expected differed from the one sent, resulting in 400 errors when performing completions. The second issue was that creating a new persona won't define a method for `response_format`. This has to be explicitly defined when we wrap it inside the Persona class. Also, There was a mismatch between the default value and what we stored in the DB. Some parts of the code expected symbols as keys and others as strings. Finally, we add a safeguard when, even if asked to, the model refuses to reply with a valid JSON. In this case, we are making a best-effort to recover and stream the raw response.
62 lines
1.8 KiB
Ruby
62 lines
1.8 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
module DiscourseAi
|
|
module Completions
|
|
class StructuredOutput
|
|
def initialize(json_schema_properties)
|
|
@property_names = json_schema_properties.keys.map(&:to_sym)
|
|
@property_cursors =
|
|
json_schema_properties.reduce({}) do |m, (k, prop)|
|
|
m[k.to_sym] = 0 if prop[:type] == "string"
|
|
m
|
|
end
|
|
|
|
@tracked = {}
|
|
|
|
@raw_response = +""
|
|
@raw_cursor = 0
|
|
|
|
@partial_json_tracker = JsonStreamingTracker.new(self)
|
|
end
|
|
|
|
attr_reader :last_chunk_buffer
|
|
|
|
def <<(raw)
|
|
@raw_response << raw
|
|
@partial_json_tracker << raw
|
|
end
|
|
|
|
def read_buffered_property(prop_name)
|
|
# Safeguard: If the model is misbehaving and generating something that's not a JSON,
|
|
# treat response as a normal string.
|
|
# This is a best-effort to recover from an unexpected scenario.
|
|
if @partial_json_tracker.broken?
|
|
unread_chunk = @raw_response[@raw_cursor..]
|
|
@raw_cursor = @raw_response.length
|
|
return unread_chunk
|
|
end
|
|
|
|
# Maybe we haven't read that part of the JSON yet.
|
|
return nil if @tracked[prop_name].blank?
|
|
|
|
# This means this property is a string and we want to return unread chunks.
|
|
if @property_cursors[prop_name].present?
|
|
unread = @tracked[prop_name][@property_cursors[prop_name]..]
|
|
@property_cursors[prop_name] = @tracked[prop_name].length
|
|
unread
|
|
else
|
|
# Ints and bools are always returned as is.
|
|
@tracked[prop_name]
|
|
end
|
|
end
|
|
|
|
def notify_progress(key, value)
|
|
key_sym = key.to_sym
|
|
return if !@property_names.include?(key_sym)
|
|
|
|
@tracked[key_sym] = value
|
|
end
|
|
end
|
|
end
|
|
end
|