diff --git a/cmd/internal/llm/main.go b/cmd/internal/llm/main.go new file mode 100644 index 0000000..1b4ac72 --- /dev/null +++ b/cmd/internal/llm/main.go @@ -0,0 +1,129 @@ +package main + +import ( + "bytes" + "encoding/json" + "flag" + "fmt" + "io" + "io/fs" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +type AIMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type AIRequest struct { + Model string `json:"model"` + MaxTokens int `json:"max_tokens"` + Messages []AIMessage `json:"messages"` +} + +type AIResponse struct { + ID string `json:"id"` + Model string `json:"model"` + Choices []struct { + Message struct { + Role string `json:"role"` + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` +} + +func fetchAllScripts() map[string]string { + files := map[string]string{} + + filepath.Walk("./assets", func(path string, info fs.FileInfo, err error) error { + if !strings.HasSuffix(path, ".lua") { + return nil + } + + res, err := os.ReadFile(path) + if err != nil { + return err + } + + files[path] = string(res) + return nil + }) + + return files +} + +func main() { + model := flag.String("model", "gpt-3.5-turbo", "The model to use") + apiKey := flag.String("api-key", "", "The API key to use") + prompt := flag.String("prompt", "", "The prompt to use") + flag.Parse() + + if *apiKey == "" { + panic("API key is required") + } + + files := fetchAllScripts() + + aiRequest := AIRequest{ + Model: *model, + MaxTokens: 500000, + Messages: []AIMessage{ + {Role: "system", Content: `You are a code generator for a deck building card game. +You are given a set of Lua scripts that define the cards and artifacts in the game. +Your task is to generate a new Lua script that defines a new card or artifact for the game. +The new card or artifact should be unique and interesting, and should fit well into the existing game mechanics. +The new card or artifact should also be balanced and fair, and should not be overpowered or underpowered. +You can use the existing Lua scripts as a reference, and you can also use your own creativity and imagination to come up with new ideas. +The new card or artifact should be well documented, with clear and concise comments that explain how it works and how it interacts with other cards and artifacts in the game.`}, + }, + } + + for key, value := range files { + aiRequest.Messages = append(aiRequest.Messages, AIMessage{Role: "system", Content: fmt.Sprintf("File: %s\n\n%s", key, value)}) + } + + if *prompt != "" { + aiRequest.Messages = append(aiRequest.Messages, AIMessage{Role: "user", Content: *prompt}) + } + + body, err := json.Marshal(aiRequest) + if err != nil { + panic(err) + } + + req, err := http.NewRequest("POST", "https://openrouter.ai/api/v1/chat/completions", bytes.NewBuffer(body)) + if err != nil { + panic(err) + } + + client := &http.Client{ + Timeout: time.Second * 60 * 10, + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+*apiKey) + req.Header.Set("HTTP-Referer", "https://github.com/BigJk/end_of_eden") + req.Header.Set("X-Title", "End of Eden") + + resp, err := client.Do(req) + if err != nil { + panic(err) + } + + var aiResp AIResponse + respBody, err := io.ReadAll(resp.Body) + if err != nil { + panic(err) + } + + err = json.Unmarshal(respBody, &aiResp) + if err != nil { + panic(err) + } + + fmt.Println(aiResp.Choices[0].Message.Content) +}