Files
budgetTool/main.go

316 lines
7.8 KiB
Go

package main
import (
"bufio"
"context"
"encoding/csv"
"encoding/json"
"fmt"
"github.com/sashabaranov/go-openai"
"os"
"strconv"
"strings"
"time"
)
type Category struct {
Name string
Transactions []Transaction
Total float64
}
func main() {
// read csv line by line
//lines, err := readLinesFromFile("./Transactions-2023-05-03.csv")
//if err != nil {
// panic(err)
//}
//transactions, err := getJSONFromFile("./completed.json")
//if err != nil {
// panic(err)
//}
transactions, err := getCategories()
if err != nil {
panic(err)
}
fmt.Println(len(transactions))
fmt.Println(transactions[0])
categories := make(map[string]Category)
for _, transaction := range transactions {
category, ok := categories[transaction.Category]
if !ok {
category = Category{
Name: transaction.Category,
Transactions: []Transaction{},
Total: 0,
}
}
category.Name = transaction.Category
category.Transactions = append(category.Transactions, transaction)
// parse amount from string to float
amount, err := strconv.ParseFloat(transaction.Amount, 64)
if err != nil {
continue
}
category.Total += amount
categories[transaction.Category] = category
}
count := 0
for _, category := range categories {
fmt.Printf("%s: Count: %v Total: %.2f\n", category.Name, len(category.Transactions), category.Total)
count += 1
}
fmt.Println(count)
fmt.Println(len(transactions))
}
func getJSONFromFile(filePath string) ([]Transaction, error) {
file, err := os.Open(filePath)
if err != nil {
return []Transaction{}, fmt.Errorf("error opening file: %w", err)
}
defer file.Close()
var transactions []Transaction
err = json.NewDecoder(file).Decode(&transactions)
if err != nil {
return []Transaction{}, fmt.Errorf("error decoding json: %w", err)
}
return transactions, nil
}
func getCategories() ([]Transaction, error) {
transactions, err := readCSV("./Transactions-2023-05-03.csv")
if err != nil {
return []Transaction{}, fmt.Errorf("error reading csv: %w", err)
}
var finalTransactions []Transaction
count := 0
for _, line := range transactions {
//if count > 100 {
// break
//}
fmt.Println(line)
transaction, err := parseTransaction(line.TransactionName)
if err != nil {
// continue if parsing error
fmt.Println(err)
}
//fmt.Printf("%+v\n", transaction)
combinedTransaction := Transaction{
Category: transaction.Category,
SubCategory: transaction.SubCategory,
Amount: line.Amount,
TransactionName: line.TransactionName,
IsoDate: line.IsoDate,
}
finalTransactions = append(finalTransactions, combinedTransaction)
count++
}
jsonTransactions, err := json.Marshal(finalTransactions)
if err != nil {
return []Transaction{}, fmt.Errorf("error marshalling json: %w", err)
}
fmt.Println(string(jsonTransactions))
err = writeBytesToFile("./raw-categories.json", jsonTransactions)
if err != nil {
return []Transaction{}, fmt.Errorf("error writing to file: %w", err)
}
return finalTransactions, nil
}
func writeBytesToFile(filePath string, bytes []byte) error {
file, err := os.Create(filePath)
if err != nil {
return fmt.Errorf("error creating file: %w", err)
}
defer file.Close()
_, err = file.Write(bytes)
if err != nil {
return fmt.Errorf("error writing to file: %w", err)
}
return nil
}
func parseTransaction(rawTransaction string) (Transaction, error) {
// save for later
// category should be one of the following: "Income", "Food", "Shopping", "Transportation", "Housing", "Utilities", "Insurance", "Medical", "Saving and Investing", "Debt Payments", "Personal Spending", "Entertainment", "Other"
categories := []string{
"Donations",
"Banking/Finance",
"Mortgage",
"Credit Cards",
"Loans",
"Bills/Utilities",
"Food/Groceries",
"Pet Care",
"Housing",
"Cloud Services",
"Health and Beauty",
"Entertainment/Shopping",
"Transportation",
"Uncategorized",
}
client := openai.NewClient(os.Getenv("OPENAI_API_KEY"))
messages := []openai.ChatCompletionMessage{
{
Role: openai.ChatMessageRoleSystem,
Content: `
You are a helpful bot that categorizes transactions. You are given a transaction name raw from the bank and you need to categorize it. Please only respond in json format.
category should be one of the following:
` + strings.Join(categories, ",\n") + `
subCategory should provide more detail on the category. For example, if the category is "Food/Groceries", the subCategory could be "Groceries" or "Restaurants".
Expected output format:
{
"category": "",
"subCategory": "",
"transactionName": ""
}`,
},
{
Role: openai.ChatMessageRoleUser,
Content: rawTransaction,
},
}
hasCategory := false
for !hasCategory {
resp, err := client.CreateChatCompletion(
context.Background(),
openai.ChatCompletionRequest{
Model: openai.GPT3Dot5Turbo,
Messages: messages,
MaxTokens: 256,
Temperature: 0,
},
)
if err != nil {
return Transaction{}, fmt.Errorf("ChatCompletion request error: %w", err)
}
fmt.Println(resp.Choices[0].Message.Content)
transaction, err := parseAIOutput(resp.Choices[0].Message.Content)
if err != nil {
return Transaction{}, fmt.Errorf("ChatCompletion parsing error: %w, got response: %v", err, resp.Choices[0].Message.Content)
}
// check if category is in the categories slice
if contains(categories, transaction.Category) {
fmt.Println("Found category!")
hasCategory = true
return transaction, nil
} else if len(messages) < 5 {
fmt.Println("Did not find category, trying again...")
messages = append(messages, openai.ChatCompletionMessage{
Role: resp.Choices[0].Message.Role,
Content: resp.Choices[0].Message.Content,
Name: "",
})
messages = append(messages, openai.ChatCompletionMessage{
Role: openai.ChatMessageRoleSystem,
Content: "Please provide a valid category.",
Name: "",
})
} else {
break
}
}
return Transaction{}, fmt.Errorf("could not find category for transaction: %v\nMessages: %+v", rawTransaction, messages)
}
func contains(categories []string, category string) bool {
for _, c := range categories {
if c == category {
return true
}
}
return false
}
type Transaction struct {
Category string `json:"category"`
SubCategory string `json:"subCategory"`
Amount string `json:"amount"`
TransactionName string `json:"transactionName"`
IsoDate string `json:"isodate"`
}
func parseAIOutput(aiOutput string) (Transaction, error) {
var transaction Transaction
aiOutputParts := strings.Split(aiOutput, "{")
aiJson := "{" + aiOutputParts[len(aiOutputParts)-1]
err := json.Unmarshal([]byte(aiJson), &transaction)
if err != nil {
return Transaction{}, err
}
return transaction, nil
}
func readLinesFromFile(filename string) ([]string, error) {
file, err := os.Open(filename)
if err != nil {
return nil, err
}
defer file.Close()
var lines []string
scanner := bufio.NewScanner(file)
for scanner.Scan() {
lines = append(lines, scanner.Text())
}
if err := scanner.Err(); err != nil {
return nil, err
}
return lines, nil
}
func readCSV(filePath string) ([]Transaction, error) {
file, err := os.Open(filePath)
if err != nil {
return []Transaction{}, fmt.Errorf("error opening file: %w", err)
}
defer file.Close()
reader := csv.NewReader(file)
records, err := reader.ReadAll()
if err != nil {
return []Transaction{}, fmt.Errorf("error reading csv: %w", err)
}
var transactions []Transaction
// Do something with the records
for _, record := range records {
//fmt.Println(record)
isoDate, err := time.Parse("01/02/2006", record[0])
if err != nil {
fmt.Println(fmt.Errorf("error parsing date: %w", err))
continue
}
t := Transaction{
IsoDate: isoDate.Format("2006-01-02"),
TransactionName: record[2],
Amount: record[7],
}
transactions = append(transactions, t)
}
return transactions, nil
}