persona-community-3/pkg/textgen/adapters/laozhang.go

package adapters

import (
	"context"
	"fmt"
	"strings"

	"git.threesix.ai/jordan/persona-community-3/pkg/laozhang"
	"git.threesix.ai/jordan/persona-community-3/pkg/textgen"
)

const (
	// IMPORTANT: Always use gemini-3-flash-preview for text generation.
	// DO NOT use gemini-2.5-flash or any 2.x models.
	defaultLaoZhangTextModel = "gemini-3-flash-preview"
)

// LaoZhangTextProvider implements textgen.TextGenerator using LaoZhang API.
type LaoZhangTextProvider struct {
	client *laozhang.Client
	model  string
}

// NewLaoZhangTextProvider creates a new LaoZhang text generation provider.
func NewLaoZhangTextProvider(client *laozhang.Client, model string) *LaoZhangTextProvider {
	if model == "" {
		model = defaultLaoZhangTextModel
	}
	return &LaoZhangTextProvider{
		client: client,
		model:  model,
	}
}

// Name implements textgen.Provider.
func (p *LaoZhangTextProvider) Name() string {
	return "laozhang"
}

// Health implements textgen.Provider.
func (p *LaoZhangTextProvider) Health(ctx context.Context) error {
	return p.client.Health(ctx)
}

// GenerateText implements textgen.TextGenerator.
func (p *LaoZhangTextProvider) GenerateText(ctx context.Context, req textgen.TextRequest) (*textgen.TextResponse, error) {
	model := req.Model
	if model == "" {
		model = p.model
	}

	// Build messages from request
	var messages []laozhang.ChatMessage

	// Add system prompt if provided
	if req.SystemPrompt != "" {
		messages = append(messages, laozhang.ChatMessage{
			Role:    "system",
			Content: req.SystemPrompt,
		})
	}

	// Add messages if provided (multi-turn conversation)
	if len(req.Messages) > 0 {
		for _, msg := range req.Messages {
			messages = append(messages, laozhang.ChatMessage{
				Role:    msg.Role,
				Content: msg.Content,
			})
		}
	} else if req.Prompt != "" {
		// Single prompt
		messages = append(messages, laozhang.ChatMessage{
			Role:    "user",
			Content: req.Prompt,
		})
	}

	// Build chat request
	chatReq := laozhang.ChatCompletionRequest{
		Model:    model,
		Messages: messages,
	}

	if req.MaxTokens > 0 {
		chatReq.MaxTokens = req.MaxTokens
	}
	if req.Temperature > 0 {
		chatReq.Temperature = req.Temperature
	}

	// Call LaoZhang API
	resp, err := p.client.ChatCompletion(ctx, chatReq)
	if err != nil {
		return nil, classifyLaoZhangError(err)
	}

	// Extract text from response
	if len(resp.Choices) == 0 {
		return nil, fmt.Errorf("empty response from LaoZhang")
	}

	responseText := resp.Choices[0].Message.Content
	if responseText == "" {
		return nil, fmt.Errorf("empty content in LaoZhang response")
	}

	// Build response
	result := &textgen.TextResponse{
		Text: responseText,
		Usage: &textgen.Usage{
			PromptTokens:     resp.Usage.PromptTokens,
			CompletionTokens: resp.Usage.CompletionTokens,
			TotalTokens:      resp.Usage.TotalTokens,
		},
	}

	return result, nil
}

// classifyLaoZhangError converts LaoZhang errors to textgen sentinel errors.
func classifyLaoZhangError(err error) error {
	if err == nil {
		return nil
	}

	errStr := strings.ToLower(err.Error())

	// Check for common error patterns
	switch {
	case strings.Contains(errStr, "quota"):
		return fmt.Errorf("%w: %v", textgen.ErrQuotaExceeded, err)
	case strings.Contains(errStr, "rate") || strings.Contains(errStr, "429"):
		return fmt.Errorf("%w: %v", textgen.ErrRateLimited, err)
	case strings.Contains(errStr, "safety") || strings.Contains(errStr, "blocked") || strings.Contains(errStr, "content"):
		return fmt.Errorf("%w: %v", textgen.ErrContentBlocked, err)
	case strings.Contains(errStr, "timeout") || strings.Contains(errStr, "deadline"):
		return fmt.Errorf("%w: %v", textgen.ErrTimeout, err)
	default:
		return err
	}
}

// GenerateStream implements textgen.TextStreamer using LaoZhang streaming API.
func (p *LaoZhangTextProvider) GenerateStream(ctx context.Context, req textgen.TextRequest, onChunk func(textgen.StreamChunk)) error {
	model := req.Model
	if model == "" {
		model = p.model
	}

	// Build messages from request
	var messages []laozhang.ChatMessage
	if req.SystemPrompt != "" {
		messages = append(messages, laozhang.ChatMessage{
			Role:    "system",
			Content: req.SystemPrompt,
		})
	}
	if len(req.Messages) > 0 {
		for _, msg := range req.Messages {
			messages = append(messages, laozhang.ChatMessage{
				Role:    msg.Role,
				Content: msg.Content,
			})
		}
	} else if req.Prompt != "" {
		messages = append(messages, laozhang.ChatMessage{
			Role:    "user",
			Content: req.Prompt,
		})
	}

	chatReq := laozhang.ChatCompletionRequest{
		Model:    model,
		Messages: messages,
	}
	if req.MaxTokens > 0 {
		chatReq.MaxTokens = req.MaxTokens
	}
	if req.Temperature > 0 {
		chatReq.Temperature = req.Temperature
	}

	return p.client.ChatCompletionStream(ctx, chatReq, func(chunk laozhang.StreamChunk) {
		onChunk(textgen.StreamChunk{
			Text:     chunk.Text,
			Done:     chunk.Done,
			Provider: func() string { if chunk.Done { return p.Name() }; return "" }(),
		})
	})
}

// Compile-time interface checks
var _ textgen.TextGenerator = (*LaoZhangTextProvider)(nil)
var _ textgen.TextStreamer = (*LaoZhangTextProvider)(nil)