cclohmar
5ca3ff7555
- 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
102 lines
2.7 KiB
Go
102 lines
2.7 KiB
Go
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
|
|
}
|