259 lines
7.4 KiB
Go
259 lines
7.4 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"inou/lib"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
)
|
|
|
|
// --- Local Structs for Prompt Processing ---
|
|
// These are defined locally to this application as per the new architecture.
|
|
|
|
type TriageResponse struct {
|
|
Category string `json:"category"`
|
|
Language string `json:"language"`
|
|
HasData bool `json:"has_data"`
|
|
Error string `json:"error,omitempty"`
|
|
}
|
|
|
|
type ExtractionResult 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"`
|
|
Entries []*EntryData `json:"entries,omitempty"`
|
|
Error string `json:"error,omitempty"`
|
|
}
|
|
|
|
type InputConfig struct {
|
|
Fields []FormField `json:"fields,omitempty"`
|
|
Groups []FormGroup `json:"groups,omitempty"`
|
|
}
|
|
|
|
type FormGroup struct {
|
|
Title string `json:"title"`
|
|
Fields []FormField `json:"fields"`
|
|
}
|
|
|
|
type FormField struct {
|
|
Key string `json:"key"`
|
|
Type string `json:"type"`
|
|
Label string `json:"label"`
|
|
Unit string `json:"unit,omitempty"`
|
|
Options []string `json:"options,omitempty"`
|
|
}
|
|
|
|
type ScheduleSlot struct {
|
|
Days []string `json:"days"`
|
|
Time string `json:"time"`
|
|
}
|
|
|
|
type EntryData struct {
|
|
Value string `json:"value"`
|
|
Data interface{} `json:"data"`
|
|
}
|
|
|
|
var ValidCategories = map[string]bool{
|
|
"vital": true, "exercise": true, "medication": true, "supplement": true,
|
|
"symptom": true, "note": true, "surgery": true, "hospitalization": true,
|
|
"consultation": true, "diagnosis": true, "device": true, "therapy": true,
|
|
"assessment": true, "birth": true, "imaging_finding": true, "eeg_finding": true,
|
|
"provider": true, "question": true, "history": true, "family_history": true,
|
|
"nutrition": true, "fertility": true, "out_of_domain": true,
|
|
}
|
|
|
|
|
|
// --- Main Application Logic ---
|
|
|
|
func main() {
|
|
if len(os.Args) < 2 {
|
|
fmt.Println("Usage: test-prompts <input>")
|
|
os.Exit(1)
|
|
}
|
|
input := strings.Join(os.Args[1:], " ")
|
|
|
|
// Initialize shared library components
|
|
lib.ConfigInit() // Sets lib.GeminiKey from .env
|
|
lib.InitPrompts("api/prompts") // Sets the prompts dir for the lib
|
|
|
|
// Triage
|
|
fmt.Println("=== TRIAGE ===")
|
|
triage, err := runTriage(input)
|
|
if err != nil {
|
|
fmt.Printf("Triage error: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
if triage.Error != "" {
|
|
fmt.Printf("Error: %s\n", triage.Error)
|
|
os.Exit(0)
|
|
}
|
|
fmt.Printf("Category: %s\nLanguage: %s\nHas data: %v\n", triage.Category, triage.Language, triage.HasData)
|
|
|
|
|
|
// Extraction
|
|
fmt.Println("\n=== EXTRACTION ===")
|
|
// For a standalone test, we don't have existing types to pass.
|
|
existingTypes := make(map[string][]string)
|
|
|
|
// Special case for testing pregnancy contradiction
|
|
if input == "I just had my period" {
|
|
existingTypes["fertility"] = []string{"pregnancy"}
|
|
fmt.Println("INFO: Injected existing type: fertility: [pregnancy]")
|
|
}
|
|
|
|
result, err := runExtraction(input, triage.Category, triage.Language, existingTypes)
|
|
if err != nil {
|
|
fmt.Printf("Extraction error: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
out, _ := json.MarshalIndent(result, "", " ")
|
|
fmt.Println(string(out))
|
|
|
|
// Summary
|
|
fmt.Println("\n=== SUMMARY ===")
|
|
if len(result.Entries) > 0 {
|
|
for i, entry := range result.Entries {
|
|
fmt.Printf("Entry %d: %s\n", i+1, entry.Value)
|
|
}
|
|
} else {
|
|
fmt.Println("Entry: (none)")
|
|
}
|
|
|
|
if len(result.Schedule) > 0 && result.Question != "" {
|
|
fmt.Printf("Prompt: %s\n", result.Question)
|
|
fmt.Printf("Schedule contains %d slot(s)\n", len(result.Schedule))
|
|
} else {
|
|
fmt.Println("Prompt: (none - one-time event)")
|
|
}
|
|
}
|
|
|
|
|
|
// --- Local Prompt Handling 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)
|
|
}
|
|
|
|
// Use same context as extraction
|
|
dossierName := "Sophia Helena"
|
|
dossierDOB := "2017-01-01"
|
|
|
|
prompt := strings.ReplaceAll(tmpl, "{{INPUT}}", userInput)
|
|
prompt = strings.ReplaceAll(prompt, "{{DOSSIER_NAME}}", dossierName)
|
|
prompt = strings.ReplaceAll(prompt, "{{DOSSIER_DOB}}", dossierDOB)
|
|
|
|
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)"
|
|
}
|
|
|
|
// Use Sophia's context for testing (child patient helps with practitioner detection)
|
|
dossierName := "Sophia Helena"
|
|
dossierDOB := "2017-01-01"
|
|
currentDate := "2026-02-08"
|
|
|
|
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)
|
|
prompt = strings.ReplaceAll(prompt, "{{DOSSIER_NAME}}", dossierName)
|
|
prompt = strings.ReplaceAll(prompt, "{{DOSSIER_DOB}}", dossierDOB)
|
|
prompt = strings.ReplaceAll(prompt, "{{CURRENT_DATE}}", currentDate)
|
|
|
|
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
|
|
} |