Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions cmd/generate/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ func (h *generateCommandHandler) CreateContextFromPrompt() (*PromptPexContext, e
}

runID := fmt.Sprintf("run_%d", time.Now().Unix())
context := &PromptPexContext{
promptContext := &PromptPexContext{
// Unique identifier for the run
RunID: runID,
// The prompt content and metadata
Expand All @@ -50,21 +50,21 @@ func (h *generateCommandHandler) CreateContextFromPrompt() (*PromptPexContext, e
} else {
sessionInfo = fmt.Sprintf("reloading session file at %s", *h.sessionFile)
// Check if prompt hashes match
if existingContext.PromptHash != context.PromptHash {
if existingContext.PromptHash != promptContext.PromptHash {
return nil, fmt.Errorf("prompt changed unable to reuse session file")
}

// Merge existing context data
if existingContext != nil {
context = mergeContexts(existingContext, context)
promptContext = mergeContexts(existingContext, promptContext)
}
}
}

h.WriteToParagraph(RenderMessagesToString(context.Prompt.Messages))
h.WriteToParagraph(RenderMessagesToString(promptContext.Prompt.Messages))
h.WriteEndBox(sessionInfo)

return context, nil
return promptContext, nil
}

// loadContextFromFile loads a PromptPexContext from a JSON file
Expand Down
6 changes: 3 additions & 3 deletions cmd/generate/generate.go
Original file line number Diff line number Diff line change
Expand Up @@ -76,14 +76,14 @@ func NewGenerateCommand(cfg *command.Config) *cobra.Command {
sessionFile: util.Ptr(sessionFile),
}

// Create context
context, err := handler.CreateContextFromPrompt()
// Create prompt context
promptContext, err := handler.CreateContextFromPrompt()
if err != nil {
return fmt.Errorf("failed to create context: %w", err)
}

// Run the PromptPex pipeline
if err := handler.RunTestGenerationPipeline(context); err != nil {
if err := handler.RunTestGenerationPipeline(promptContext); err != nil {
// Disable usage help for pipeline failures
cmd.SilenceUsage = true
return fmt.Errorf("pipeline failed: %w", err)
Expand Down
17 changes: 13 additions & 4 deletions cmd/generate/llm.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,10 @@ func (h *generateCommandHandler) callModelWithRetry(step string, req azuremodels
for attempt := 0; attempt <= maxRetries; attempt++ {
sp := spinner.New(spinner.CharSets[14], 100*time.Millisecond, spinner.WithWriter(h.cfg.ErrOut))
sp.Start()
//nolint:gocritic,revive // TODO
defer sp.Stop()

resp, err := h.client.GetChatCompletionStream(ctx, req, h.org)
if err != nil {
sp.Stop()
var rateLimitErr *azuremodels.RateLimitError
if errors.As(err, &rateLimitErr) {
if attempt < maxRetries {
Expand All @@ -53,8 +52,6 @@ func (h *generateCommandHandler) callModelWithRetry(step string, req azuremodels
return "", err
}
reader := resp.Reader
//nolint:gocritic,revive // TODO
defer reader.Close()

var content strings.Builder
for {
Expand All @@ -63,6 +60,11 @@ func (h *generateCommandHandler) callModelWithRetry(step string, req azuremodels
if errors.Is(err, context.Canceled) || strings.Contains(err.Error(), "EOF") {
break
}
if closeErr := reader.Close(); closeErr != nil {
// Log close error but don't override the original error
fmt.Fprintf(h.cfg.ErrOut, "Warning: failed to close reader: %v\n", closeErr)
}
sp.Stop()
return "", err
}
for _, choice := range completion.Choices {
Expand All @@ -75,6 +77,13 @@ func (h *generateCommandHandler) callModelWithRetry(step string, req azuremodels
}
}

// Properly close reader and stop spinner before returning success
err = reader.Close()
sp.Stop()
if err != nil {
return "", fmt.Errorf("failed to close reader: %w", err)
}

res := strings.TrimSpace(content.String())
h.LogLLMResponse(res)
return res, nil
Expand Down
Loading