package main

import (
	"context"
	"fmt"
	"log"
	"os"

	"github.com/joho/godotenv"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/openai"
)

var (
	apiKey  string
	baseUrl string
	model   = "stepfun-ai/Step-3.5-Flash"
)

func init() {
	err := godotenv.Load()
	if err != nil {
		fmt.Println(".env 文件加载失败")
	}
	apiKey = os.Getenv("API_KEY")
	baseUrl = os.Getenv("BASE_URL")
	fmt.Printf("API Key: %s\n", apiKey)
	fmt.Printf("Base URL: %s\n", baseUrl)
}

func main() {
	llm, err := openai.New(
		openai.WithBaseURL(baseUrl),
		openai.WithToken(apiKey),
		openai.WithModel(model),
	)

	if err != nil {
		log.Fatal(err)
	}

	ctx := context.Background()

	response, err := llms.GenerateFromSinglePrompt(
		ctx,
		llm,
		"Hello! My name is wonder!",
	)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("AI:", response)

}