Week 1 Day 5：Structured Output - 苏格拉底教学#

// internal/agent/answer.go
package agent

import (
	"encoding/json"
	"fmt"
	"strings"
)

// Confidence 置信度枚举
type Confidence string

const (
	ConfidenceLow    Confidence = "low"
	ConfidenceMedium Confidence = "medium"
	ConfidenceHigh   Confidence = "high"
)

// Source 引用的文档来源
type Source struct {
	DocID   string `json:"doc_id"`
	ChunkID string `json:"chunk_id"`
	Title   string `json:"title,omitempty"`
	URL     string `json:"url,omitempty"`
	Snippet string `json:"snippet,omitempty"` // 引用的原文片段
}

// ToolCall Agent调用的工具记录
type ToolCall struct {
	Name      string          `json:"name"`
	Arguments json.RawMessage `json:"arguments"`
	Result    json.RawMessage `json:"result,omitempty"`
	Error     string          `json:"error,omitempty"`
}

// AgentAnswer Agent的完整结构化回答
type AgentAnswer struct {
	Answer            string     `json:"answer"`
	Confidence        Confidence `json:"confidence"`
	Sources           []Source   `json:"sources"`
	ToolCalls         []ToolCall `json:"tool_calls"`
	NeedsHumanReview  bool       `json:"needs_human_review"`
	Reasoning         string     `json:"reasoning"`
}

// internal/agent/validate.go
package agent

// Violation 一条违规记录
type Violation struct {
	Field   string `json:"field"`
	Rule    string `json:"rule"`
	Message string `json:"message"`
}

// Validate 返回所有违规；为空表示合法
func (a *AgentAnswer) Validate() []Violation {
	var violations []Violation

	// 规则1：answer不能为空
	if strings.TrimSpace(a.Answer) == "" {
		violations = append(violations, Violation{
			Field:   "answer",
			Rule:    "required",
			Message: "answer must not be empty",
		})
	}

	// 规则2：confidence必须是枚举值
	switch a.Confidence {
	case ConfidenceLow, ConfidenceMedium, ConfidenceHigh:
	default:
		violations = append(violations, Violation{
			Field:   "confidence",
			Rule:    "enum",
			Message: fmt.Sprintf("confidence must be low/medium/high, got %q", a.Confidence),
		})
	}

	// 规则3：sources为空 → confidence不能是high
	if len(a.Sources) == 0 && a.Confidence == ConfidenceHigh {
		violations = append(violations, Violation{
			Field:   "confidence",
			Rule:    "consistency",
			Message: "confidence cannot be high when sources is empty (no grounding)",
		})
	}

	// 规则4：创建工单 → needs_human_review必须true
	for _, tc := range a.ToolCalls {
		if tc.Name == "create_ticket" && !a.NeedsHumanReview {
			violations = append(violations, Violation{
				Field:   "needs_human_review",
				Rule:    "policy",
				Message: "creating a ticket requires human review",
			})
		}
	}

	// 规则5：confidence=low → needs_human_review必须true
	if a.Confidence == ConfidenceLow && !a.NeedsHumanReview {
		violations = append(violations, Violation{
			Field:   "needs_human_review",
			Rule:    "consistency",
			Message: "low confidence answers must be marked for review",
		})
	}

	// 规则6：每个source必须有doc_id和chunk_id
	for i, src := range a.Sources {
		if src.DocID == "" || src.ChunkID == "" {
			violations = append(violations, Violation{
				Field:   fmt.Sprintf("sources[%d]", i),
				Rule:    "required",
				Message: "each source must have doc_id and chunk_id",
			})
		}
	}

	return violations
}

// IsValid 简便方法
func (a *AgentAnswer) IsValid() bool {
	return len(a.Validate()) == 0
}

func TestValidate_SourcesEmptyHighConfidence(t *testing.T) {
	a := AgentAnswer{
		Answer:     "SSO密码可以通过portal重置",
		Confidence: ConfidenceHigh,
		Sources:    []Source{}, // 空！
	}
	vs := a.Validate()
	if len(vs) == 0 {
		t.Fatal("expected violations for empty sources with high confidence")
	}
}

// internal/agent/schema.go
package agent

import "encoding/json"

// AnswerSchema 返回给OpenAI response_format用的schema
func AnswerSchema() json.RawMessage {
	schema := map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"answer": map[string]interface{}{
				"type":        "string",
				"description": "Final answer shown to the user.",
			},
			"confidence": map[string]interface{}{
				"type": "string",
				"enum": []string{"low", "medium", "high"},
				"description": "high=grounded in sources; medium=partial; low=guessing or no sources.",
			},
			"sources": map[string]interface{}{
				"type": "array",
				"items": map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"doc_id":   map[string]interface{}{"type": "string"},
						"chunk_id": map[string]interface{}{"type": "string"},
						"title":    map[string]interface{}{"type": "string"},
						"snippet":  map[string]interface{}{"type": "string"},
					},
					"required":             []string{"doc_id", "chunk_id"},
					"additionalProperties": false,
				},
			},
			"tool_calls": map[string]interface{}{
				"type": "array",
				"items": map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"name":      map[string]interface{}{"type": "string"},
						"arguments": map[string]interface{}{"type": "object"},
					},
					"required": []string{"name", "arguments"},
				},
			},
			"needs_human_review": map[string]interface{}{
				"type": "boolean",
			},
			"reasoning": map[string]interface{}{
				"type":        "string",
				"description": "Brief chain-of-thought, not shown to user.",
			},
		},
		"required": []string{
			"answer", "confidence", "sources",
			"tool_calls", "needs_human_review", "reasoning",
		},
		"additionalProperties": false,
	}
	b, _ := json.Marshal(schema)
	return b
}

// internal/agent/generate.go
package agent

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/sashabaranov/go-openai"
)

type Generator struct {
	client *openai.Client
	model  string
}

func NewGenerator(apiKey, model string) *Generator {
	return &Generator{
		client: openai.NewClient(apiKey),
		model:  model,
	}
}

// Generate 调用LLM并解析成AgentAnswer
func (g *Generator) Generate(ctx context.Context, userQuery string, retrievedChunks []Source) (*AgentAnswer, error) {
	systemPrompt := buildSystemPrompt(retrievedChunks)

	resp, err := g.client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
		Model: g.model,
		Messages: []openai.ChatCompletionMessage{
			{Role: openai.ChatMessageRoleSystem, Content: systemPrompt},
			{Role: openai.ChatMessageRoleUser, Content: userQuery},
		},
		ResponseFormat: &openai.ChatCompletionResponseFormat{
			Type: openai.ChatCompletionResponseFormatTypeJSONObject,
		},
		Temperature: 0.2, // 结构化输出建议低温度
	})
	if err != nil {
		return nil, fmt.Errorf("openai call: %w", err)
	}

	if len(resp.Choices) == 0 {
		return nil, fmt.Errorf("no choices returned")
	}

	var answer AgentAnswer
	raw := resp.Choices[0].Message.Content
	if err := json.Unmarshal([]byte(raw), &answer); err != nil {
		return nil, fmt.Errorf("parse answer json: %w (raw=%s)", err, raw)
	}

	// 验证
	if violations := answer.Validate(); len(violations) > 0 {
		return &answer, &ValidationError{Violations: violations, Raw: raw}
	}

	return &answer, nil
}

type ValidationError struct {
	Violations []Violation
	Raw        string
}

func (e *ValidationError) Error() string {
	return fmt.Sprintf("validation failed: %d violations", len(e.Violations))
}

// internal/agent/prompt.go
package agent

import (
	"fmt"
	"strings"
)

const systemPromptTemplate = `You are a customer support agent. You MUST respond in strict JSON following this schema:

{
  "answer": "<final answer to user>",
  "confidence": "low" | "medium" | "high",
  "sources": [{"doc_id": "...", "chunk_id": "...", "title": "...", "snippet": "..."}],
  "tool_calls": [{"name": "...", "arguments": {...}}],
  "needs_human_review": true | false,
  "reasoning": "<brief thought>"
}

## CRITICAL RULES (violating these causes your response to be REJECTED):

1. If "sources" is empty, "confidence" MUST be "low". Never guess.
2. If you use tool "create_ticket", "needs_human_review" MUST be true.
3. If "confidence" is "low", "needs_human_review" MUST be true.
4. Every source must include both "doc_id" and "chunk_id".
5. "answer" must be non-empty.

## Grounding

Answer ONLY based on the following retrieved context. If the context does not
contain the answer, say so honestly and set confidence=low.

## Retrieved Context

%s

## Examples

Good response (grounded):
{"answer":"To reset SSO, visit the portal...","confidence":"high","sources":[{"doc_id":"faq_001","chunk_id":"faq_001_chunk_2","snippet":"Click Forgot Password"}],"tool_calls":[],"needs_human_review":false,"reasoning":"Found direct answer in FAQ."}

Good response (no grounding):
{"answer":"I don't have information about this. Let me escalate.","confidence":"low","sources":[],"tool_calls":[],"needs_human_review":true,"reasoning":"No relevant docs retrieved."}

Now respond in strict JSON.`

func buildSystemPrompt(chunks []Source) string {
	var sb strings.Builder
	for i, c := range chunks {
		sb.WriteString(fmt.Sprintf("[%d] doc_id=%s chunk_id=%s\n%s\n\n",
			i+1, c.DocID, c.ChunkID, c.Snippet))
	}
	if sb.Len() == 0 {
		sb.WriteString("(no relevant context retrieved)")
	}
	return fmt.Sprintf(systemPromptTemplate, sb.String())
}

// internal/agent/repair.go
func (g *Generator) GenerateWithRepair(ctx context.Context, query string, chunks []Source) (*AgentAnswer, error) {
	answer, err := g.Generate(ctx, query, chunks)
	var vErr *ValidationError
	if !errors.As(err, &vErr) {
		return answer, err // 成功或非validation错误
	}

	// 让LLM自己修复
	repairPrompt := fmt.Sprintf(
		"Your previous response violated these rules:\n%s\n\nOriginal response:\n%s\n\nFix it and return strict JSON again.",
		formatViolations(vErr.Violations), vErr.Raw,
	)

	resp, err := g.client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
		Model: g.model,
		Messages: []openai.ChatCompletionMessage{
			{Role: openai.ChatMessageRoleUser, Content: repairPrompt},
		},
		ResponseFormat: &openai.ChatCompletionResponseFormat{
			Type: openai.ChatCompletionResponseFormatTypeJSONObject,
		},
	})
	if err != nil {
		return nil, err
	}

	var repaired AgentAnswer
	if err := json.Unmarshal([]byte(resp.Choices[0].Message.Content), &repaired); err != nil {
		return nil, err
	}
	return &repaired, nil
}

// 给定'1'（陆地）'0'（水）的2D grid，返回岛屿数
func numIslands(grid [][]byte) int {
	if len(grid) == 0 {
		return 0
	}
	m, n := len(grid), len(grid[0])
	count := 0

	var dfs func(i, j int)
	dfs = func(i, j int) {
		if i < 0 || i >= m || j < 0 || j >= n || grid[i][j] != '1' {
			return
		}
		grid[i][j] = '0' // 标记访问
		dfs(i+1, j)
		dfs(i-1, j)
		dfs(i, j+1)
		dfs(i, j-1)
	}

	for i := 0; i < m; i++ {
		for j := 0; j < n; j++ {
			if grid[i][j] == '1' {
				count++
				dfs(i, j)
			}
		}
	}
	return count
}

func maxAreaOfIsland(grid [][]int) int {
	m, n := len(grid), len(grid[0])
	maxArea := 0

	var dfs func(i, j int) int
	dfs = func(i, j int) int {
		if i < 0 || i >= m || j < 0 || j >= n || grid[i][j] != 1 {
			return 0
		}
		grid[i][j] = 0
		return 1 + dfs(i+1, j) + dfs(i-1, j) + dfs(i, j+1) + dfs(i, j-1)
	}

	for i := 0; i < m; i++ {
		for j := 0; j < n; j++ {
			if grid[i][j] == 1 {
				if a := dfs(i, j); a > maxArea {
					maxArea = a
				}
			}
		}
	}
	return maxArea
}

type Node struct {
	Val       int
	Neighbors []*Node
}

func cloneGraph(node *Node) *Node {
	if node == nil {
		return nil
	}
	visited := map[*Node]*Node{}

	var clone func(n *Node) *Node
	clone = func(n *Node) *Node {
		if cp, ok := visited[n]; ok {
			return cp
		}
		cp := &Node{Val: n.Val}
		visited[n] = cp // 必须先登记再递归，防止环
		for _, nb := range n.Neighbors {
			cp.Neighbors = append(cp.Neighbors, clone(nb))
		}
		return cp
	}
	return clone(node)
}