gonamer/gonamer.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++
}
}