feat: add configurable generation timeout
All checks were successful
continuous-integration/drone/push Build is passing
All checks were successful
continuous-integration/drone/push Build is passing
This commit is contained in:
@@ -44,6 +44,34 @@ func New(s store.Store, dataDir string, logger *logrus.Logger, llmEndpoint, llmK
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeSettings(settings *store.Settings) {
|
||||
// Default Text Generation Timeout
|
||||
if settings.TextGenerationTimeoutMinutes == 0 {
|
||||
settings.TextGenerationTimeoutMinutes = 5
|
||||
}
|
||||
|
||||
// Validate Text Generation Timeout
|
||||
switch settings.TextGenerationTimeoutMinutes {
|
||||
case 1, 5, 10, 15, 30:
|
||||
return
|
||||
default:
|
||||
settings.TextGenerationTimeoutMinutes = 5
|
||||
}
|
||||
}
|
||||
|
||||
func (a *API) textGenerationTimeout() time.Duration {
|
||||
// Load Settings
|
||||
settings, err := a.store.GetSettings()
|
||||
if err != nil {
|
||||
a.logger.WithError(err).Error("failed to retrieve settings for text generation timeout")
|
||||
return 5 * time.Minute
|
||||
}
|
||||
|
||||
// Normalize Timeout
|
||||
normalizeSettings(settings)
|
||||
return time.Duration(settings.TextGenerationTimeoutMinutes) * time.Minute
|
||||
}
|
||||
|
||||
func (a *API) GetSettings(w http.ResponseWriter, r *http.Request) {
|
||||
log := a.logger.WithField("handler", "GetSettingsHandler")
|
||||
|
||||
@@ -54,6 +82,9 @@ func (a *API) GetSettings(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
// Normalize Settings
|
||||
normalizeSettings(settings)
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(settings); err != nil {
|
||||
log.WithError(err).Error("failed to encode application settings response")
|
||||
@@ -72,6 +103,9 @@ func (a *API) PostSettings(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
// Normalize Settings
|
||||
normalizeSettings(&newSettings)
|
||||
|
||||
if err := a.store.SaveSettings(&newSettings); err != nil {
|
||||
log.WithError(err).Error("failed to save settings")
|
||||
http.Error(w, "Failed to save application settings", http.StatusInternalServerError)
|
||||
@@ -589,7 +623,7 @@ func (a *API) startMessageGeneration(chatID uuid.UUID, chatModel, userMessage st
|
||||
|
||||
func (a *API) runMessageGeneration(apiClient *client.Client, chat *store.Chat, assistantMsg *store.Message, chatModel string, gen *generation) {
|
||||
// Create Generation Context
|
||||
ctx, cancel := context.WithTimeout(gen.ctx, time.Minute*5)
|
||||
ctx, cancel := context.WithTimeout(gen.ctx, a.textGenerationTimeout())
|
||||
defer cancel()
|
||||
|
||||
// Send Message
|
||||
|
||||
@@ -141,9 +141,10 @@ func TestInMemoryStore_SaveSettings(t *testing.T) {
|
||||
store := NewInMemoryStore()
|
||||
|
||||
settings := &Settings{
|
||||
ImageEditSelector: ".image-edit",
|
||||
ImageGenerationSelector: ".image-gen",
|
||||
TextGenerationSelector: ".text-gen",
|
||||
ImageEditSelector: ".image-edit",
|
||||
ImageGenerationSelector: ".image-gen",
|
||||
TextGenerationSelector: ".text-gen",
|
||||
TextGenerationTimeoutMinutes: 10,
|
||||
}
|
||||
|
||||
err := store.SaveSettings(settings)
|
||||
@@ -160,9 +161,10 @@ func TestInMemoryStore_GetSettings(t *testing.T) {
|
||||
|
||||
// Set some settings
|
||||
settings = &Settings{
|
||||
ImageEditSelector: ".image-edit",
|
||||
ImageGenerationSelector: ".image-gen",
|
||||
TextGenerationSelector: ".text-gen",
|
||||
ImageEditSelector: ".image-edit",
|
||||
ImageGenerationSelector: ".image-gen",
|
||||
TextGenerationSelector: ".text-gen",
|
||||
TextGenerationTimeoutMinutes: 10,
|
||||
}
|
||||
err = store.SaveSettings(settings)
|
||||
require.NoError(t, err)
|
||||
@@ -171,4 +173,5 @@ func TestInMemoryStore_GetSettings(t *testing.T) {
|
||||
settings, err = store.GetSettings()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, ".image-edit", settings.ImageEditSelector)
|
||||
assert.Equal(t, 10, settings.TextGenerationTimeoutMinutes)
|
||||
}
|
||||
|
||||
@@ -15,9 +15,10 @@ var _ Store = (*FileStore)(nil)
|
||||
|
||||
// Settings represents the application settings
|
||||
type Settings struct {
|
||||
ImageEditSelector string `json:"image_edit_selector,omitempty"`
|
||||
ImageGenerationSelector string `json:"image_generation_selector,omitempty"`
|
||||
TextGenerationSelector string `json:"text_generation_selector,omitempty"`
|
||||
ImageEditSelector string `json:"image_edit_selector,omitempty"`
|
||||
ImageGenerationSelector string `json:"image_generation_selector,omitempty"`
|
||||
TextGenerationSelector string `json:"text_generation_selector,omitempty"`
|
||||
TextGenerationTimeoutMinutes int `json:"text_generation_timeout_minutes,omitempty"`
|
||||
}
|
||||
|
||||
// FileStore implements the Store interface using a file-based storage
|
||||
|
||||
@@ -209,9 +209,10 @@ func TestFileStore_SaveSettings(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
|
||||
settings := &Settings{
|
||||
ImageEditSelector: ".image-edit",
|
||||
ImageGenerationSelector: ".image-gen",
|
||||
TextGenerationSelector: ".text-gen",
|
||||
ImageEditSelector: ".image-edit",
|
||||
ImageGenerationSelector: ".image-gen",
|
||||
TextGenerationSelector: ".text-gen",
|
||||
TextGenerationTimeoutMinutes: 10,
|
||||
}
|
||||
|
||||
err = store.SaveSettings(settings)
|
||||
@@ -236,9 +237,10 @@ func TestFileStore_GetSettings(t *testing.T) {
|
||||
|
||||
// Set some settings
|
||||
settings = &Settings{
|
||||
ImageEditSelector: ".image-edit",
|
||||
ImageGenerationSelector: ".image-gen",
|
||||
TextGenerationSelector: ".text-gen",
|
||||
ImageEditSelector: ".image-edit",
|
||||
ImageGenerationSelector: ".image-gen",
|
||||
TextGenerationSelector: ".text-gen",
|
||||
TextGenerationTimeoutMinutes: 10,
|
||||
}
|
||||
err = store.SaveSettings(settings)
|
||||
require.NoError(t, err)
|
||||
@@ -247,4 +249,5 @@ func TestFileStore_GetSettings(t *testing.T) {
|
||||
settings, err = store.GetSettings()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, ".image-edit", settings.ImageEditSelector)
|
||||
assert.Equal(t, 10, settings.TextGenerationTimeoutMinutes)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user