kirogo/proxy/identity.go

package proxy

import (
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"unicode"

	"github.com/google/uuid"
)

// identityPatterns covers common ways users ask about the AI's identity.
var identityPatterns = []*regexp.Regexp{
	// English
	regexp.MustCompile(`(?i)\bwho are you\b`),
	regexp.MustCompile(`(?i)\bwhat are you\b`),
	regexp.MustCompile(`(?i)\bwhat model\b`),
	regexp.MustCompile(`(?i)\bwhich model\b`),
	regexp.MustCompile(`(?i)\byour (name|identity|model|version)\b`),
	regexp.MustCompile(`(?i)\btell me about yourself\b`),
	regexp.MustCompile(`(?i)\bidentify yourself\b`),
	regexp.MustCompile(`(?i)\bwhat (llm|language model) are you\b`),
	regexp.MustCompile(`(?i)\bwhat (ai|assistant) are you\b`),
	// Chinese
	regexp.MustCompile(`你是谁`),
	regexp.MustCompile(`你是什么`),
	regexp.MustCompile(`你叫什么`),
	regexp.MustCompile(`什么模型`),
	regexp.MustCompile(`哪个模型`),
	regexp.MustCompile(`你基于什么`),
	regexp.MustCompile(`你是哪个`),
	regexp.MustCompile(`你是哪款`),
	regexp.MustCompile(`你的身份`),
	regexp.MustCompile(`你的名字`),
	regexp.MustCompile(`什么大模型`),
	regexp.MustCompile(`什么AI`),
}

// isIdentityQuestion returns true when the text appears to be asking about AI identity.
func isIdentityQuestion(text string) bool {
	for _, re := range identityPatterns {
		if re.MatchString(text) {
			return true
		}
	}
	return false
}

func hasChinese(s string) bool {
	for _, r := range s {
		if unicode.Is(unicode.Han, r) {
			return true
		}
	}
	return false
}

// friendlyModelName converts a raw model ID to a human-readable Claude model name.
func friendlyModelName(model string) string {
	m := strings.ToLower(model)
	for _, suf := range []string{"-thinking", "-thought"} {
		if strings.HasSuffix(m, suf) {
			m = m[:len(m)-len(suf)]
			break
		}
	}
	switch {
	case strings.Contains(m, "opus-4.7") || strings.Contains(m, "opus-4-7"):
		return "Claude Opus 4.7"
	case strings.Contains(m, "opus-4.6") || strings.Contains(m, "opus-4-6"):
		return "Claude Opus 4.6"
	case strings.Contains(m, "opus-4.5") || strings.Contains(m, "opus-4-5"):
		return "Claude Opus 4.5"
	case strings.Contains(m, "sonnet-4.7") || strings.Contains(m, "sonnet-4-7"):
		return "Claude Sonnet 4.7"
	case strings.Contains(m, "sonnet-4.6") || strings.Contains(m, "sonnet-4-6"):
		return "Claude Sonnet 4.6"
	case strings.Contains(m, "sonnet-4.5") || strings.Contains(m, "sonnet-4-5"):
		return "Claude Sonnet 4.5"
	case strings.Contains(m, "sonnet-4"):
		return "Claude Sonnet 4"
	case strings.Contains(m, "haiku-4.7") || strings.Contains(m, "haiku-4-7"):
		return "Claude Haiku 4.7"
	case strings.Contains(m, "haiku-4.5") || strings.Contains(m, "haiku-4-5"):
		return "Claude Haiku 4.5"
	default:
		return "Claude"
	}
}

// claudeIdentityText returns a Claude-style identity response, language-matched to the question.
func claudeIdentityText(model, question string) string {
	friendly := friendlyModelName(model)
	if hasChinese(question) {
		return "我是 Claude，由 Anthropic 公司开发的 AI 助手。\n\n我目前使用的是 " + friendly + " 模型，可以帮你完成代码编写与调试、分析推理、多语言对话等任务。有什么可以帮你的吗？"
	}
	return "I'm Claude, an AI assistant made by Anthropic.\n\nI'm currently running on the " + friendly + " model. I can help with coding, debugging, analysis, writing, and much more. What can I help you with?"
}

// getLastUserTextClaude extracts the text of the last user message.
func getLastUserTextClaude(msgs []ClaudeMessage) string {
	for i := len(msgs) - 1; i >= 0; i-- {
		if msgs[i].Role == "user" {
			text, _, _ := extractClaudeUserContent(msgs[i].Content)
			return text
		}
	}
	return ""
}

// getLastUserTextOpenAI extracts the text of the last user message.
func getLastUserTextOpenAI(msgs []OpenAIMessage) string {
	for i := len(msgs) - 1; i >= 0; i-- {
		if msgs[i].Role == "user" {
			text, _ := extractOpenAIUserContent(msgs[i].Content)
			return text
		}
	}
	return ""
}

// sendClaudeIdentityNonStream writes a non-streaming Claude identity response.
func (h *Handler) sendClaudeIdentityNonStream(w http.ResponseWriter, model, question string, estimatedInputTokens int) {
	text := claudeIdentityText(model, question)
	outTokens := estimateTextTokens(text)
	resp := KiroToClaudeResponse(text, "", false, nil, estimatedInputTokens, outTokens, model)
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	json.NewEncoder(w).Encode(resp)
}

// sendClaudeIdentityStream writes a streaming Claude identity response.
func (h *Handler) sendClaudeIdentityStream(w http.ResponseWriter, model, question string, estimatedInputTokens int) {
	w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	flusher, ok := w.(http.Flusher)
	if !ok {
		h.sendClaudeError(w, 500, "api_error", "Streaming not supported")
		return
	}

	text := claudeIdentityText(model, question)
	outTokens := estimateTextTokens(text)
	msgID := "msg_" + uuid.New().String()

	send := func(event string, data interface{}) {
		b, _ := json.Marshal(data)
		fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event, b)
		flusher.Flush()
	}

	send("message_start", map[string]interface{}{
		"type": "message_start",
		"message": map[string]interface{}{
			"id":            msgID,
			"type":          "message",
			"role":          "assistant",
			"content":       []interface{}{},
			"model":         model,
			"stop_reason":   nil,
			"stop_sequence": nil,
			"usage":         map[string]int{"input_tokens": estimatedInputTokens, "output_tokens": 0},
		},
	})
	send("content_block_start", map[string]interface{}{
		"type":  "content_block_start",
		"index": 0,
		"content_block": map[string]string{
			"type": "text",
			"text": "",
		},
	})
	send("ping", map[string]string{"type": "ping"})
	send("content_block_delta", map[string]interface{}{
		"type":  "content_block_delta",
		"index": 0,
		"delta": map[string]string{"type": "text_delta", "text": text},
	})
	send("content_block_stop", map[string]interface{}{
		"type":  "content_block_stop",
		"index": 0,
	})
	send("message_delta", map[string]interface{}{
		"type":  "message_delta",
		"delta": map[string]string{"stop_reason": "end_turn"},
		"usage": map[string]int{"output_tokens": outTokens},
	})
	send("message_stop", map[string]string{"type": "message_stop"})
}

// sendOpenAIIdentityNonStream writes a non-streaming OpenAI identity response.
func (h *Handler) sendOpenAIIdentityNonStream(w http.ResponseWriter, model, question string, estimatedInputTokens int) {
	text := claudeIdentityText(model, question)
	outTokens := estimateTextTokens(text)
	resp := OpenAIResponse{
		ID:      "chatcmpl-" + uuid.New().String(),
		Object:  "chat.completion",
		Created: 0,
		Model:   model,
		Choices: []OpenAIChoice{{
			Index:        0,
			Message:      OpenAIMessage{Role: "assistant", Content: text},
			FinishReason: "stop",
		}},
		Usage: OpenAIUsage{
			PromptTokens:     estimatedInputTokens,
			CompletionTokens: outTokens,
			TotalTokens:      estimatedInputTokens + outTokens,
		},
	}
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	json.NewEncoder(w).Encode(resp)
}

// estimateTextTokens estimates token count as word-count * 1.3.
func estimateTextTokens(text string) int {
	n := len(strings.Fields(text))
	if n < 1 {
		n = 1
	}
	return int(float64(n)*1.3) + 1
}