Skip to content

Commit

Permalink
fix: add context timeout to LLM query and limit the return length to …
Browse files Browse the repository at this point in the history
…less than 10000 characters.
  • Loading branch information
wwcchh0123 committed Dec 9, 2024
1 parent a78f42f commit 9c7bd8d
Showing 1 changed file with 18 additions and 2 deletions.
20 changes: 18 additions & 2 deletions internal/llm/llm.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@ import (
"context"
"errors"
"fmt"
"time"

"github.com/qiniu/reviewbot/internal/util"
"github.com/tmc/langchaingo/llms"
)

Expand All @@ -30,11 +32,17 @@ func New(ctx context.Context, config Config) (llms.Model, error) {
}

func Query(ctx context.Context, model llms.Model, query string, extraContext string) (string, error) {
log := util.FromContext(ctx)
ragQuery := fmt.Sprintf(ragTemplateStr, query, extraContext)
respText, err := llms.GenerateFromSinglePrompt(ctx, model, ragQuery)

ctxWithTimeout, cancel := context.WithTimeout(ctx, 5*time.Minute)
defer cancel()

respText, err := llms.GenerateFromSinglePrompt(ctxWithTimeout, model, ragQuery)
if err != nil {
return "", err
}
log.Debugf("length of LLM respText: %d", len(respText))
return respText, nil
}

Expand All @@ -60,20 +68,28 @@ Context:
`

func QueryForReference(ctx context.Context, model llms.Model, linterOutput string) (string, error) {
log := util.FromContext(ctx)
if model == nil {
return "", ErrModelIsNil
}
ragQuery := fmt.Sprintf(referenceTemplateStr, linterOutput)
respText, err := llms.GenerateFromSinglePrompt(ctx, model, ragQuery)

ctxWithTimeout, cancel := context.WithTimeout(ctx, 5*time.Minute)
defer cancel()

respText, err := llms.GenerateFromSinglePrompt(ctxWithTimeout, model, ragQuery)
if err != nil {
return "", err
}
log.Debugf("length of LLM respText: %d", len(respText))

return respText, nil
}

const referenceTemplateStr = `
You are a lint expert, you can explain in detail the meaning of the lint result according to the content of the given context.
You will follow the format of the example in <format> to answer in Chinese, firstly, you will explain the lint, secondly, you will give the incorrect usage (it can be a code example or text description), and finally, you will give the correct usage (it can be a code example or a text description), and finally add a blank line before the result.
Please make sure that the output must not exceed 10,000 characters.
<format>
Expand Down

0 comments on commit 9c7bd8d

Please sign in to comment.