inou/api/api_llm.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"inou/lib"
)

// --- Local Structs for Prompt Processing ---
// These are defined in api/llm_types.go and should be used from there.
// They are commented out here to prevent redeclaration.

/*
type TriageResponse struct { ... }
type ExtractionResult struct { ... }
type InputConfig struct { ... }
type FormGroup struct { ... }
type FormField struct { ... }
type ScheduleSlot struct { ... }
type EntryData struct { ... }
var ValidCategories = map[string]bool{ ... }
*/


// --- API-Specific Logic ---

func loadLLMConfig() {
	// Load GeminiKey from file or environment
	data, err := os.ReadFile("anthropic.env")
	if err != nil {
		log.Printf("Warning: anthropic.env not found. Looking for GEMINI_API_KEY in environment.")
	}
	for _, line := range strings.Split(string(data), "\n") {
		parts := strings.SplitN(line, "=", 2)
		if len(parts) == 2 && parts[0] == "GEMINI_API_KEY" {
			lib.GeminiKey = strings.TrimSpace(parts[1])
		}
	}
	if lib.GeminiKey == "" {
		lib.GeminiKey = os.Getenv("GEMINI_API_KEY")
	}
	if lib.GeminiKey != "" {
		log.Println("Gemini API key loaded.")
	} else {
		log.Println("Warning: Gemini API key not found.")
	}

	// Initialize prompts directory
	exe, _ := os.Executable()
	promptsDir := filepath.Join(filepath.Dir(exe), "..", "api", "prompts")
	if _, err := os.Stat(promptsDir); os.IsNotExist(err) {
		promptsDir = "prompts" // Dev fallback
	}
	lib.InitPrompts(promptsDir)
	log.Printf("Prompts directory set to: %s", lib.PromptsDir())
}

// callLLMForPrompt is the main entry point for turning user text into a structured prompt.
func callLLMForPrompt(userInput string, dossierID string) (*ExtractionResult, error) {
	triage, err := runTriage(userInput)
	if err != nil {
		return nil, err
	}
	if triage.Error != "" {
		return &ExtractionResult{Error: triage.Error}, nil
	}

	existingTypes := getExistingPromptTypes(dossierID) // Assuming db is accessible in api/main
	return runExtraction(userInput, triage.Category, triage.Language, existingTypes)
}


// --- Local Prompt Handling & DB Functions ---

func loadPrompt(name string) (string, error) {
	path := filepath.Join(lib.PromptsDir(), name+".md")
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func runTriage(userInput string) (*TriageResponse, error) {
	tmpl, err := loadPrompt("triage")
	if err != nil {
		return nil, fmt.Errorf("failed to load triage prompt: %v", err)
	}
	prompt := strings.ReplaceAll(tmpl, "{{INPUT}}", userInput)
	
	respText, err := lib.CallGemini(prompt)
	if err != nil {
		return nil, err
	}

	var result TriageResponse
	if err := json.Unmarshal([]byte(respText), &result); err != nil {
		var errMap map[string]string
		if json.Unmarshal([]byte(respText), &errMap) == nil {
			if errMsg, ok := errMap["error"]; ok {
				result.Error = errMsg
				return &result, nil
			}
		}
		return nil, fmt.Errorf("failed to parse triage JSON: %v (raw: %s)", err, respText)
	}
	
	if _, ok := ValidCategories[result.Category]; !ok && result.Error == "" {
		result.Category = "note"
	}
	return &result, nil
}

func runExtraction(userInput, category, language string, existingTypes map[string][]string) (*ExtractionResult, error) {
	tmpl, err := loadPrompt(category)
	if err != nil {
		tmpl, err = loadPrompt("default")
		if err != nil {
			return nil, fmt.Errorf("failed to load prompt: %v", err)
		}
	}

	var existingStr string
	for cat, types := range existingTypes {
		if len(types) > 0 {
			existingStr += fmt.Sprintf("- %s: %v\n", cat, types)
		}
	}
	if existingStr == "" {
		existingStr = "(none yet)"
	}

	prompt := tmpl
	prompt = strings.ReplaceAll(prompt, "{{INPUT}}", userInput)
	prompt = strings.ReplaceAll(prompt, "{{LANGUAGE}}", language)
	prompt = strings.ReplaceAll(prompt, "{{CATEGORY}}", category)
	prompt = strings.ReplaceAll(prompt, "{{EXISTING_TYPES}}", existingStr)
	
	respText, err := lib.CallGemini(prompt)
	if err != nil {
		return nil, err
	}

	var result ExtractionResult
	if err := json.Unmarshal([]byte(respText), &result); err != nil {
		// Fallback for single entry to maintain compatibility with older prompts
		var singleEntryResult struct {
			Question    string      `json:"question"`
			Category    string      `json:"category"`
			Type        string      `json:"type"`
			InputType   string      `json:"input_type"`
			InputConfig InputConfig `json:"input_config"`
			Schedule    []ScheduleSlot `json:"schedule"`
			Entry       *EntryData     `json:"entry,omitempty"`
			Error       string      `json:"error,omitempty"`
		}
		if err2 := json.Unmarshal([]byte(respText), &singleEntryResult); err2 == nil {
			result.Question = singleEntryResult.Question
			result.Category = singleEntryResult.Category
			result.Type = singleEntryResult.Type
			result.InputType = singleEntryResult.InputType
			result.InputConfig = singleEntryResult.InputConfig
			result.Schedule = singleEntryResult.Schedule
			result.Error = singleEntryResult.Error
			if singleEntryResult.Entry != nil {
				result.Entries = []*EntryData{singleEntryResult.Entry}
			}
			if result.Category == "" {
				result.Category = category
			}
			return &result, nil
		}
		return nil, fmt.Errorf("failed to parse extraction JSON: %v (raw: %s)", err, respText)
	}

	if result.Category == "" {
		result.Category = category
	}
	return &result, nil
}


func getExistingPromptTypes(dossierID string) map[string][]string {
	result, err := lib.PromptDistinctTypes(dossierID)
	if err != nil {
		log.Printf("Failed to get existing prompt types: %v", err)
		return make(map[string][]string)
	}
	return result
}

// --- Deprecated Anthropic/Sonnet Functions ---
// Kept for reference, but no longer used in the main flow.

var anthropicKey string

func callSonnet(prompt string) (string, error) {
	return callSonnetWithRetry(prompt, 5, 15*time.Second)
}

func callSonnetWithRetry(prompt string, maxRetries int, baseDelay time.Duration) (string, error) {
	// ... implementation remains the same, but is not called by the main prompt generation logic.
	return "", fmt.Errorf("callSonnet is deprecated")
}