WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

scripts/graph_extraction_probe/main.go

@@ -0,0 +1,170 @@
+// 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)
+	}
+}