gooseek/backend/internal/llm/timeweb.go

package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

type TimewebClient struct {
	baseClient
	httpClient    *http.Client
	baseURL       string
	agentAccessID string
	apiKey        string
	proxySource   string
}

type TimewebConfig struct {
	ProviderID    string
	ModelKey      string
	BaseURL       string
	AgentAccessID string
	APIKey        string
	ProxySource   string
}

func NewTimewebClient(cfg TimewebConfig) (*TimewebClient, error) {
	if cfg.AgentAccessID == "" {
		return nil, errors.New("agent_access_id is required for Timeweb")
	}
	if cfg.APIKey == "" {
		return nil, errors.New("api_key is required for Timeweb")
	}

	baseURL := cfg.BaseURL
	if baseURL == "" {
		baseURL = "https://api.timeweb.cloud"
	}

	proxySource := cfg.ProxySource
	if proxySource == "" {
		proxySource = "gooseek"
	}

	return &TimewebClient{
		baseClient: baseClient{
			providerID: cfg.ProviderID,
			modelKey:   cfg.ModelKey,
		},
		httpClient: &http.Client{
			Timeout: 120 * time.Second,
		},
		baseURL:       baseURL,
		agentAccessID: cfg.AgentAccessID,
		apiKey:        cfg.APIKey,
		proxySource:   proxySource,
	}, nil
}

type timewebChatRequest struct {
	Model       string                   `json:"model,omitempty"`
	Messages    []timewebMessage         `json:"messages"`
	Stream      bool                     `json:"stream,omitempty"`
	Temperature float64                  `json:"temperature,omitempty"`
	MaxTokens   int                      `json:"max_tokens,omitempty"`
	TopP        float64                  `json:"top_p,omitempty"`
	Tools       []timewebTool            `json:"tools,omitempty"`
	Stop        []string                 `json:"stop,omitempty"`
}

type timewebMessage struct {
	Role       string                   `json:"role"`
	Content    interface{}              `json:"content"`
	Name       string                   `json:"name,omitempty"`
	ToolCalls  []timewebToolCall        `json:"tool_calls,omitempty"`
	ToolCallID string                   `json:"tool_call_id,omitempty"`
}

type timewebTool struct {
	Type     string              `json:"type"`
	Function timewebFunction     `json:"function"`
}

type timewebFunction struct {
	Name        string      `json:"name"`
	Description string      `json:"description"`
	Parameters  interface{} `json:"parameters"`
}

type timewebToolCall struct {
	ID       string `json:"id"`
	Type     string `json:"type"`
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"`
	} `json:"function"`
}

type timewebChatResponse struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	Model   string `json:"model"`
	Choices []struct {
		Index   int `json:"index"`
		Message struct {
			Role      string            `json:"role"`
			Content   string            `json:"content"`
			ToolCalls []timewebToolCall `json:"tool_calls,omitempty"`
		} `json:"message"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
}

type timewebStreamResponse struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	Model   string `json:"model"`
	Choices []struct {
		Index int `json:"index"`
		Delta struct {
			Role      string            `json:"role,omitempty"`
			Content   string            `json:"content,omitempty"`
			ToolCalls []timewebToolCall `json:"tool_calls,omitempty"`
		} `json:"delta"`
		FinishReason string `json:"finish_reason,omitempty"`
	} `json:"choices"`
}

func (c *TimewebClient) StreamText(ctx context.Context, req StreamRequest) (<-chan StreamChunk, error) {
	messages := make([]timewebMessage, 0, len(req.Messages))
	for _, m := range req.Messages {
		msg := timewebMessage{
			Role:    string(m.Role),
			Content: m.Content,
		}
		if m.Name != "" {
			msg.Name = m.Name
		}
		if m.ToolCallID != "" {
			msg.ToolCallID = m.ToolCallID
		}
		if len(m.ToolCalls) > 0 {
			msg.ToolCalls = make([]timewebToolCall, len(m.ToolCalls))
			for i, tc := range m.ToolCalls {
				args, _ := json.Marshal(tc.Arguments)
				msg.ToolCalls[i] = timewebToolCall{
					ID:   tc.ID,
					Type: "function",
				}
				msg.ToolCalls[i].Function.Name = tc.Name
				msg.ToolCalls[i].Function.Arguments = string(args)
			}
		}
		messages = append(messages, msg)
	}

	chatReq := timewebChatRequest{
		Model:    c.modelKey,
		Messages: messages,
		Stream:   true,
	}

	if req.Options.MaxTokens > 0 {
		chatReq.MaxTokens = req.Options.MaxTokens
	}
	if req.Options.Temperature > 0 {
		chatReq.Temperature = req.Options.Temperature
	}
	if req.Options.TopP > 0 {
		chatReq.TopP = req.Options.TopP
	}
	if len(req.Options.StopWords) > 0 {
		chatReq.Stop = req.Options.StopWords
	}

	if len(req.Tools) > 0 {
		chatReq.Tools = make([]timewebTool, len(req.Tools))
		for i, t := range req.Tools {
			chatReq.Tools[i] = timewebTool{
				Type: "function",
				Function: timewebFunction{
					Name:        t.Name,
					Description: t.Description,
					Parameters:  t.Schema,
				},
			}
		}
	}

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

	url := fmt.Sprintf("%s/api/v1/cloud-ai/agents/%s/v1/chat/completions", c.baseURL, c.agentAccessID)
	httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
	httpReq.Header.Set("x-proxy-source", c.proxySource)

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		defer resp.Body.Close()
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("Timeweb API error: status %d, body: %s", resp.StatusCode, string(body))
	}

	ch := make(chan StreamChunk, 100)
	go func() {
		defer close(ch)
		defer resp.Body.Close()

		toolCalls := make(map[int]*ToolCall)
		reader := bufio.NewReader(resp.Body)

		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					return
				}
				if len(toolCalls) > 0 {
					calls := make([]ToolCall, 0, len(toolCalls))
					for _, tc := range toolCalls {
						calls = append(calls, *tc)
					}
					ch <- StreamChunk{ToolCallChunk: calls}
				}
				return
			}

			line = strings.TrimSpace(line)
			if line == "" {
				continue
			}

			if !strings.HasPrefix(line, "data: ") {
				continue
			}

			data := strings.TrimPrefix(line, "data: ")
			if data == "[DONE]" {
				if len(toolCalls) > 0 {
					calls := make([]ToolCall, 0, len(toolCalls))
					for _, tc := range toolCalls {
						calls = append(calls, *tc)
					}
					ch <- StreamChunk{ToolCallChunk: calls}
				}
				return
			}

			var streamResp timewebStreamResponse
			if err := json.Unmarshal([]byte(data), &streamResp); err != nil {
				continue
			}

			if len(streamResp.Choices) == 0 {
				continue
			}

			delta := streamResp.Choices[0].Delta

			if delta.Content != "" {
				ch <- StreamChunk{ContentChunk: delta.Content}
			}

			for _, tc := range delta.ToolCalls {
				idx := 0
				if _, ok := toolCalls[idx]; !ok {
					toolCalls[idx] = &ToolCall{
						ID:        tc.ID,
						Name:      tc.Function.Name,
						Arguments: make(map[string]interface{}),
					}
				}

				if tc.Function.Arguments != "" {
					existing := toolCalls[idx]
					var args map[string]interface{}
					if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err == nil {
						for k, v := range args {
							existing.Arguments[k] = v
						}
					}
				}
			}

			if streamResp.Choices[0].FinishReason != "" {
				ch <- StreamChunk{FinishReason: streamResp.Choices[0].FinishReason}
			}
		}
	}()

	return ch, nil
}

func (c *TimewebClient) GenerateText(ctx context.Context, req StreamRequest) (string, error) {
	messages := make([]timewebMessage, 0, len(req.Messages))
	for _, m := range req.Messages {
		msg := timewebMessage{
			Role:    string(m.Role),
			Content: m.Content,
		}
		if m.Name != "" {
			msg.Name = m.Name
		}
		if m.ToolCallID != "" {
			msg.ToolCallID = m.ToolCallID
		}
		messages = append(messages, msg)
	}

	chatReq := timewebChatRequest{
		Model:    c.modelKey,
		Messages: messages,
		Stream:   false,
	}

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

	if len(req.Tools) > 0 {
		chatReq.Tools = make([]timewebTool, len(req.Tools))
		for i, t := range req.Tools {
			chatReq.Tools[i] = timewebTool{
				Type: "function",
				Function: timewebFunction{
					Name:        t.Name,
					Description: t.Description,
					Parameters:  t.Schema,
				},
			}
		}
	}

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

	url := fmt.Sprintf("%s/api/v1/cloud-ai/agents/%s/v1/chat/completions", c.baseURL, c.agentAccessID)
	httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
	if err != nil {
		return "", fmt.Errorf("failed to create request: %w", err)
	}

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
	httpReq.Header.Set("x-proxy-source", c.proxySource)

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return "", fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("Timeweb API error: status %d, body: %s", resp.StatusCode, string(body))
	}

	var chatResp timewebChatResponse
	if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
		return "", fmt.Errorf("failed to decode response: %w", err)
	}

	if len(chatResp.Choices) == 0 {
		return "", errors.New("no choices in response")
	}

	return chatResp.Choices[0].Message.Content, nil
}