descrybe/apps/api/cmd/_chat_hang_probe_tmp/main.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/descrybe/descrybe-v2/apps/api/internal/aiprovider"
	"github.com/descrybe/descrybe-v2/apps/api/internal/config"
	"github.com/descrybe/descrybe-v2/apps/api/internal/db"
	"github.com/descrybe/descrybe-v2/apps/api/internal/platformsettings"
	"github.com/descrybe/descrybe-v2/apps/api/internal/processing"
	"github.com/google/uuid"
)

func main() {
	cfg, err := config.Load()
	if err != nil {
		fmt.Fprintf(os.Stderr, "config: %v\n", err)
		os.Exit(2)
	}
	ctx := context.Background()
	pool, err := db.NewPool(ctx, cfg.DatabaseURL, db.PoolOptions{
		MaxConns: int32(cfg.DBMaxConns), MinConns: int32(cfg.DBMinConns),
		MaxConnLifetime: cfg.DBMaxConnLifetime, MaxConnLifetimeJitter: cfg.DBMaxConnLifetimeJitter,
		MaxConnIdleTime: cfg.DBMaxConnIdleTime, HealthCheckPeriod: cfg.DBHealthCheckPeriod,
		StatementTimeout: cfg.DBStatementTimeout,
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "db: %v\n", err)
		os.Exit(2)
	}
	defer pool.Close()

	platEnv := platformsettings.EnvConfig{
		AppEncryptionKey: cfg.AppEncryptionKey, CredentialsEncryptionKey: cfg.CredentialsEncryptionKey,
		TokenSigningSecret: cfg.TokenSigningSecret, DatabaseURL: cfg.DatabaseURL,
		OpenAIAPIKey: cfg.OpenAIAPIKey, OpenAIBaseURL: cfg.OpenAIBaseURL, OpenAIModel: cfg.OpenAIModel,
	}
	plat := platformsettings.NewService(pool, platEnv)
	ai := aiprovider.NewService(pool, aiprovider.EnvConfig{
		AppEncryptionKey: cfg.AppEncryptionKey, CredentialsEncryptionKey: cfg.CredentialsEncryptionKey,
		TokenSigningSecret: cfg.TokenSigningSecret, DatabaseURL: cfg.DatabaseURL,
		OpenAIAPIKey: cfg.OpenAIAPIKey, OpenAIBaseURL: cfg.OpenAIBaseURL, OpenAIModel: cfg.OpenAIModel,
	})
	ai.Platform = plat

	companyID := uuid.MustParse("604f23a8-b66e-4b21-8b45-0d72b68f4790")
	roleCfg, err := plat.ResolveAIConfig(ctx, platformsettings.AIRoleProcessing)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ResolveAIConfig: %v\n", err)
		os.Exit(2)
	}
	completer, mode, byok, err := ai.ResolveCompleterForRole(ctx, companyID, aiprovider.RoleProcessing)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ResolveCompleter: %v\n", err)
		os.Exit(2)
	}
	client, ok := completer.(*processing.OpenAIClient)
	if !ok {
		fmt.Fprintf(os.Stderr, "not OpenAIClient: %T\n", completer)
		os.Exit(2)
	}
	base := strings.TrimRight(strings.TrimSpace(client.BaseURL), "/")
	model := strings.TrimSpace(client.Model)
	if processing.IsMockOrLoopbackBaseURL(base) {
		fmt.Println("FAIL mock/loopback")
		os.Exit(3)
	}

	fmt.Printf("resolved source=%s provider=%s base=%s model=%s mode=%s byok=%v key_len=%d app_max_tokens_enhance=%d retry=%d http_timeout=240s\n",
		roleCfg.Source, roleCfg.Provider, base, model, mode, byok, len(client.APIKey),
		processing.MaxTokensEnhance, processing.MaxTokensEnhanceRetry)

	// Optional: adjust model id if gateway wants green/ prefix (prior resolve used both).
	models := []string{model}
	if !strings.HasPrefix(model, "green/") {
		models = append(models, "green/"+model)
	}

	size := "tiny"
	if len(os.Args) > 1 {
		size = strings.ToLower(strings.TrimSpace(os.Args[1]))
	}
	maxTok := 256
	prompt := "Reply with exactly: OK"
	timeout := 60 * time.Second
	switch size {
	case "tiny":
		maxTok = 256
		prompt = "Reply with exactly one word: OK"
		timeout = 60 * time.Second
	case "medium":
		maxTok = 3072
		prompt = "Write a short JSON object with keys title and description for product Sony WH-1000XM5 headphones. description must be 2 short HTML paragraphs. Keep under 400 words."
		timeout = 90 * time.Second
	case "large":
		maxTok = 16384
		prompt = "Write a long JSON object with keys title, description, attributes. description must be multi-section Slovenian HTML with several h2+p and a ul list for Sony WH-1000XM5. Be thorough."
		timeout = 120 * time.Second
	default:
		fmt.Fprintf(os.Stderr, "usage: %s [tiny|medium|large]\n", os.Args[0])
		os.Exit(2)
	}

	useModel := models[0]
	if len(os.Args) > 2 && strings.TrimSpace(os.Args[2]) != "" {
		useModel = strings.TrimSpace(os.Args[2])
	}

	bodyObj := map[string]any{
		"model":       useModel,
		"temperature": 0.2,
		"max_tokens":  maxTok,
		"messages": []map[string]string{
			{"role": "system", "content": "You are a concise assistant."},
			{"role": "user", "content": prompt},
		},
	}
	raw, _ := json.Marshal(bodyObj)
	url := base + "/chat/completions"
	fmt.Printf("probe size=%s model=%s max_tokens=%d timeout=%s url=%s prompt_len=%d\n",
		size, useModel, maxTok, timeout, url, len(prompt))

	reqCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, url, bytes.NewReader(raw))
	if err != nil {
		fmt.Fprintf(os.Stderr, "new request: %v\n", err)
		os.Exit(4)
	}
	req.Header.Set("Authorization", "Bearer "+client.APIKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	// Separate dial vs header wait so we can tell "hung with no headers" from slow body.
	headerTO := 25 * time.Second
	if timeout < headerTO {
		headerTO = timeout
	}
	httpClient := &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			ResponseHeaderTimeout: headerTO,
			IdleConnTimeout:       30 * time.Second,
		},
	}
	start := time.Now()
	resp, err := httpClient.Do(req)
	elapsed := time.Since(start).Round(time.Millisecond)
	if err != nil {
		fmt.Printf("RESULT ok=false elapsed=%s header_timeout=%s err=%v\n", elapsed, headerTO, err)
		os.Exit(5)
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
	snip := strings.TrimSpace(string(b))
	if len(snip) > 500 {
		snip = snip[:500] + "…"
	}
	fmt.Printf("RESULT ok=%t status=%d elapsed=%s body_len=%d snippet=%s\n",
		resp.StatusCode >= 200 && resp.StatusCode < 300, resp.StatusCode, elapsed, len(b), snip)
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		os.Exit(6)
	}
}