website/llm.go at main · gosuda/website · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
package main

import (
	"context"
	"os"

	"github.com/lemon-mint/coord"
	"github.com/lemon-mint/coord/llm"
	"github.com/lemon-mint/coord/pconf"
	"github.com/lemon-mint/coord/provider"
	_ "github.com/lemon-mint/coord/provider/aistudio"
	_ "github.com/lemon-mint/coord/provider/anthropic"
	_ "github.com/lemon-mint/coord/provider/openai"
	_ "github.com/lemon-mint/coord/provider/vertexai"
	"github.com/pemistahl/lingua-go"
	"github.com/rs/zerolog/log"
	"golang.org/x/time/rate"
)

var llmClient provider.LLMClient
var llmModel llm.Model
var languageDetector lingua.LanguageDetector

func init() {
	languages := []lingua.Language{
		lingua.English,
		lingua.Spanish,
		lingua.Chinese,
		lingua.Korean,
		lingua.Japanese,
		lingua.German,
		lingua.Russian,
		lingua.French,
		lingua.Dutch,
		lingua.Italian,
		lingua.Indonesian,
		lingua.Portuguese,
		lingua.Swedish,
		lingua.Czech,
	}

	languageDetector = lingua.NewLanguageDetectorBuilder().
		FromLanguages(languages...).
		Build()

	if os.Getenv("LLM_INIT") == "false" || os.Getenv("LLM_INIT") == "0" {
		log.Info().Msg("llm init skipped")
		return
	}

	var err error
	var client provider.LLMClient

	if os.Getenv("PROVIDER") == "aistudio" {
		log.Debug().Msg("initializing llm client")
		client, err = coord.NewLLMClient(
			context.Background(),
			"aistudio",
			pconf.WithAPIKey(os.Getenv("AI_STUDIO_API_KEY")),
		)
		if err != nil {
			log.Fatal().Err(err).Msg("failed to create llm client")
		}
	} else {
		log.Debug().Str("location", os.Getenv("LOCATION")).Str("project_id", os.Getenv("PROJECT_ID")).Msg("initializing llm client")
		client, err = coord.NewLLMClient(
			context.Background(),
			"vertexai",
			pconf.WithLocation(os.Getenv("LOCATION")),
			pconf.WithProjectID(os.Getenv("PROJECT_ID")),
		)
		if err != nil {
			log.Fatal().Err(err).Msg("failed to create llm client")
		}
	}
	llmClient = client
	log.Debug().Msg("llm client initialized")

	llmModel, err = llmClient.NewLLM("gemini-2.5-flash", &llm.Config{
		Temperature:           Ptr(float32(0.7)),
		MaxOutputTokens:       Ptr(65535),
		SafetyFilterThreshold: llm.BlockOff,
		ThinkingConfig: &llm.ThinkingConfig{
			ThinkingBudget: pconf.Ptrify(0),
		},
	})

	if err != nil {
		log.Fatal().Err(err).Msg("failed to create llm model")
	}

	//llmModel = newRateLimitModel(llmModel, rate.Every(time.Minute/9))
	log.Debug().Msg("llm model initialized")
}

func Ptr[T any](t T) *T {
	return &t
}

type rateLimitModel struct {
	llm.Model
	limit   rate.Limit
	limiter *rate.Limiter
}

func newRateLimitModel(model llm.Model, limit rate.Limit) *rateLimitModel {
	return &rateLimitModel{
		Model:   model,
		limit:   limit,
		limiter: rate.NewLimiter(limit, 1),
	}
}

func (r *rateLimitModel) GenerateStream(ctx context.Context, chat *llm.ChatContext, input *llm.Content) *llm.StreamContent {
	if err := r.limiter.Wait(ctx); err != nil {
		ch := make(chan llm.Segment)
		close(ch)
		return &llm.StreamContent{
			Err:          err,
			Content:      &llm.Content{},
			FinishReason: llm.FinishReasonError,
			Stream:       ch,
		}
	}
	return r.Model.GenerateStream(ctx, chat, input)
}