Perplexity api (#1432)

I have added Perplexity to the default AI models. I see Anthropic models are becoming part of the default as well, so I thought I should add a model that is specific for web search. This pull request is a work in progress; reviews and edit recommendations are welcome. --------- Co-authored-by: sawka <mike@commandline.dev>
2024-12-21 16:38:23 +01:00 · 2024-12-09 15:48:33 -08:00 · 2024-12-09 15:48:33 -08:00 · c071cc04c3
commit c071cc04c3
parent b706d4524b
4 changed files with 248 additions and 28 deletions
--- a/docs/docs/faq.mdx
+++ b/docs/docs/faq.mdx
@ -66,6 +66,23 @@ Set these keys:

 Note: we set the `ai:*` key to true to clear all the existing "ai" keys so this config is a clean slate.

+### How can I connect to Perplexity?
+
+Open your [config file](./config) in Wave using `wsh editconfig`.
+
+Set these keys:
+
+```json
+{
+  "ai:*": true,
+  "ai:apitype": "perplexity",
+  "ai:model": "llama-3.1-sonar-small-128k-online",
+  "ai:apitoken": "<your perplexity API key>"
+}
+```
+
+Note: we set the `ai:*` key to true to clear all the existing "ai" keys so this config is a clean slate.
+
 To switch between models, consider [adding AI Presets](./presets) instead.

 ### How can I see the block numbers?
--- a/frontend/app/view/waveai/waveai.tsx
+++ b/frontend/app/view/waveai/waveai.tsx
@ -180,40 +180,54 @@ export class WaveAiModel implements ViewModel {
            const presetKey = get(this.presetKey);
            const presetName = presets[presetKey]?.["display:name"] ?? "";
            const isCloud = isBlank(aiOpts.apitoken) && isBlank(aiOpts.baseurl);
-            if (aiOpts?.apitype == "anthropic") {
-                const modelName = aiOpts.model;
-                viewTextChildren.push({
-                    elemtype: "iconbutton",
-                    icon: "globe",
-                    title: "Using Remote Antropic API (" + modelName + ")",
-                    noAction: true,
-                });
-            } else if (isCloud) {
-                viewTextChildren.push({
-                    elemtype: "iconbutton",
-                    icon: "cloud",
-                    title: "Using Wave's AI Proxy (gpt-4o-mini)",
-                    noAction: true,
-                });
-            } else {
-                const baseUrl = aiOpts.baseurl ?? "OpenAI Default Endpoint";
-                const modelName = aiOpts.model;
-                if (baseUrl.startsWith("http://localhost") || baseUrl.startsWith("http://127.0.0.1")) {
-                    viewTextChildren.push({
-                        elemtype: "iconbutton",
-                        icon: "location-dot",
-                        title: "Using Local Model @ " + baseUrl + " (" + modelName + ")",
-                        noAction: true,
-                    });
-                } else {
+
+            // Handle known API providers
+            switch (aiOpts?.apitype) {
+                case "anthropic":
                    viewTextChildren.push({
                        elemtype: "iconbutton",
                        icon: "globe",
-                        title: "Using Remote Model @ " + baseUrl + " (" + modelName + ")",
+                        title: `Using Remote Anthropic API (${aiOpts.model})`,
                        noAction: true,
                    });
-                }
+                    break;
+                case "perplexity":
+                    viewTextChildren.push({
+                        elemtype: "iconbutton",
+                        icon: "globe",
+                        title: `Using Remote Perplexity API (${aiOpts.model})`,
+                        noAction: true,
+                    });
+                    break;
+                default:
+                    if (isCloud) {
+                        viewTextChildren.push({
+                            elemtype: "iconbutton",
+                            icon: "cloud",
+                            title: "Using Wave's AI Proxy (gpt-4o-mini)",
+                            noAction: true,
+                        });
+                    } else {
+                        const baseUrl = aiOpts.baseurl ?? "OpenAI Default Endpoint";
+                        const modelName = aiOpts.model;
+                        if (baseUrl.startsWith("http://localhost") || baseUrl.startsWith("http://127.0.0.1")) {
+                            viewTextChildren.push({
+                                elemtype: "iconbutton",
+                                icon: "location-dot",
+                                title: `Using Local Model @ ${baseUrl} (${modelName})`,
+                                noAction: true,
+                            });
+                        } else {
+                            viewTextChildren.push({
+                                elemtype: "iconbutton",
+                                icon: "globe",
+                                title: `Using Remote Model @ ${baseUrl} (${modelName})`,
+                                noAction: true,
+                            });
+                        }
+                    }
            }
+
            const dropdownItems = Object.entries(presets)
                .sort((a, b) => ((a[1]["display:order"] ?? 0) > (b[1]["display:order"] ?? 0) ? 1 : -1))
                .map(
--- a/pkg/waveai/perplexitybackend.go
+++ b/pkg/waveai/perplexitybackend.go
@ -0,0 +1,179 @@
+// Copyright 2024, Command Line Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package waveai
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/wavetermdev/waveterm/pkg/panichandler"
+	"github.com/wavetermdev/waveterm/pkg/wshrpc"
+)
+
+type PerplexityBackend struct{}
+
+var _ AIBackend = PerplexityBackend{}
+
+// Perplexity API request types
+type perplexityMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type perplexityRequest struct {
+	Model    string              `json:"model"`
+	Messages []perplexityMessage `json:"messages"`
+	Stream   bool                `json:"stream"`
+}
+
+// Perplexity API response types
+type perplexityResponseDelta struct {
+	Content string `json:"content"`
+}
+
+type perplexityResponseChoice struct {
+	Delta        perplexityResponseDelta `json:"delta"`
+	FinishReason string                  `json:"finish_reason"`
+}
+
+type perplexityResponse struct {
+	ID      string                     `json:"id"`
+	Choices []perplexityResponseChoice `json:"choices"`
+	Model   string                     `json:"model"`
+}
+
+func (PerplexityBackend) StreamCompletion(ctx context.Context, request wshrpc.OpenAiStreamRequest) chan wshrpc.RespOrErrorUnion[wshrpc.OpenAIPacketType] {
+	rtn := make(chan wshrpc.RespOrErrorUnion[wshrpc.OpenAIPacketType])
+
+	go func() {
+		defer func() {
+			panicErr := panichandler.PanicHandler("PerplexityBackend.StreamCompletion")
+			if panicErr != nil {
+				rtn <- makeAIError(panicErr)
+			}
+			close(rtn)
+		}()
+
+		if request.Opts == nil {
+			rtn <- makeAIError(errors.New("no perplexity opts found"))
+			return
+		}
+
+		model := request.Opts.Model
+		if model == "" {
+			model = "llama-3.1-sonar-small-128k-online"
+		}
+
+		// Convert messages format
+		var messages []perplexityMessage
+		for _, msg := range request.Prompt {
+			role := "user"
+			if msg.Role == "assistant" {
+				role = "assistant"
+			} else if msg.Role == "system" {
+				role = "system"
+			}
+
+			messages = append(messages, perplexityMessage{
+				Role:    role,
+				Content: msg.Content,
+			})
+		}
+
+		perplexityReq := perplexityRequest{
+			Model:    model,
+			Messages: messages,
+			Stream:   true,
+		}
+
+		reqBody, err := json.Marshal(perplexityReq)
+		if err != nil {
+			rtn <- makeAIError(fmt.Errorf("failed to marshal perplexity request: %v", err))
+			return
+		}
+
+		req, err := http.NewRequestWithContext(ctx, "POST", "https://api.perplexity.ai/chat/completions", strings.NewReader(string(reqBody)))
+		if err != nil {
+			rtn <- makeAIError(fmt.Errorf("failed to create perplexity request: %v", err))
+			return
+		}
+
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("Authorization", "Bearer "+request.Opts.APIToken)
+
+		client := &http.Client{}
+		resp, err := client.Do(req)
+		if err != nil {
+			rtn <- makeAIError(fmt.Errorf("failed to send perplexity request: %v", err))
+			return
+		}
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusOK {
+			bodyBytes, _ := io.ReadAll(resp.Body)
+			rtn <- makeAIError(fmt.Errorf("Perplexity API error: %s - %s", resp.Status, string(bodyBytes)))
+			return
+		}
+
+		reader := bufio.NewReader(resp.Body)
+		sentHeader := false
+
+		for {
+			// Check for context cancellation
+			select {
+			case <-ctx.Done():
+				rtn <- makeAIError(fmt.Errorf("request cancelled: %v", ctx.Err()))
+				return
+			default:
+			}
+
+			line, err := reader.ReadString('\n')
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				rtn <- makeAIError(fmt.Errorf("error reading stream: %v", err))
+				break
+			}
+
+			line = strings.TrimSpace(line)
+			if !strings.HasPrefix(line, "data: ") {
+				continue
+			}
+
+			data := strings.TrimPrefix(line, "data: ")
+			if data == "[DONE]" {
+				break
+			}
+
+			var response perplexityResponse
+			if err := json.Unmarshal([]byte(data), &response); err != nil {
+				rtn <- makeAIError(fmt.Errorf("error unmarshaling response: %v", err))
+				break
+			}
+
+			if !sentHeader {
+				pk := MakeOpenAIPacket()
+				pk.Model = response.Model
+				rtn <- wshrpc.RespOrErrorUnion[wshrpc.OpenAIPacketType]{Response: *pk}
+				sentHeader = true
+			}
+
+			for _, choice := range response.Choices {
+				pk := MakeOpenAIPacket()
+				pk.Text = choice.Delta.Content
+				pk.FinishReason = choice.FinishReason
+				rtn <- wshrpc.RespOrErrorUnion[wshrpc.OpenAIPacketType]{Response: *pk}
+			}
+		}
+	}()
+
+	return rtn
+}
--- a/pkg/waveai/waveai.go
+++ b/pkg/waveai/waveai.go
@ -17,6 +17,7 @@ const OpenAICloudReqStr = "openai-cloudreq"
 const PacketEOFStr = "EOF"
 const DefaultAzureAPIVersion = "2023-05-15"
 const ApiType_Anthropic = "anthropic"
+const ApiType_Perplexity = "perplexity"

 type OpenAICmdInfoPacketOutputType struct {
 	Model        string `json:"model,omitempty"`
@ -74,6 +75,15 @@ func RunAICommand(ctx context.Context, request wshrpc.OpenAiStreamRequest) chan
 		anthropicBackend := AnthropicBackend{}
 		return anthropicBackend.StreamCompletion(ctx, request)
 	}
+	if request.Opts.APIType == ApiType_Perplexity {
+		endpoint := request.Opts.BaseURL
+		if endpoint == "" {
+			endpoint = "default"
+		}
+		log.Printf("sending ai chat message to perplexity endpoint %q using model %s\n", endpoint, request.Opts.Model)
+		perplexityBackend := PerplexityBackend{}
+		return perplexityBackend.StreamCompletion(ctx, request)
+	}
 	if IsCloudAIRequest(request.Opts) {
 		log.Print("sending ai chat message to default waveterm cloud endpoint\n")
 		cloudBackend := WaveAICloudBackend{}