package llm

import (
	"agent/src/utils"
	"context"
	"fmt"
	"testing"

	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/ollama"
)

func TestOllama(t *testing.T) {
	// llm, err := ollama.New(ollama.WithServerURL("http://192.168.1.5:11434"), ollama.WithModel("deepseek-r1:32b"))
	llm, err := ollama.New(ollama.WithServerURL("http://ishangsf.com:11434"), ollama.WithModel("qwen2.5-coder:7b"))
	if err != nil {
		panic(err)
	}
	ctx := context.Background()

	content := []llms.MessageContent{
		llms.TextParts(llms.ChatMessageTypeSystem, "You are a company branding design wizard."),
		llms.TextParts(llms.ChatMessageTypeHuman, "What would be a good company name a company that makes colorful socks?"),
	}
	llm.GenerateContent(ctx, content, llms.WithStreamingFunc(func(ctx context.Context, chunk []byte) error {
		fmt.Print(string(chunk))
		return nil
	}))
}

func TestInitLocalModel(t *testing.T) {
	err := InitLocalModel([]string{"http://ishangsf.com:11434"})
	if err != nil {
		panic(err)
	}
	utils.PrintJson(LLMNames)
}