responseHandler.go

package deepseek

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

// ChatCompletionResponse represents a response from the chat completion endpoint.
type ChatCompletionResponse struct {
	ID                string   `json:"id"`                           // Unique identifier for the chat completion.
	Object            string   `json:"object"`                       // Type of the object, typically "chat.completion".
	Created           int64    `json:"created"`                      // Timestamp when the chat completion was created.
	Model             string   `json:"model"`                        // The model used for generating the completion.
	Choices           []Choice `json:"choices"`                      // List of completion choices generated by the model.
	Usage             Usage    `json:"usage"`                        // Token usage statistics.
	SystemFingerprint *string  `json:"system_fingerprint,omitempty"` // Fingerprint of the system configuration.
}

// Choice represents a completion choice generated by the model.
type Choice struct {
	Index        int       `json:"index"`              // Index of the choice in the list of choices.
	Message      Message   `json:"message"`            // The message generated by the model.
	Logprobs     *Logprobs `json:"logprobs,omitempty"` // Log probabilities of the tokens, if available.
	FinishReason string    `json:"finish_reason"`      // Reason why the completion finished.
}

// ToolCallFunction represents a function call in the tool.
type ToolCallFunction struct {
	Name      string `json:"name"`      // Name of the function (required)
	Arguments string `json:"arguments"` // JSON string of arguments passed to the function (required)
}

// ToolCall represents a tool call in the completion.
type ToolCall struct {
	Index    int              `json:"index"`    // Index of the tool call
	ID       string           `json:"id"`       // Unique identifier for the tool call
	Type     string           `json:"type"`     // Type of the tool call, e.g., "function"
	Function ToolCallFunction `json:"function"` // The function details for the call
}

// Message represents a message generated by the model.
type Message struct {
	Role             string     `json:"role"`                        // Role of the message sender (e.g., "user", "assistant").
	Content          string     `json:"content"`                     // Content of the message.
	ReasoningContent string     `json:"reasoning_content,omitempty"` // Optional reasoning content.
	ToolCalls        []ToolCall `json:"tool_calls,omitempty"`        // Optional tool calls.
}

// Logprobs represents log probability information for a choice or token.
type Logprobs struct {
	Content     []ContentToken    `json:"content"`      // A list of message content tokens with log probability information.
	TopLogprobs []TopLogprobToken `json:"top_logprobs"` // List of the most likely tokens and their log probability.
}

// ContentToken represents a single token within the content with its log probability and byte information.
type ContentToken struct {
	Token   string  `json:"token"`           // The token string.
	Logprob float64 `json:"logprob"`         // The log probability of this token. -9999.0 if not in top 20.
	Bytes   []int   `json:"bytes,omitempty"` // UTF-8 byte representation of the token. Can be nil.
}

// TopLogprobToken represents a single token within the top log probabilities with its log probability and byte information.
type TopLogprobToken struct {
	Token   string  `json:"token"`           // The token string.
	Logprob float64 `json:"logprob"`         // The log probability of this token. -9999.0 if not in top 20.
	Bytes   []int   `json:"bytes,omitempty"` // UTF-8 byte representation of the token. Can be nil.
}

// Usage represents token usage statistics.
type Usage struct {
	PromptTokens          int `json:"prompt_tokens"`            // Number of tokens used in the prompt.
	CompletionTokens      int `json:"completion_tokens"`        // Number of tokens used in the completion.
	TotalTokens           int `json:"total_tokens"`             // Total number of tokens used.
	PromptCacheHitTokens  int `json:"prompt_cache_hit_tokens"`  // Number of tokens served from cache.
	PromptCacheMissTokens int `json:"prompt_cache_miss_tokens"` // Number of tokens not served from cache.
}

// HandleChatCompletionResponse parses the response from the chat completion endpoint.
func HandleChatCompletionResponse(resp *http.Response) (*ChatCompletionResponse, error) {
	body, err := io.ReadAll(resp.Body) //Do not re read the body hereafter.
	if err != nil {
		return nil, fmt.Errorf("failed to read response body: %w", err)
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			fmt.Printf("failed to close response body: %s\n", err)
		}
	}()

	var parsedResponse ChatCompletionResponse
	if err := json.Unmarshal(body, &parsedResponse); err != nil {
		return nil, handleAPIError(body)
	}

	if err := validateChatCompletionResponse(&parsedResponse); err != nil {
		return nil, fmt.Errorf("invalid response: %w", err)
	}
	return &parsedResponse, nil
}

// HandleFIMCompletionRequest parses the response from the FIM completion endpoint.
func HandleFIMCompletionRequest(resp *http.Response) (*FIMCompletionResponse, error) {
	body, err := io.ReadAll(resp.Body) //Do not re read the body hereafter.
	if err != nil {
		return nil, fmt.Errorf("failed to read response body: %w", err)
	}
	defer func() {
		if err := resp.Body.Close(); err != nil {
			fmt.Printf("failed to close response body: %s\n", err)
		}
	}()

	var parsedResponse FIMCompletionResponse
	if err := json.Unmarshal(body, &parsedResponse); err != nil {
		return nil, handleAPIError(body)
	}
	return &parsedResponse, nil
}

// handleAPIError handles API errors by parsing the response body.
func handleAPIError(body []byte) error {
	responseBody := string(body)

	if len(responseBody) == 0 {
		return fmt.Errorf("failed to parse response JSON: empty response body")
	}
	if strings.HasPrefix(responseBody, "<!DOCTYPE html>") {
		return fmt.Errorf("unexpected HTML response (model may not exist). This is likely an issue with the how some external servers return html responses for error. Make sure you are calling the right path or models")
	}
	return fmt.Errorf("failed to parse response JSON: unexpected end of JSON input. %s", responseBody)
}

func validateChatCompletionResponse(parsedResponse *ChatCompletionResponse) error {
	if parsedResponse == nil {
		return fmt.Errorf("nil response")
	}

	// Validate required fields
	if parsedResponse.ID == "" {
		return fmt.Errorf("missing response ID")
	}

	if len(parsedResponse.Choices) == 0 {
		return fmt.Errorf("no choices in response")
	}

	return nil
}