tm_back/pkg/glm/service.go

package glm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"tm/pkg/logger"
)

// Service implements the GLMService interface
type Service struct {
	client *Client
	config *Config
	logger logger.Logger
}

// NewService creates a new GLM service
func NewService(client *Client) GLMService {
	return &Service{
		client: client,
		config: client.config,
		logger: client.logger,
	}
}

// ChatCompletion performs a chat completion request
func (s *Service) ChatCompletion(ctx context.Context, req *ChatCompletionRequest) (*ChatCompletionResponse, error) {
	if s.logger != nil {
		s.logger.Info("Performing chat completion", map[string]interface{}{
			"model":         req.Model,
			"message_count": len(req.Messages),
			"max_tokens":    req.MaxTokens,
			"stream":        req.Stream,
		})
	}

	var resp ChatCompletionResponse
	err := s.client.Post(ctx, "/api/coding/paas/v4/chat/completions", req, &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Chat completion failed", map[string]interface{}{
				"model": req.Model,
				"error": err.Error(),
			})
		}
		return nil, ErrChatCompletion("chat completion request failed", err)
	}

	if s.logger != nil {
		s.logger.Info("Chat completion completed", map[string]interface{}{
			"model":        req.Model,
			"choice_count": len(resp.Choices),
			"usage":        resp.Usage,
		})
	}

	return &resp, nil
}

// StreamChatCompletion performs a streaming chat completion request
func (s *Service) StreamChatCompletion(ctx context.Context, req *ChatCompletionRequest, onChunk func(*StreamChatCompletionResponse) error) error {
	if s.logger != nil {
		s.logger.Info("Starting streaming chat completion", map[string]interface{}{
			"model":         req.Model,
			"message_count": len(req.Messages),
			"max_tokens":    req.MaxTokens,
		})
	}

	// Force streaming
	req.Stream = true

	return s.streamRequest(ctx, req, onChunk)
}

// streamRequest handles streaming requests
func (s *Service) streamRequest(ctx context.Context, req *ChatCompletionRequest, onChunk func(*StreamChatCompletionResponse) error) error {
	// Create a custom HTTP request for streaming
	httpReq, err := s.client.createStreamingRequest(ctx, req)
	if err != nil {
		return ErrStreamError("failed to create streaming request", err)
	}

	resp, err := s.client.httpClient.Do(httpReq)
	if err != nil {
		return ErrStreamError("failed to execute streaming request", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		body, _ := io.ReadAll(resp.Body)
		return s.client.handleHTTPError(resp.StatusCode, body)
	}

	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()

		// Skip empty lines
		if strings.TrimSpace(line) == "" {
			continue
		}

		// Parse SSE format
		if strings.HasPrefix(line, "data: ") {
			data := strings.TrimPrefix(line, "data: ")

			// Check for end of stream
			if data == "[DONE]" {
				break
			}

			var chunk StreamChatCompletionResponse
			if err := json.Unmarshal([]byte(data), &chunk); err != nil {
				if s.logger != nil {
					s.logger.Warn("Failed to parse streaming chunk", map[string]interface{}{
						"error": err.Error(),
						"data":  data,
					})
				}
				continue
			}

			if err := onChunk(&chunk); err != nil {
				return ErrStreamError("chunk callback failed", err)
			}
		}
	}

	if err := scanner.Err(); err != nil {
		return ErrStreamError("streaming scanner error", err)
	}

	if s.logger != nil {
		s.logger.Info("Streaming chat completion finished", map[string]interface{}{})
	}

	return nil
}

// ListModels lists available models
func (s *Service) ListModels(ctx context.Context) (*ModelsResponse, error) {
	if s.logger != nil {
		s.logger.Info("Listing available models", map[string]interface{}{})
	}

	var resp ModelsResponse
	err := s.client.Get(ctx, "/api/coding/paas/v4/models", &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Failed to list models", map[string]interface{}{
				"error": err.Error(),
			})
		}
		return nil, fmt.Errorf("failed to list models: %w", err)
	}

	if s.logger != nil {
		s.logger.Info("Models listed successfully", map[string]interface{}{
			"model_count": len(resp.Data),
		})
	}

	return &resp, nil
}

// Health checks if the GLM service is available
func (s *Service) Health(ctx context.Context) error {
	if s.logger != nil {
		s.logger.Debug("Performing health check", map[string]interface{}{})
	}

	// Simple health check by listing models
	_, err := s.ListModels(ctx)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Health check failed", map[string]interface{}{
				"error": err.Error(),
			})
		}
		return fmt.Errorf("health check failed: %w", err)
	}

	if s.logger != nil {
		s.logger.Debug("Health check passed", map[string]interface{}{})
	}

	return nil
}

// createStreamingRequest creates an HTTP request for streaming
func (c *Client) createStreamingRequest(ctx context.Context, req *ChatCompletionRequest) (*http.Request, error) {
	url := c.config.BaseURL + "/api/coding/paas/v4/chat/completions"

	bodyBytes, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(bodyBytes))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	// Set headers
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("User-Agent", c.config.UserAgent)
	if c.config.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+c.config.APIKey)
	}

	return httpReq, nil
}