xinmi/scripts/graph_extraction_probe/main.go

// graph_extraction_probe calls the production extractor without writing graph data.
// The model configuration and AES key arrive on stdin and are never reported.
package main

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"time"
	"unicode/utf8"

	chatpipeline "github.com/Tencent/WeKnora/internal/application/service/chat_pipeline"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/utils"
	"gopkg.in/yaml.v3"
)

type request struct {
	ExecutionID   string              `json:"execution_id"`
	MaxHTTPSends  *int                `json:"max_http_sends"`
	AESKey        string              `json:"aes_key"`
	Model         types.Model         `json:"model"`
	Chunk         types.Chunk         `json:"chunk"`
	ExtractConfig types.ExtractConfig `json:"extract_config"`
}

type chatDelegate interface{ chat.Chat }

type probeCall struct {
	Number           int              `json:"number"`
	ElapsedMS        int64            `json:"elapsed_ms"`
	MaxTokens        int              `json:"max_tokens"`
	Thinking         *bool            `json:"thinking_requested"`
	ResponseReceived bool             `json:"response_received"`
	Failed           bool             `json:"failed"`
	FinishReason     string           `json:"finish_reason"`
	ContentRunes     int              `json:"content_runes"`
	ReasoningRunes   int              `json:"reasoning_runes"`
	Usage            types.TokenUsage `json:"usage"`
}

type recordingChat struct {
	chatDelegate
	response *types.ChatResponse
	calls    int
	records  []probeCall
}

func (c *recordingChat) Chat(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
	started := time.Now()
	response, err := c.chatDelegate.Chat(ctx, messages, options)
	c.response = response
	c.calls++
	record := probeCall{
		Number: c.calls, ElapsedMS: time.Since(started).Milliseconds(),
		ResponseReceived: response != nil, Failed: err != nil,
	}
	if options != nil {
		record.MaxTokens = options.MaxTokens
		if options.Thinking != nil {
			thinking := *options.Thinking
			record.Thinking = &thinking
		}
	}
	if response != nil {
		record.FinishReason = response.FinishReason
		record.ContentRunes = utf8.RuneCountInString(response.Content)
		record.ReasoningRunes = utf8.RuneCountInString(response.ReasoningContent)
		record.Usage = response.Usage
	}
	c.records = append(c.records, record)
	return response, err
}

func run() error {
	var input request
	if err := json.NewDecoder(os.Stdin).Decode(&input); err != nil {
		return fmt.Errorf("invalid probe input")
	}
	capValue, capSet := os.LookupEnv("GRAPH_PROBE_MAX_SENDS")
	maxSends, err := probeSendLimit(input.MaxHTTPSends, capValue, capSet)
	if err != nil {
		return err
	}
	if input.ExecutionID == "" || input.Chunk.ID == "" || input.Chunk.Content == "" || input.Model.ID == "" || input.Model.Name == "" {
		return fmt.Errorf("probe requires a source chunk and model")
	}
	if err := os.Setenv("SYSTEM_AES_KEY", input.AESKey); err != nil {
		return err
	}
	input.AESKey = ""
	key, err := utils.DecryptStoredSecret(input.Model.Parameters.APIKey)
	if err != nil {
		return fmt.Errorf("cannot decrypt configured model credential")
	}
	input.Model.Parameters.APIKey = key
	defer os.Unsetenv("SYSTEM_AES_KEY")
	chatConfig := chat.ConfigFromModel(&input.Model, "", "")
	model, err := chat.NewChat(chatConfig, nil)
	if err != nil {
		return fmt.Errorf("cannot initialize configured remote model")
	}
	base, err := os.ReadFile("config/config.yaml")
	if err != nil {
		return err
	}
	var config struct {
		Extract struct {
			Graph struct {
				Description string `yaml:"description"`
			} `yaml:"extract_graph"`
		} `yaml:"extract"`
	}
	if err := yaml.Unmarshal(base, &config); err != nil {
		return err
	}
	recorder := &recordingChat{chatDelegate: model}
	extractor := chatpipeline.NewGraphExtractor(recorder, &types.PromptTemplateStructured{
		Description: types.AppendCustomPromptInstructions(config.Extract.Graph.Description, input.ExtractConfig.CustomInstructions, "graph_extraction"),
		Tags:        input.ExtractConfig.Tags,
	})
	ctx, guard, cancel, err := bindProbeExecution(context.Background(), input.ExecutionID, input.Model.ID, maxSends)
	if err != nil {
		return err
	}
	defer cancel()
	started := time.Now()
	graph, extractionErr := extractor.Extract(ctx, input.Chunk.Content)
	hash := sha256.Sum256([]byte(input.Chunk.Content))
	report := map[string]any{
		"checked_at": started.UTC().Format(time.RFC3339), "elapsed_ms": time.Since(started).Milliseconds(),
		"chunk_id": input.Chunk.ID, "source_sha256": hex.EncodeToString(hash[:]),
		"model_name": input.Model.Name, "allowed_relation_types": input.ExtractConfig.Tags,
		"passed": extractionErr == nil, "graph": graph, "production_write": false,
		"model_calls":                recorder.calls,
		"execution_id":               input.ExecutionID,
		"http_send_reservations":     guard.Reservations(),
		"http_send_limit":            maxSends,
		"call_diagnostics":           recorder.records,
		"effective_thinking_control": chat.EffectiveThinkingControl(chatConfig),
		"diagnostics_version":        2,
	}
	if recorder.response != nil {
		report["finish_reason"] = recorder.response.FinishReason
		report["usage"] = recorder.response.Usage
		if extractionErr != nil {
			report["rejected_model_content"] = recorder.response.Content
		}
	}
	if extractionErr != nil {
		report["error"] = extractionErr.Error()
	}
	encoded, err := json.Marshal(report)
	if err != nil {
		return err
	}
	fmt.Printf("GRAPH_EXTRACTION_PROBE_RESULT=%s\n", encoded)
	return nil
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}