385 lines
11 KiB
Go
385 lines
11 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"io/ioutil"
|
|
"log"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"regexp"
|
|
"strings"
|
|
"unicode"
|
|
|
|
"gopkg.in/yaml.v3"
|
|
)
|
|
|
|
// ImageExtensionRegex is a regular expression to validate image file extensions.
|
|
var ImageExtensionRegex = regexp.MustCompile(`(?i)\.(jpg|jpeg|png|gif)$`)
|
|
|
|
// OpenAIRequest represents the structure of the request to the OpenAI API.
|
|
type OpenAIRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []Message `json:"messages"`
|
|
Temperature float64 `json:"temperature,omitempty"`
|
|
}
|
|
|
|
// Message represents a message in the OpenAI API request.
|
|
type Message struct {
|
|
Role string `json:"role"`
|
|
Content []Content `json:"content"`
|
|
}
|
|
|
|
type Content struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text,omitempty"`
|
|
ImageURL *ImageURL `json:"image_url,omitempty"`
|
|
}
|
|
|
|
type ImageURL struct {
|
|
URL string `json:"url"`
|
|
}
|
|
|
|
// OpenAIResponse represents the structure of the response from the OpenAI API.
|
|
type OpenAIResponse struct {
|
|
Choices []struct {
|
|
Message struct {
|
|
Content string `json:"content"`
|
|
} `json:"message"`
|
|
} `json:"choices"`
|
|
}
|
|
|
|
// Config represents the structure of the configuration file.
|
|
type Config struct {
|
|
APIKey string `yaml:"apikey"`
|
|
Model string `yaml:"model"`
|
|
Endpoint string `yaml:"endpoint"`
|
|
Temperature float64 `yaml:"temperature"`
|
|
}
|
|
|
|
func main() {
|
|
var configFile string
|
|
flag.StringVar(&configFile, "c", "", "Path to config file")
|
|
flag.Parse()
|
|
|
|
// Get the config file path
|
|
configFilePath, err := getConfigFilePath(configFile)
|
|
if err != nil {
|
|
log.Fatalf("Error getting config file path: %v", err)
|
|
}
|
|
|
|
// Load the configuration from the file
|
|
config, err := loadConfig(configFilePath)
|
|
if err != nil {
|
|
log.Fatalf("Error loading config: %v", err)
|
|
}
|
|
|
|
// Get the image filename from the command-line arguments
|
|
if len(flag.Args()) < 1 {
|
|
fmt.Println("Usage: gonamer [-c config.yaml] <image_filename>")
|
|
flag.PrintDefaults()
|
|
return
|
|
}
|
|
imageFilename := flag.Arg(0)
|
|
|
|
// Validate the image filename extension
|
|
if !ImageExtensionRegex.MatchString(imageFilename) {
|
|
fmt.Println("Error: Invalid image file extension. Only .jpg, .jpeg, .png, and .gif are supported.")
|
|
return
|
|
}
|
|
|
|
// Get the suggested filename from the LLM
|
|
newFilename, err := getSuggestedFilename(imageFilename, config)
|
|
if err != nil {
|
|
fmt.Println("Error getting suggested filename:", err)
|
|
return
|
|
}
|
|
|
|
// Rename the file
|
|
err = renameFile(imageFilename, newFilename)
|
|
if err != nil {
|
|
fmt.Println("Error renaming file:", err)
|
|
return
|
|
}
|
|
|
|
fmt.Printf("Successfully renamed '%s' to '%s'\n", imageFilename, newFilename)
|
|
}
|
|
|
|
// getConfigFilePath returns the path to the config file.
|
|
func getConfigFilePath(configFile string) (string, error) {
|
|
if configFile != "" {
|
|
return configFile, nil
|
|
}
|
|
homeDir, err := os.UserHomeDir()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return filepath.Join(homeDir, ".config", "gonamer.yaml"), nil
|
|
}
|
|
|
|
// loadConfig loads the configuration from the specified file.
|
|
func loadConfig(configFilePath string) (*Config, error) {
|
|
data, err := ioutil.ReadFile(configFilePath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var config Config
|
|
err = yaml.Unmarshal(data, &config)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &config, nil
|
|
}
|
|
|
|
// getSuggestedFilename sends the image to the LLM and returns the suggested filename.
|
|
func getSuggestedFilename(imageFilename string, config *Config) (string, error) {
|
|
// 1. Encode the image to base64
|
|
imageData, err := os.ReadFile(imageFilename)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error reading image file: %w", err)
|
|
}
|
|
base64Image := base64.StdEncoding.EncodeToString(imageData)
|
|
|
|
// 2. Create the request payload
|
|
requestData := OpenAIRequest{
|
|
Model: config.Model,
|
|
Messages: []Message{
|
|
{
|
|
Role: "user",
|
|
Content: []Content{
|
|
{
|
|
Type: "text",
|
|
Text: "Suggest a useful, descriptive, and unique filename for this image, considering its content and context. Only provide the filename, no other text or explanation. Do not include the file extension.",
|
|
},
|
|
{
|
|
Type: "image_url",
|
|
ImageURL: &ImageURL{
|
|
URL: fmt.Sprintf("data:image/%s;base64,%s", strings.TrimPrefix(filepath.Ext(imageFilename), "."), base64Image),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
Temperature: config.Temperature,
|
|
}
|
|
|
|
requestBody, err := json.Marshal(requestData)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error marshalling request data: %w", err)
|
|
}
|
|
|
|
// 3. Send the request to the LLM (using JSON request as default)
|
|
originalExtension := filepath.Ext(imageFilename)
|
|
newFilename, err := sendJSONRequest(requestBody, originalExtension, config)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
return newFilename, nil
|
|
}
|
|
|
|
// sendJSONRequest sends the request as JSON to the LLM endpoint
|
|
func sendJSONRequest(requestBody []byte, originalExtension string, config *Config) (string, error) {
|
|
req, err := http.NewRequest("POST", config.Endpoint, bytes.NewBuffer(requestBody))
|
|
if err != nil {
|
|
return "", fmt.Errorf("error creating HTTP request: %w", err)
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+config.APIKey)
|
|
|
|
client := &http.Client{}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error sending HTTP request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
// 4. Process the response
|
|
responseBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error reading response body: %w", err)
|
|
}
|
|
|
|
// Print the raw response body
|
|
//fmt.Println("Raw Response Body:")
|
|
//fmt.Println(string(responseBody))
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return "", fmt.Errorf("API error (status %d): %s", resp.StatusCode, responseBody)
|
|
}
|
|
|
|
var apiResponse OpenAIResponse
|
|
err = json.Unmarshal(responseBody, &apiResponse)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error unmarshalling response: %w", err)
|
|
}
|
|
|
|
if len(apiResponse.Choices) == 0 || apiResponse.Choices[0].Message.Content == "" {
|
|
return "", fmt.Errorf("empty response from API")
|
|
}
|
|
|
|
suggestedFilename := strings.TrimSpace(apiResponse.Choices[0].Message.Content)
|
|
|
|
// Remove any existing extension from the suggested filename
|
|
suggestedFilename = strings.TrimSuffix(suggestedFilename, filepath.Ext(suggestedFilename))
|
|
|
|
// Ensure the suggested filename has the correct extension
|
|
if !strings.HasSuffix(strings.ToLower(suggestedFilename), strings.ToLower(originalExtension)) {
|
|
suggestedFilename += originalExtension
|
|
}
|
|
|
|
return suggestedFilename, nil
|
|
}
|
|
|
|
// sendMultipartRequest sends the request as multipart/form-data to the LLM endpoint (alternative)
|
|
func sendMultipartRequest(requestBody []byte, originalExtension string, config *Config) (string, error) {
|
|
// Prepare a form
|
|
body := &bytes.Buffer{}
|
|
writer := multipart.NewWriter(body)
|
|
|
|
// Add the JSON part
|
|
jsonPart, err := writer.CreateFormField("json")
|
|
if err != nil {
|
|
return "", fmt.Errorf("error creating form field: %w", err)
|
|
}
|
|
_, err = jsonPart.Write(requestBody)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error writing to form field: %w", err)
|
|
}
|
|
|
|
// Close the multipart writer
|
|
err = writer.Close()
|
|
if err != nil {
|
|
return "", fmt.Errorf("error closing multipart writer: %w", err)
|
|
}
|
|
|
|
// Create a new request
|
|
req, err := http.NewRequest("POST", config.Endpoint, body)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error creating HTTP request: %w", err)
|
|
}
|
|
|
|
// Set headers
|
|
req.Header.Set("Content-Type", writer.FormDataContentType())
|
|
req.Header.Set("Authorization", "Bearer "+config.APIKey)
|
|
|
|
// Send the request
|
|
client := &http.Client{}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error sending HTTP request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
// Process the response (same as in sendJSONRequest)
|
|
responseBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error reading response body: %w", err)
|
|
}
|
|
|
|
// Print the raw response body
|
|
//fmt.Println("Raw Response Body:")
|
|
//fmt.Println(string(responseBody))
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return "", fmt.Errorf("API error (status %d): %s", resp.StatusCode, responseBody)
|
|
}
|
|
|
|
var apiResponse OpenAIResponse
|
|
err = json.Unmarshal(responseBody, &apiResponse)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error unmarshalling response: %w", err)
|
|
}
|
|
|
|
if len(apiResponse.Choices) == 0 || apiResponse.Choices[0].Message.Content == "" {
|
|
return "", fmt.Errorf("empty response from API")
|
|
}
|
|
|
|
suggestedFilename := strings.TrimSpace(apiResponse.Choices[0].Message.Content)
|
|
|
|
// Remove any existing extension from the suggested filename
|
|
suggestedFilename = strings.TrimSuffix(suggestedFilename, filepath.Ext(suggestedFilename))
|
|
|
|
// Ensure the suggested filename has the correct extension
|
|
if !strings.HasSuffix(strings.ToLower(suggestedFilename), strings.ToLower(originalExtension)) {
|
|
suggestedFilename += originalExtension
|
|
}
|
|
|
|
return suggestedFilename, nil
|
|
}
|
|
|
|
// renameFile renames the file, handling potential issues like existing files.
|
|
func renameFile(oldFilename, newFilename string) error {
|
|
// Sanitize the new filename (remove invalid characters, etc.)
|
|
newFilename = sanitizeFilename(newFilename)
|
|
|
|
// Check if a file with the new name already exists
|
|
if _, err := os.Stat(newFilename); err == nil {
|
|
// If it exists, generate a unique name (e.g., by adding a number)
|
|
newFilename = makeUniqueFilename(newFilename)
|
|
}
|
|
|
|
// Perform the rename operation
|
|
return os.Rename(oldFilename, newFilename)
|
|
}
|
|
|
|
// sanitizeFilename removes invalid characters from a filename and ensures the first character is a letter.
|
|
func sanitizeFilename(filename string) string {
|
|
// Extract the extension
|
|
ext := filepath.Ext(filename)
|
|
filenameWithoutExt := filename[:len(filename)-len(ext)]
|
|
|
|
// Define a regular expression to match invalid filename characters
|
|
invalidChars := regexp.MustCompile(`[\\/:*?"<>|]`)
|
|
|
|
// Replace invalid characters with underscores
|
|
sanitizedFilename := invalidChars.ReplaceAllString(filenameWithoutExt, "_")
|
|
|
|
// Remove non-letter characters from the beginning of the filename
|
|
for len(sanitizedFilename) > 0 && !unicode.IsLetter(rune(sanitizedFilename[0])) {
|
|
sanitizedFilename = sanitizedFilename[1:]
|
|
}
|
|
|
|
// Ensure the first character is a letter
|
|
if len(sanitizedFilename) == 0 || !unicode.IsLetter(rune(sanitizedFilename[0])) {
|
|
sanitizedFilename = "A_" + sanitizedFilename
|
|
}
|
|
|
|
// Remove leading dots and underscores
|
|
for len(sanitizedFilename) > 0 && (sanitizedFilename[0] == '.' || sanitizedFilename[0] == '_') {
|
|
sanitizedFilename = sanitizedFilename[1:]
|
|
}
|
|
|
|
// If the filename is now empty (e.g., it was only dots or underscores), set it to "A"
|
|
if len(sanitizedFilename) == 0 {
|
|
sanitizedFilename = "A"
|
|
}
|
|
|
|
return sanitizedFilename + ext
|
|
}
|
|
|
|
// makeUniqueFilename adds a number to the filename to make it unique.
|
|
func makeUniqueFilename(filename string) string {
|
|
ext := filepath.Ext(filename)
|
|
name := filename[:len(filename)-len(ext)]
|
|
|
|
counter := 1
|
|
newFilename := filename
|
|
for {
|
|
if _, err := os.Stat(newFilename); os.IsNotExist(err) {
|
|
return newFilename
|
|
}
|
|
newFilename = fmt.Sprintf("%s_%d%s", name, counter, ext)
|
|
counter++
|
|
}
|
|
}
|