fix: handle no content for gemini provider

Kujtim Hoxha created 3 months ago

Change summary

internal/llm/agent/agent.go     | 10 ++++++++--
internal/llm/provider/gemini.go |  9 +++++++++
2 files changed, 17 insertions(+), 2 deletions(-)

Detailed changes

internal/llm/agent/agent.go

@@ -395,7 +395,7 @@ func (a *agent) processGeneration(ctx context.Context, sessionID, content string
 			defer log.RecoverPanic("agent.Run", func() {
 				slog.Error("panic while generating title")
 			})
-			titleErr := a.generateTitle(context.Background(), sessionID, content)
+			titleErr := a.generateTitle(ctx, sessionID, content)
 			if titleErr != nil && !errors.Is(titleErr, context.Canceled) && !errors.Is(titleErr, context.DeadlineExceeded) {
 				slog.Error("failed to generate title", "error", titleErr)
 			}
@@ -996,11 +996,17 @@ func (a *agent) UpdateModel() error {
 		return fmt.Errorf("provider %s not found in config", largeModelCfg.Provider)
 	}
 
+	var maxTitleTokens int64 = 40
+
+	// if the max output is too low for the gemini provider it won't return anything
+	if smallModelCfg.Provider == "gemini" {
+		maxTitleTokens = 1000
+	}
 	// Recreate title provider
 	titleOpts := []provider.ProviderClientOption{
 		provider.WithModel(config.SelectedModelTypeSmall),
 		provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),
-		provider.WithMaxTokens(40),
+		provider.WithMaxTokens(maxTitleTokens),
 	}
 	newTitleProvider, err := provider.NewProvider(smallModelProviderCfg, titleOpts...)
 	if err != nil {

internal/llm/provider/gemini.go

@@ -322,6 +322,7 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
 			for _, part := range lastMsg.Parts {
 				lastMsgParts = append(lastMsgParts, *part)
 			}
+
 			for resp, err := range chat.SendMessageStream(ctx, lastMsgParts...) {
 				if err != nil {
 					retry, after, retryErr := g.shouldRetry(attempts, err)
@@ -385,6 +386,9 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
 							}
 						}
 					}
+				} else {
+					// no content received
+					break
 				}
 			}
 
@@ -408,6 +412,11 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
 					},
 				}
 				return
+			} else {
+				eventChan <- ProviderEvent{
+					Type:  EventError,
+					Error: errors.New("no content received"),
+				}
 			}
 		}
 	}()