perplexity.go

package goai

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

func init() {
	register("perplexity", newPerplexityProvider)
}

// perplexityProvider implements Provider for Perplexity models
type perplexityProvider struct {
	apiKey      string
	modelID     string
	apiEndpoint string
}

// newPerplexityProvider creates a new Perplexity provider
func newPerplexityProvider(model string, opts Options) (Provider, error) {
	if opts.APIKey == "" {
		return nil, fmt.Errorf("%w: perplexity", ErrNoAPIKey)
	}
	
	endpoint := opts.Endpoint
	if endpoint == "" {
		endpoint = "https://api.perplexity.ai/chat/completions"
	}
	
	return &perplexityProvider{
		apiKey:      opts.APIKey,
		modelID:     model,
		apiEndpoint: endpoint,
	}, nil
}

// chat sends a prompt to Perplexity and returns the response
func (p *perplexityProvider) chat(ctx context.Context, prompt string, opts Options) (*Response, error) {
	start := time.Now()
	
	requestBody := map[string]interface{}{
		"model": p.modelID,
		"messages": []map[string]string{
			{"role": "user", "content": prompt},
		},
		"max_tokens":  opts.MaxTokens,
		"temperature": opts.Temperature,
	}
	
	jsonData, err := json.Marshal(requestBody)
	if err != nil {
		return nil, &Error{
			Provider: "perplexity",
			Message:  fmt.Sprintf("failed to marshal request: %v", err),
			Err:      err,
		}
	}
	
	req, err := http.NewRequestWithContext(ctx, "POST", p.apiEndpoint, bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, &Error{
			Provider: "perplexity",
			Message:  fmt.Sprintf("failed to create request: %v", err),
			Err:      err,
		}
	}
	
	req.Header.Set("Authorization", "Bearer "+p.apiKey)
	req.Header.Set("Content-Type", "application/json")
	
	client := &http.Client{Timeout: opts.Timeout}
	resp, err := client.Do(req)
	if err != nil {
		return nil, &Error{
			Provider: "perplexity",
			Message:  fmt.Sprintf("request failed: %v", err),
			Err:      err,
		}
	}
	defer resp.Body.Close()
	
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, &Error{
			Provider: "perplexity",
			Message:  fmt.Sprintf("failed to read response: %v", err),
			Err:      err,
		}
	}
	
	if resp.StatusCode != http.StatusOK {
		return nil, &Error{
			Provider:   "perplexity",
			StatusCode: resp.StatusCode,
			Message:    string(body),
		}
	}
	
	var apiResponse struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
		Usage struct {
			TotalTokens int `json:"total_tokens"`
		} `json:"usage"`
	}
	
	if err := json.Unmarshal(body, &apiResponse); err != nil {
		return nil, &Error{
			Provider: "perplexity",
			Message:  fmt.Sprintf("failed to parse response: %v", err),
			Err:      err,
		}
	}
	
	if len(apiResponse.Choices) == 0 {
		return nil, &Error{
			Provider: "perplexity",
			Message:  "no response from Perplexity",
			Err:      ErrNoResponse,
		}
	}
	
	return &Response{
		Content:    apiResponse.Choices[0].Message.Content,
		TokensUsed: apiResponse.Usage.TotalTokens,
		LatencyMs:  time.Since(start).Milliseconds(),
		Model:      p.modelID,
		Provider:   "perplexity",
		Metadata:   map[string]interface{}{"model": p.modelID},
	}, nil
}

// name returns the provider name
func (p *perplexityProvider) name() string {
	return "perplexity"
}

// model returns the model identifier
func (p *perplexityProvider) model() string {
	return p.modelID
}