Добавил интеграцию с LLM
This commit is contained in:
@@ -0,0 +1,81 @@
|
||||
package llm_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.vakhrushev.me/av/jellybit/internal/llm"
|
||||
)
|
||||
|
||||
// TestIntegration_OpenAICompat бьётся в реальный эндпоинт. По умолчанию
|
||||
// пропускается; включается переменными окружения:
|
||||
//
|
||||
// JELLYBIT_LLM_BASE_URL=https://bothub.chat/api/v2/openai/v1 \
|
||||
// JELLYBIT_LLM_API_KEY=... \
|
||||
// JELLYBIT_LLM_MODEL=deepseek-v4-flash \
|
||||
// go test ./internal/llm/ -run Integration -v
|
||||
func TestIntegration_OpenAICompat(t *testing.T) {
|
||||
base := os.Getenv("JELLYBIT_LLM_BASE_URL")
|
||||
key := os.Getenv("JELLYBIT_LLM_API_KEY")
|
||||
model := os.Getenv("JELLYBIT_LLM_MODEL")
|
||||
if base == "" || model == "" {
|
||||
t.Skip("set JELLYBIT_LLM_BASE_URL and JELLYBIT_LLM_MODEL to run")
|
||||
}
|
||||
|
||||
p, err := llm.New(llm.Config{
|
||||
Type: "openai-compat",
|
||||
BaseURL: base,
|
||||
APIKey: key,
|
||||
Model: model,
|
||||
Timeout: 90 * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
|
||||
defer cancel()
|
||||
|
||||
zero := 0.0
|
||||
resp, err := p.Complete(ctx, llm.Request{
|
||||
JSONMode: true,
|
||||
Temperature: &zero,
|
||||
MaxTokens: 2000,
|
||||
Messages: []llm.Message{
|
||||
{Role: llm.RoleSystem, Content: `Распознай медиа. Ответь только JSON по схеме: ` +
|
||||
`{"kind":"movie|series","title":string,"year":number,"season":number|null}`},
|
||||
{Role: llm.RoleUser, Content: "Аватар: Легенда об Аанге / Avatar: The Last Airbender / " +
|
||||
"Книга 2: Земля [2006, США, DVDRip-AVC]"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Complete: %v", err)
|
||||
}
|
||||
t.Logf("model=%s usage=%+v content=%s", resp.Model, resp.Usage, resp.Content)
|
||||
|
||||
raw, err := llm.ExtractJSONObject(resp.Content)
|
||||
if err != nil {
|
||||
t.Fatalf("ExtractJSONObject: %v (content: %q)", err, resp.Content)
|
||||
}
|
||||
var plan struct {
|
||||
Kind string `json:"kind"`
|
||||
Title string `json:"title"`
|
||||
Year int `json:"year"`
|
||||
Season *int `json:"season"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(raw), &plan); err != nil {
|
||||
t.Fatalf("unmarshal plan: %v (raw: %s)", err, raw)
|
||||
}
|
||||
if plan.Kind != "series" {
|
||||
t.Errorf("kind = %q, want series", plan.Kind)
|
||||
}
|
||||
if plan.Year != 2006 {
|
||||
t.Errorf("year = %d, want 2006", plan.Year)
|
||||
}
|
||||
if plan.Season == nil || *plan.Season != 2 {
|
||||
t.Errorf("season = %v, want 2", plan.Season)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user