126 lines
3.5 KiB
Go
126 lines
3.5 KiB
Go
package api_test
|
|
|
|
import (
|
|
"context"
|
|
openroute2 "easy_rag/internal/llm/openroute"
|
|
"fmt"
|
|
"testing"
|
|
)
|
|
|
|
func TestFetchChatCompletions(t *testing.T) {
|
|
client := openroute2.NewOpenRouterClient("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")
|
|
|
|
request := openroute2.Request{
|
|
Model: "qwen/qwen2.5-vl-72b-instruct:free",
|
|
Messages: []openroute2.MessageRequest{
|
|
{openroute2.RoleUser, "Привет!", "", ""},
|
|
},
|
|
}
|
|
|
|
output, err := client.FetchChatCompletions(request)
|
|
if err != nil {
|
|
t.Errorf("error %v", err)
|
|
}
|
|
|
|
t.Logf("output: %v", output.Choices[0].Message.Content)
|
|
}
|
|
|
|
func TestFetchChatCompletionsStreaming(t *testing.T) {
|
|
client := openroute2.NewOpenRouterClient("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")
|
|
|
|
request := openroute2.Request{
|
|
Model: "qwen/qwen2.5-vl-72b-instruct:free",
|
|
Messages: []openroute2.MessageRequest{
|
|
{openroute2.RoleUser, "Привет!", "", ""},
|
|
},
|
|
Stream: true,
|
|
}
|
|
|
|
outputChan := make(chan openroute2.Response)
|
|
processingChan := make(chan interface{})
|
|
errChan := make(chan error)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
go client.FetchChatCompletionsStream(request, outputChan, processingChan, errChan, ctx)
|
|
|
|
for {
|
|
select {
|
|
case output := <-outputChan:
|
|
if len(output.Choices) > 0 {
|
|
t.Logf("%s", output.Choices[0].Delta.Content)
|
|
}
|
|
case <-processingChan:
|
|
t.Logf("Обработка\n")
|
|
case err := <-errChan:
|
|
if err != nil {
|
|
t.Errorf("Ошибка: %v", err)
|
|
return
|
|
}
|
|
return
|
|
case <-ctx.Done():
|
|
fmt.Println("Контекст отменен:", ctx.Err())
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestFetchChatCompletionsAgentStreaming(t *testing.T) {
|
|
client := openroute2.NewOpenRouterClient("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")
|
|
agent := openroute2.NewRouterAgent(client, "qwen/qwen2.5-vl-72b-instruct:freet", openroute2.RouterAgentConfig{
|
|
Temperature: 0.7,
|
|
MaxTokens: 100,
|
|
})
|
|
|
|
outputChan := make(chan openroute2.Response)
|
|
processingChan := make(chan interface{})
|
|
errChan := make(chan error)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
chat := []openroute2.MessageRequest{
|
|
{Role: openroute2.RoleSystem, Content: "Вы полезный помощник."},
|
|
{Role: openroute2.RoleUser, Content: "Привет!"},
|
|
}
|
|
|
|
go agent.ChatStream(chat, outputChan, processingChan, errChan, ctx)
|
|
|
|
for {
|
|
select {
|
|
case output := <-outputChan:
|
|
if len(output.Choices) > 0 {
|
|
t.Logf("%s", output.Choices[0].Delta.Content)
|
|
}
|
|
case <-processingChan:
|
|
t.Logf("Обработка\n")
|
|
case err := <-errChan:
|
|
if err != nil {
|
|
t.Errorf("Ошибка: %v", err)
|
|
return
|
|
}
|
|
return
|
|
case <-ctx.Done():
|
|
fmt.Println("Контекст отменен:", ctx.Err())
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestFetchChatCompletionsAgentSimpleChat(t *testing.T) {
|
|
client := openroute2.NewOpenRouterClient("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")
|
|
agent := openroute2.NewRouterAgentChat(client, "qwen/qwen2.5-vl-72b-instruct:free", openroute2.RouterAgentConfig{
|
|
Temperature: 0.0,
|
|
MaxTokens: 100,
|
|
}, "Вы полезный помощник, отвечайте короткими словами.")
|
|
|
|
agent.Chat("Запомни это: \"wojtess\"")
|
|
agent.Chat("Что я просил вас запомнить?")
|
|
|
|
for _, msg := range agent.Messages {
|
|
content, ok := msg.Content.(string)
|
|
if ok {
|
|
t.Logf("%s: %s", msg.Role, content)
|
|
}
|
|
}
|
|
}
|