From 00d8f7c94a5a70ec08315ca338007b34ecdd45d7 Mon Sep 17 00:00:00 2001
From: James <james@jongsma.me>
Date: Sun, 15 Feb 2026 12:00:36 -0500
Subject: [PATCH] chore: auto-commit uncommitted changes

---
 ai.go | 103 +++++++++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 99 insertions(+), 4 deletions(-)

diff --git a/ai.go b/ai.go
index fd382ab..c20d32f 100644
--- a/ai.go
+++ b/ai.go
@@ -170,6 +170,7 @@ Respond in JSON ONLY:
 		"model":      "accounts/fireworks/models/kimi-k2p5",
 		"max_tokens": 4096,
 		"messages": []map[string]interface{}{
+			{"role": "system", "content": "You are a document analysis API. You MUST respond with raw JSON only. No markdown, no code fences, no explanation text. Start your response with { and end with }."},
 			{
 				"role": "user",
 				"content": []map[string]interface{}{
@@ -180,7 +181,28 @@ Respond in JSON ONLY:
 		},
 	}
 
-	return callFireworks(reqBody)
+	analysis, err := callFireworks(reqBody)
+	if err != nil {
+		// Retry once with a simpler prompt that's harder for the model to misinterpret
+		log.Printf("  [AI] First attempt failed, retrying with simplified prompt...")
+		retryBody := map[string]interface{}{
+			"model":      "accounts/fireworks/models/kimi-k2p5",
+			"max_tokens": 4096,
+			"messages": []map[string]interface{}{
+				{"role": "system", "content": "Output valid JSON only. No other text."},
+				{
+					"role": "user",
+					"content": []map[string]interface{}{
+						{"type": "image_url", "image_url": map[string]string{"url": "data:image/png;base64," + b64}},
+						{"type": "text", "text": `Look at this document. Return ONLY this JSON (fill in values):
+{"category":"uncategorized","doc_type":"unknown","date":"","vendor":"","amount":"","title":"Short Title Here","summary":"One sentence.","full_text":"All visible text here"}`},
+					},
+				},
+			},
+		}
+		return callFireworks(retryBody)
+	}
+	return analysis, nil
 }
 
 // AnalyzeText uses K2 text model for plain text files
@@ -241,14 +263,17 @@ func callFireworks(reqBody map[string]interface{}) (*DocumentAnalysis, error) {
 		return nil, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body))
 	}
 
+	respBody, _ := io.ReadAll(resp.Body)
+
 	var result struct {
 		Choices []struct {
 			Message struct {
-				Content string `json:"content"`
+				Content          string `json:"content"`
+				ReasoningContent string `json:"reasoning_content"`
 			} `json:"message"`
 		} `json:"choices"`
 	}
-	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+	if err := json.Unmarshal(respBody, &result); err != nil {
 		return nil, err
 	}
 
@@ -257,6 +282,28 @@ func callFireworks(reqBody map[string]interface{}) (*DocumentAnalysis, error) {
 	}
 
 	content := result.Choices[0].Message.Content
+	reasoning := result.Choices[0].Message.ReasoningContent
+
+	// K2.5 reasoning mode: actual JSON may be in content or reasoning_content
+	// Try content first, if it doesn't look like JSON, try reasoning_content
+	if !strings.Contains(content, "{") && reasoning != "" && strings.Contains(reasoning, "{") {
+		log.Printf("  [AI] Using reasoning_content (content had no JSON)")
+		content = reasoning
+	}
+
+	// Strip markdown code fences (```json ... ``` or ``` ... ```)
+	content = strings.TrimSpace(content)
+	if strings.HasPrefix(content, "```") {
+		// Remove opening fence (```json or ```)
+		if idx := strings.Index(content, "\n"); idx >= 0 {
+			content = content[idx+1:]
+		}
+		// Remove closing fence
+		if idx := strings.LastIndex(content, "```"); idx >= 0 {
+			content = content[:idx]
+		}
+		content = strings.TrimSpace(content)
+	}
 
 	// Extract JSON from response
 	if idx := strings.Index(content, "{"); idx >= 0 {
@@ -267,7 +314,17 @@ func callFireworks(reqBody map[string]interface{}) (*DocumentAnalysis, error) {
 
 	var analysis DocumentAnalysis
 	if err := json.Unmarshal([]byte(content), &analysis); err != nil {
-		return nil, fmt.Errorf("failed to parse response: %w", err)
+		// Last resort: try to find a JSON object with braces matching
+		cleaned := extractJSONObject(content)
+		if cleaned != "" {
+			if err2 := json.Unmarshal([]byte(cleaned), &analysis); err2 != nil {
+				log.Printf("  [AI debug] Failed to parse even after cleanup. Content starts: %.200s", content)
+				return nil, fmt.Errorf("failed to parse response: %w", err)
+			}
+		} else {
+			log.Printf("  [AI debug] No JSON object found in response. Content starts: %.200s", content)
+			return nil, fmt.Errorf("failed to parse response: %w", err)
+		}
 	}
 
 	// Validate category
@@ -279,6 +336,44 @@ func callFireworks(reqBody map[string]interface{}) (*DocumentAnalysis, error) {
 	return &analysis, nil
 }
 
+// extractJSONObject tries to find a balanced JSON object in a string
+func extractJSONObject(s string) string {
+	start := strings.Index(s, "{")
+	if start < 0 {
+		return ""
+	}
+	depth := 0
+	inString := false
+	escaped := false
+	for i := start; i < len(s); i++ {
+		c := s[i]
+		if escaped {
+			escaped = false
+			continue
+		}
+		if c == '\\' && inString {
+			escaped = true
+			continue
+		}
+		if c == '"' {
+			inString = !inString
+			continue
+		}
+		if inString {
+			continue
+		}
+		if c == '{' {
+			depth++
+		} else if c == '}' {
+			depth--
+			if depth == 0 {
+				return s[start : i+1]
+			}
+		}
+	}
+	return ""
+}
+
 // GenerateEmbedding creates a vector embedding using Fireworks
 func GenerateEmbedding(text string) ([]float32, error) {
 	if fireworksAPIKey == "" {