feat: configurable AI provider system (Gemini, OpenAI, Ollama)
- New provider architecture with common interface - Provider selected via AI_PROVIDER env var (gemini/openai/ollama) - Gemini (default): existing implementation, uses GEMINI_API_KEY - OpenAI-compatible: uses OPENAI_API_KEY + AI_MODEL + AI_BASE_URL - Works with OpenAI, Perplexity, Together AI, Groq, etc. - Ollama: local LLM, uses AI_BASE_URL + AI_MODEL - Supports llava, bakllava, and other vision models - deepseek.go renamed to llm.go (cleanup) - .env.example updated with all AI provider options
This commit is contained in:
parent
646378df64
commit
5ca3ff7555
7 changed files with 514 additions and 266 deletions
17
.env.example
17
.env.example
|
|
@ -7,8 +7,21 @@ SMTP_PORT=587
|
|||
SMTP_USER=post@2-4-h.app
|
||||
SMTP_PASS=D9AW8JP74r1V
|
||||
|
||||
# DeepSeek Vision API Key
|
||||
DEEPSEEK_API_KEY=sk-e9362165d2694883a52a5142811aa422
|
||||
# --- AI Provider Configuration ---
|
||||
# Choose one: gemini (default), openai, ollama
|
||||
AI_PROVIDER=gemini
|
||||
|
||||
# Gemini (default, used when AI_PROVIDER=gemini)
|
||||
GEMINI_API_KEY=AQ.Ab8RN6IQjKTQofuKOW2TT5mZ0zwt8rFa8X3SHGyYyce4DrbBJw
|
||||
|
||||
# OpenAI / Compatible (used when AI_PROVIDER=openai)
|
||||
# OPENAI_API_KEY=sk-...
|
||||
# AI_MODEL=gpt-4o-mini
|
||||
# AI_BASE_URL=https://api.openai.com/v1
|
||||
|
||||
# Ollama - Local LLM (used when AI_PROVIDER=ollama)
|
||||
# AI_BASE_URL=http://localhost:11434
|
||||
# AI_MODEL=llava
|
||||
|
||||
# Base URL for generating absolute links in emails
|
||||
BASE_URL=http://localhost:8080
|
||||
|
|
|
|||
|
|
@ -1,264 +0,0 @@
|
|||
package ai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ReceiptData represents the structured data extracted from a receipt image.
|
||||
type ReceiptData struct {
|
||||
Amount float64 `json:"amount"`
|
||||
Currency string `json:"currency"`
|
||||
Merchant string `json:"merchant"`
|
||||
Category string `json:"category"`
|
||||
Date string `json:"date"`
|
||||
}
|
||||
|
||||
const (
|
||||
geminiAPIURL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-flash-lite:generateContent"
|
||||
requestTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
// --- Gemini API types ---
|
||||
|
||||
type geminiRequest struct {
|
||||
Contents []geminiContent `json:"contents"`
|
||||
}
|
||||
|
||||
type geminiContent struct {
|
||||
Parts []geminiPart `json:"parts"`
|
||||
}
|
||||
|
||||
type geminiPart struct {
|
||||
Text string `json:"text,omitempty"`
|
||||
InlineData *geminiFileData `json:"inline_data,omitempty"`
|
||||
}
|
||||
|
||||
type geminiFileData struct {
|
||||
MimeType string `json:"mime_type"`
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
type geminiResponse struct {
|
||||
Candidates []geminiCandidate `json:"candidates"`
|
||||
Error *geminiError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type geminiCandidate struct {
|
||||
Content geminiResponseContent `json:"content"`
|
||||
}
|
||||
|
||||
type geminiResponseContent struct {
|
||||
Parts []geminiResponsePart `json:"parts"`
|
||||
}
|
||||
|
||||
type geminiResponsePart struct {
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
type geminiError struct {
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
// ExtractReceipt sends a receipt image to Google Gemini Vision API and
|
||||
// returns structured receipt data extracted from the image.
|
||||
//
|
||||
// Gemini supports native vision processing via inline_data, so the image
|
||||
// is sent directly as base64-encoded data with its MIME type.
|
||||
func ExtractReceipt(imagePath string) (*ReceiptData, error) {
|
||||
// 1. Read the image file.
|
||||
imageData, err := readImageFile(imagePath)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, err
|
||||
}
|
||||
|
||||
// 2. Detect MIME type from magic bytes.
|
||||
mimeType := detectMimeType(imageData)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/jpeg" // safe default
|
||||
}
|
||||
|
||||
// 3. Get the API key.
|
||||
apiKey := os.Getenv("GEMINI_API_KEY")
|
||||
if apiKey == "" {
|
||||
// Fall back to the key from .env.example.
|
||||
apiKey = "AQ.Ab8RN6IQjKTQofuKOW2TT5mZ0zwt8rFa8X3SHGyYyce4DrbBJw"
|
||||
}
|
||||
|
||||
// 4. Build the Gemini request with the image as inline_data.
|
||||
b64Data := base64.StdEncoding.EncodeToString(imageData)
|
||||
|
||||
payload := geminiRequest{
|
||||
Contents: []geminiContent{
|
||||
{
|
||||
Parts: []geminiPart{
|
||||
{
|
||||
Text: "Analyze this receipt image. Extract the following fields as a strict JSON object with these exact keys: \"merchant\" (string, store or business name), \"amount\" (number, total paid), \"currency\" (string, 3-letter code like KES, USD, EUR), \"category\" (string, one of: Food, Travel, Lodging, Software, Other), \"date\" (string, YYYY-MM-DD format). Return ONLY valid JSON. No markdown, no explanation, no code fences.",
|
||||
},
|
||||
{
|
||||
InlineData: &geminiFileData{
|
||||
MimeType: mimeType,
|
||||
Data: b64Data,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("marshal request: %w", err)
|
||||
}
|
||||
|
||||
// 5. Send to Gemini API.
|
||||
req, err := http.NewRequest(http.MethodPost, geminiAPIURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("X-goog-api-key", apiKey)
|
||||
|
||||
client := &http.Client{Timeout: requestTimeout}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("API request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("read response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return &ReceiptData{}, fmt.Errorf("Gemini API returned status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
// 6. Parse the response.
|
||||
var apiResp geminiResponse
|
||||
if err := json.Unmarshal(respBody, &apiResp); err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("parse response: %w", err)
|
||||
}
|
||||
|
||||
if apiResp.Error != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("Gemini API error: %s", apiResp.Error.Message)
|
||||
}
|
||||
|
||||
if len(apiResp.Candidates) == 0 {
|
||||
return &ReceiptData{}, errors.New("Gemini returned no candidates")
|
||||
}
|
||||
|
||||
parts := apiResp.Candidates[0].Content.Parts
|
||||
if len(parts) == 0 {
|
||||
return &ReceiptData{}, errors.New("Gemini response has no parts")
|
||||
}
|
||||
|
||||
contentStr := strings.TrimSpace(parts[0].Text)
|
||||
contentStr = stripMarkdownFences(contentStr)
|
||||
|
||||
// 7. Parse the JSON into ReceiptData.
|
||||
var receipt ReceiptData
|
||||
if err := json.Unmarshal([]byte(contentStr), &receipt); err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("parse receipt JSON: %w (content: %s)", err, contentStr)
|
||||
}
|
||||
|
||||
log.Printf("ExtractReceipt: merchant=%q amount=%.2f %s category=%q date=%q",
|
||||
receipt.Merchant, receipt.Amount, receipt.Currency, receipt.Category, receipt.Date)
|
||||
return &receipt, nil
|
||||
}
|
||||
|
||||
// detectMimeType determines the MIME type of an image from its magic bytes.
|
||||
func detectMimeType(data []byte) string {
|
||||
if len(data) < 4 {
|
||||
return ""
|
||||
}
|
||||
// JPEG
|
||||
if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
|
||||
return "image/jpeg"
|
||||
}
|
||||
// PNG
|
||||
if data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4E && data[3] == 0x47 {
|
||||
return "image/png"
|
||||
}
|
||||
// WebP
|
||||
if len(data) >= 12 && data[0] == 0x52 && data[1] == 0x49 && data[2] == 0x46 &&
|
||||
data[3] == 0x46 && data[8] == 0x57 && data[9] == 0x45 && data[10] == 0x42 && data[11] == 0x50 {
|
||||
return "image/webp"
|
||||
}
|
||||
// GIF
|
||||
if data[0] == 0x47 && data[1] == 0x49 && data[2] == 0x46 {
|
||||
return "image/gif"
|
||||
}
|
||||
// BMP
|
||||
if data[0] == 0x42 && data[1] == 0x4D {
|
||||
return "image/bmp"
|
||||
}
|
||||
// TIFF
|
||||
if (data[0] == 0x49 && data[1] == 0x49 && data[2] == 0x2A && data[3] == 0x00) ||
|
||||
(data[0] == 0x4D && data[1] == 0x4D && data[2] == 0x00 && data[3] == 0x2A) {
|
||||
return "image/tiff"
|
||||
}
|
||||
// PDF: 25 50 44 46 (%PDF)
|
||||
if len(data) >= 4 && data[0] == 0x25 && data[1] == 0x50 && data[2] == 0x44 && data[3] == 0x46 {
|
||||
return "application/pdf"
|
||||
}
|
||||
|
||||
// HEIC/HEIF (ftyp box at offset 4)
|
||||
if len(data) >= 12 && data[4] == 0x66 && data[5] == 0x74 && data[6] == 0x79 && data[7] == 0x70 {
|
||||
brand := string(data[8:12])
|
||||
switch brand {
|
||||
case "heic", "heix", "hevc", "hevx", "mif1", "msf1":
|
||||
return "image/heic"
|
||||
case "avif":
|
||||
return "image/avif"
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// stripMarkdownFences removes markdown code fences from model output.
|
||||
func stripMarkdownFences(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if strings.HasPrefix(s, "```") {
|
||||
s = s[3:]
|
||||
if idx := strings.Index(s, "\n"); idx != -1 {
|
||||
s = s[idx+1:]
|
||||
}
|
||||
}
|
||||
if strings.HasSuffix(s, "```") {
|
||||
s = s[:len(s)-3]
|
||||
}
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
|
||||
// readImageFile reads the full contents of an image file from disk.
|
||||
func readImageFile(path string) ([]byte, error) {
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, fmt.Errorf("image file does not exist: %s", path)
|
||||
}
|
||||
return nil, fmt.Errorf("cannot stat image file %s: %w", path, err)
|
||||
}
|
||||
if info.IsDir() {
|
||||
return nil, fmt.Errorf("path is a directory: %s", path)
|
||||
}
|
||||
if info.Size() > 10<<20 {
|
||||
return nil, fmt.Errorf("image too large: %d bytes (max 10MB)", info.Size())
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot read image file %s: %w", path, err)
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
126
internal/ai/gemini.go
Normal file
126
internal/ai/gemini.go
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
package ai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
geminiAPIURL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-flash-lite:generateContent"
|
||||
geminiTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
type geminiRequest struct {
|
||||
Contents []geminiContent `json:"contents"`
|
||||
}
|
||||
|
||||
type geminiContent struct {
|
||||
Parts []geminiPart `json:"parts"`
|
||||
}
|
||||
|
||||
type geminiPart struct {
|
||||
Text string `json:"text,omitempty"`
|
||||
InlineData *geminiFileData `json:"inline_data,omitempty"`
|
||||
}
|
||||
|
||||
type geminiFileData struct {
|
||||
MimeType string `json:"mime_type"`
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
type geminiResponse struct {
|
||||
Candidates []geminiCandidate `json:"candidates"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type geminiCandidate struct {
|
||||
Content geminiResponseContent `json:"content"`
|
||||
}
|
||||
|
||||
type geminiResponseContent struct {
|
||||
Parts []struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"parts"`
|
||||
}
|
||||
|
||||
type geminiProvider struct{}
|
||||
|
||||
func (p *geminiProvider) ExtractReceipt(imagePath string) (*ReceiptData, error) {
|
||||
imageData, err := readFile(imagePath)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("read file: %w", err)
|
||||
}
|
||||
|
||||
mimeType := detectMimeType(imageData)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/jpeg"
|
||||
}
|
||||
|
||||
apiKey := os.Getenv("GEMINI_API_KEY")
|
||||
if apiKey == "" {
|
||||
apiKey = "AQ.Ab8RN6IQjKTQofuKOW2TT5mZ0zwt8rFa8X3SHGyYyce4DrbBJw"
|
||||
}
|
||||
|
||||
b64Data := base64.StdEncoding.EncodeToString(imageData)
|
||||
|
||||
payload := geminiRequest{
|
||||
Contents: []geminiContent{{
|
||||
Parts: []geminiPart{
|
||||
{Text: "Analyze this receipt. Extract as strict JSON with keys: \"merchant\" (string), \"amount\" (number), \"currency\" (3-letter code), \"category\" (Food/Travel/Lodging/Software/Other), \"date\" (YYYY-MM-DD). Return ONLY valid JSON. No markdown."},
|
||||
{InlineData: &geminiFileData{MimeType: mimeType, Data: b64Data}},
|
||||
},
|
||||
}},
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(payload)
|
||||
req, _ := http.NewRequest(http.MethodPost, geminiAPIURL, bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("X-goog-api-key", apiKey)
|
||||
|
||||
client := &http.Client{Timeout: geminiTimeout}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("API request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return &ReceiptData{}, fmt.Errorf("Gemini status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var apiResp geminiResponse
|
||||
json.Unmarshal(respBody, &apiResp)
|
||||
|
||||
if apiResp.Error != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("Gemini error: %s", apiResp.Error.Message)
|
||||
}
|
||||
if len(apiResp.Candidates) == 0 {
|
||||
return &ReceiptData{}, errors.New("no candidates")
|
||||
}
|
||||
|
||||
parts := apiResp.Candidates[0].Content.Parts
|
||||
if len(parts) == 0 {
|
||||
return &ReceiptData{}, errors.New("no response text")
|
||||
}
|
||||
|
||||
contentStr := stripMarkdownFences(parts[0].Text)
|
||||
var receipt ReceiptData
|
||||
if err := json.Unmarshal([]byte(contentStr), &receipt); err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("parse JSON: %w (content: %s)", err, contentStr)
|
||||
}
|
||||
|
||||
log.Printf("ExtractReceipt [gemini]: merchant=%q amount=%.2f %s category=%q date=%q",
|
||||
receipt.Merchant, receipt.Amount, receipt.Currency, receipt.Category, receipt.Date)
|
||||
return &receipt, nil
|
||||
}
|
||||
6
internal/ai/llm.go
Normal file
6
internal/ai/llm.go
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
// This file intentionally left blank.
|
||||
// Gemini provider moved to gemini.go.
|
||||
// Shared types and provider factory are in receipt.go.
|
||||
// OpenAI provider is in openai.go.
|
||||
// Ollama provider is in ollama.go.
|
||||
package ai
|
||||
102
internal/ai/ollama.go
Normal file
102
internal/ai/ollama.go
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
package ai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const ollamaTimeout = 60 * time.Second
|
||||
|
||||
type ollamaRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []ollamaMessage `json:"messages"`
|
||||
Stream bool `json:"stream"`
|
||||
Options map[string]any `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
type ollamaMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
Images []string `json:"images,omitempty"`
|
||||
}
|
||||
|
||||
type ollamaResponse struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type ollamaProvider struct{}
|
||||
|
||||
func (p *ollamaProvider) ExtractReceipt(imagePath string) (*ReceiptData, error) {
|
||||
imageData, err := readFile(imagePath)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("read file: %w", err)
|
||||
}
|
||||
|
||||
baseURL := os.Getenv("AI_BASE_URL")
|
||||
if baseURL == "" {
|
||||
baseURL = "http://localhost:11434"
|
||||
}
|
||||
baseURL = strings.TrimRight(baseURL, "/")
|
||||
|
||||
model := os.Getenv("AI_MODEL")
|
||||
if model == "" {
|
||||
model = "llava"
|
||||
}
|
||||
|
||||
b64Data := base64.StdEncoding.EncodeToString(imageData)
|
||||
|
||||
payload := ollamaRequest{
|
||||
Model: model,
|
||||
Stream: false,
|
||||
Messages: []ollamaMessage{{
|
||||
Role: "user",
|
||||
Content: "Analyze this receipt image. Extract the following fields as a strict JSON object with these exact keys: \"merchant\" (string), \"amount\" (number), \"currency\" (string, 3-letter code), \"category\" (string, one of: Food, Travel, Lodging, Software, Other), \"date\" (string, YYYY-MM-DD). Return ONLY valid JSON. No markdown.",
|
||||
Images: []string{b64Data},
|
||||
}},
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(payload)
|
||||
apiURL := baseURL + "/api/chat"
|
||||
req, _ := http.NewRequest(http.MethodPost, apiURL, bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := &http.Client{Timeout: ollamaTimeout}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("Ollama request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return &ReceiptData{}, fmt.Errorf("Ollama status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var apiResp ollamaResponse
|
||||
json.Unmarshal(respBody, &apiResp)
|
||||
|
||||
if apiResp.Error != "" {
|
||||
return &ReceiptData{}, fmt.Errorf("Ollama error: %s", apiResp.Error)
|
||||
}
|
||||
|
||||
contentStr := stripMarkdownFences(apiResp.Message.Content)
|
||||
var receipt ReceiptData
|
||||
if err := json.Unmarshal([]byte(contentStr), &receipt); err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("parse JSON: %w (content: %s)", err, contentStr)
|
||||
}
|
||||
|
||||
log.Printf("ExtractReceipt [ollama-%s]: merchant=%q amount=%.2f %s",
|
||||
model, receipt.Merchant, receipt.Amount, receipt.Currency)
|
||||
return &receipt, nil
|
||||
}
|
||||
132
internal/ai/openai.go
Normal file
132
internal/ai/openai.go
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
package ai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const openaiTimeout = 30 * time.Second
|
||||
|
||||
type openaiRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []openaiMessage `json:"messages"`
|
||||
Temperature float64 `json:"temperature"`
|
||||
}
|
||||
|
||||
type openaiMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content []openaiContent `json:"content"`
|
||||
}
|
||||
|
||||
type openaiContent struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"`
|
||||
ImageURL *openaiImage `json:"image_url,omitempty"`
|
||||
}
|
||||
|
||||
type openaiImage struct {
|
||||
URL string `json:"url"`
|
||||
}
|
||||
|
||||
type openaiResponse struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type openaiProvider struct{}
|
||||
|
||||
func (p *openaiProvider) ExtractReceipt(imagePath string) (*ReceiptData, error) {
|
||||
imageData, err := readFile(imagePath)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("read file: %w", err)
|
||||
}
|
||||
|
||||
mimeType := detectMimeType(imageData)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/jpeg"
|
||||
}
|
||||
|
||||
apiKey := os.Getenv("OPENAI_API_KEY")
|
||||
if apiKey == "" {
|
||||
return &ReceiptData{}, errors.New("OPENAI_API_KEY environment variable not set")
|
||||
}
|
||||
|
||||
baseURL := os.Getenv("AI_BASE_URL")
|
||||
if baseURL == "" {
|
||||
baseURL = "https://api.openai.com/v1"
|
||||
}
|
||||
baseURL = strings.TrimRight(baseURL, "/")
|
||||
|
||||
model := os.Getenv("AI_MODEL")
|
||||
if model == "" {
|
||||
model = "gpt-4o-mini"
|
||||
}
|
||||
|
||||
b64Data := base64.StdEncoding.EncodeToString(imageData)
|
||||
dataURL := fmt.Sprintf("data:%s;base64,%s", mimeType, b64Data)
|
||||
|
||||
payload := openaiRequest{
|
||||
Model: model,
|
||||
Temperature: 0.1,
|
||||
Messages: []openaiMessage{{
|
||||
Role: "user",
|
||||
Content: []openaiContent{
|
||||
{Type: "text", Text: "Analyze this receipt image. Extract the following fields as a strict JSON object with these exact keys: \"merchant\" (string, store or business name), \"amount\" (number, total paid), \"currency\" (string, 3-letter code like KES, USD, EUR), \"category\" (string, one of: Food, Travel, Lodging, Software, Other), \"date\" (string, YYYY-MM-DD format). Return ONLY valid JSON. No markdown, no explanation, no code fences."},
|
||||
{Type: "image_url", ImageURL: &openaiImage{URL: dataURL}},
|
||||
},
|
||||
}},
|
||||
}
|
||||
|
||||
body, _ := json.Marshal(payload)
|
||||
apiURL := baseURL + "/chat/completions"
|
||||
req, _ := http.NewRequest(http.MethodPost, apiURL, bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
|
||||
client := &http.Client{Timeout: openaiTimeout}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("API request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return &ReceiptData{}, fmt.Errorf("API status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var apiResp openaiResponse
|
||||
json.Unmarshal(respBody, &apiResp)
|
||||
|
||||
if apiResp.Error != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("API error: %s", apiResp.Error.Message)
|
||||
}
|
||||
if len(apiResp.Choices) == 0 {
|
||||
return &ReceiptData{}, errors.New("no response choices")
|
||||
}
|
||||
|
||||
contentStr := stripMarkdownFences(apiResp.Choices[0].Message.Content)
|
||||
var receipt ReceiptData
|
||||
if err := json.Unmarshal([]byte(contentStr), &receipt); err != nil {
|
||||
return &ReceiptData{}, fmt.Errorf("parse JSON: %w (content: %s)", err, contentStr)
|
||||
}
|
||||
|
||||
log.Printf("ExtractReceipt [openai-%s]: merchant=%q amount=%.2f %s",
|
||||
model, receipt.Merchant, receipt.Amount, receipt.Currency)
|
||||
return &receipt, nil
|
||||
}
|
||||
133
internal/ai/receipt.go
Normal file
133
internal/ai/receipt.go
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
// Package ai provides receipt data extraction from images/PDFs using
|
||||
// configurable AI providers (Gemini, OpenAI-compatible, or Ollama).
|
||||
//
|
||||
// Provider selection is done via environment variables:
|
||||
//
|
||||
// AI_PROVIDER=gemini (default, uses GEMINI_API_KEY)
|
||||
// AI_PROVIDER=openai (uses OPENAI_API_KEY, AI_MODEL, AI_BASE_URL)
|
||||
// AI_PROVIDER=ollama (uses AI_BASE_URL, AI_MODEL)
|
||||
package ai
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ReceiptData represents the structured data extracted from a receipt image.
|
||||
type ReceiptData struct {
|
||||
Amount float64 `json:"amount"`
|
||||
Currency string `json:"currency"`
|
||||
Merchant string `json:"merchant"`
|
||||
Category string `json:"category"`
|
||||
Date string `json:"date"`
|
||||
}
|
||||
|
||||
// Provider is the interface that wraps receipt extraction.
|
||||
// Each provider (Gemini, OpenAI, Ollama) implements this interface.
|
||||
type Provider interface {
|
||||
ExtractReceipt(imagePath string) (*ReceiptData, error)
|
||||
}
|
||||
|
||||
// ExtractReceipt dispatches to the configured AI provider.
|
||||
// The provider is selected based on the AI_PROVIDER environment variable.
|
||||
func ExtractReceipt(imagePath string) (*ReceiptData, error) {
|
||||
provider := getProvider()
|
||||
return provider.ExtractReceipt(imagePath)
|
||||
}
|
||||
|
||||
// getProvider returns the appropriate Provider based on environment config.
|
||||
func getProvider() Provider {
|
||||
providerName := strings.ToLower(strings.TrimSpace(os.Getenv("AI_PROVIDER")))
|
||||
|
||||
switch providerName {
|
||||
case "openai":
|
||||
return &openaiProvider{}
|
||||
case "ollama":
|
||||
return &ollamaProvider{}
|
||||
case "gemini":
|
||||
fallthrough
|
||||
default:
|
||||
return &geminiProvider{}
|
||||
}
|
||||
}
|
||||
|
||||
// stripMarkdownFences removes markdown code fences from model output.
|
||||
func stripMarkdownFences(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if strings.HasPrefix(s, "```") {
|
||||
s = s[3:]
|
||||
if idx := strings.Index(s, "\n"); idx != -1 {
|
||||
s = s[idx+1:]
|
||||
}
|
||||
}
|
||||
if strings.HasSuffix(s, "```") {
|
||||
s = s[:len(s)-3]
|
||||
}
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
|
||||
// readFile reads the full contents of a file from disk.
|
||||
func readFile(path string) ([]byte, error) {
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, err
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if info.IsDir() {
|
||||
return nil, err
|
||||
}
|
||||
if info.Size() > 10<<20 {
|
||||
return nil, err
|
||||
}
|
||||
return os.ReadFile(path)
|
||||
}
|
||||
|
||||
// detectMimeType determines the MIME type from magic bytes.
|
||||
func detectMimeType(data []byte) string {
|
||||
if len(data) < 4 {
|
||||
return ""
|
||||
}
|
||||
// JPEG
|
||||
if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
|
||||
return "image/jpeg"
|
||||
}
|
||||
// PNG
|
||||
if data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4E && data[3] == 0x47 {
|
||||
return "image/png"
|
||||
}
|
||||
// WebP
|
||||
if len(data) >= 12 && data[0] == 0x52 && data[1] == 0x49 && data[2] == 0x46 &&
|
||||
data[3] == 0x46 && data[8] == 0x57 && data[9] == 0x45 && data[10] == 0x42 && data[11] == 0x50 {
|
||||
return "image/webp"
|
||||
}
|
||||
// GIF
|
||||
if data[0] == 0x47 && data[1] == 0x49 && data[2] == 0x46 {
|
||||
return "image/gif"
|
||||
}
|
||||
// BMP
|
||||
if data[0] == 0x42 && data[1] == 0x4D {
|
||||
return "image/bmp"
|
||||
}
|
||||
// TIFF
|
||||
if (data[0] == 0x49 && data[1] == 0x49 && data[2] == 0x2A && data[3] == 0x00) ||
|
||||
(data[0] == 0x4D && data[1] == 0x4D && data[2] == 0x00 && data[3] == 0x2A) {
|
||||
return "image/tiff"
|
||||
}
|
||||
// PDF
|
||||
if data[0] == 0x25 && data[1] == 0x50 && data[2] == 0x44 && data[3] == 0x46 {
|
||||
return "application/pdf"
|
||||
}
|
||||
// HEIC/HEIF (ftyp box at offset 4)
|
||||
if len(data) >= 12 && data[4] == 0x66 && data[5] == 0x74 && data[6] == 0x79 && data[7] == 0x70 {
|
||||
brand := string(data[8:12])
|
||||
switch brand {
|
||||
case "heic", "heix", "hevc", "hevx", "mif1", "msf1":
|
||||
return "image/heic"
|
||||
case "avif":
|
||||
return "image/avif"
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
Loading…
Reference in a new issue