120 lines
3.7 KiB
Go
120 lines
3.7 KiB
Go
// Command superagent-probe performs a safe, explicit Open API connectivity check.
|
||
package main
|
||
|
||
import (
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"log"
|
||
"os"
|
||
"os/signal"
|
||
"syscall"
|
||
|
||
"fire-safety-ymd/internal/config"
|
||
"fire-safety-ymd/internal/integration/superagent"
|
||
)
|
||
|
||
const probeMessage = "请只回复“SuperAgent connectivity OK”,不要调用业务工具。"
|
||
|
||
type probeResult struct {
|
||
Status string `json:"status"`
|
||
SessionID string `json:"session_id"`
|
||
RunID string `json:"run_id,omitempty"`
|
||
ProfileID string `json:"profile_id,omitempty"`
|
||
ProfileVersionID string `json:"profile_version_id,omitempty"`
|
||
ModelName string `json:"model_name,omitempty"`
|
||
Answer string `json:"answer"`
|
||
Usage superagent.TokenUsage `json:"usage"`
|
||
EventTypes []string `json:"event_types"`
|
||
}
|
||
|
||
func main() {
|
||
log.SetFlags(0)
|
||
cfg, err := config.Load()
|
||
if err != nil {
|
||
log.Fatalf("load configuration: %v", err)
|
||
}
|
||
if !cfg.SuperAgent.Enabled {
|
||
log.Fatalf("connectivity probe is disabled; set %s=true with a project-specific test credential", config.SuperAgentEnabledEnv)
|
||
}
|
||
|
||
client, err := superagent.NewHTTPClient(superagent.Config{
|
||
Enabled: cfg.SuperAgent.Enabled,
|
||
BaseURL: cfg.SuperAgent.BaseURL,
|
||
APIKey: cfg.SuperAgent.OpenAPIKey,
|
||
IncludeTrace: cfg.SuperAgent.IncludeTrace,
|
||
ConnectTimeout: cfg.SuperAgent.ConnectTimeout,
|
||
RecoveryMaxAttempts: cfg.SuperAgent.RecoveryMaxAttempts,
|
||
RecoveryInitialBackoff: cfg.SuperAgent.RecoveryInitialBackoff,
|
||
MaxMessageBytes: cfg.SuperAgent.MaxMessageBytes,
|
||
})
|
||
if err != nil {
|
||
log.Fatalf("create SuperAgent client: %v", err)
|
||
}
|
||
|
||
suffix, err := randomSuffix()
|
||
if err != nil {
|
||
log.Fatal("create probe correlation identifiers")
|
||
}
|
||
correlationID := "fsymd-probe-" + suffix
|
||
|
||
signalContext, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
||
defer stop()
|
||
ctx, cancel := context.WithTimeout(signalContext, cfg.SuperAgent.ProbeTimeout)
|
||
defer cancel()
|
||
|
||
session, err := client.CreateSession(ctx, superagent.CreateSessionRequest{
|
||
ExternalSubjectID: cfg.SuperAgent.ProbeSubjectID,
|
||
IdempotencyKey: correlationID + "-session",
|
||
RequestID: correlationID,
|
||
Metadata: map[string]any{
|
||
"source": "fire-safety-ymd",
|
||
"purpose": "connectivity-probe",
|
||
},
|
||
})
|
||
if err != nil {
|
||
log.Fatalf("create SuperAgent probe session: %v", err)
|
||
}
|
||
|
||
result, err := client.StreamMessage(ctx, superagent.StreamMessageRequest{
|
||
SessionID: session.ID,
|
||
Message: probeMessage,
|
||
IdempotencyKey: correlationID + "-message",
|
||
RequestID: correlationID,
|
||
Metadata: map[string]any{
|
||
"source": "fire-safety-ymd",
|
||
"purpose": "connectivity-probe",
|
||
},
|
||
}, nil)
|
||
if err != nil {
|
||
log.Fatalf("stream SuperAgent probe message: %v", err)
|
||
}
|
||
|
||
output := probeResult{
|
||
Status: "ok",
|
||
SessionID: result.SessionID,
|
||
RunID: result.RunID,
|
||
ProfileID: result.ProfileID,
|
||
ProfileVersionID: result.ProfileVersionID,
|
||
ModelName: result.ModelName,
|
||
Answer: result.Answer,
|
||
Usage: result.Usage,
|
||
EventTypes: result.EventTypes,
|
||
}
|
||
encoder := json.NewEncoder(os.Stdout)
|
||
encoder.SetIndent("", " ")
|
||
if err := encoder.Encode(output); err != nil {
|
||
log.Fatalf("write probe result: %v", err)
|
||
}
|
||
}
|
||
|
||
func randomSuffix() (string, error) {
|
||
value := make([]byte, 12)
|
||
if _, err := rand.Read(value); err != nil {
|
||
return "", fmt.Errorf("generate random suffix: %w", err)
|
||
}
|
||
return base64.RawURLEncoding.EncodeToString(value), nil
|
||
}
|