From a21e1213ccdd11d76f6338436c431e459ab8e574 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Fri, 26 Apr 2024 06:57:23 +0800
Subject: [PATCH] feat: support perplexity (#444)

---
 config.example.yaml      |  4 ++++
 src/client/mod.rs        |  1 +
 src/client/perplexity.rs | 19 +++++++++++++++++++
 3 files changed, 24 insertions(+)
 create mode 100644 src/client/perplexity.rs

diff --git a/config.example.yaml b/config.example.yaml
index 79e369a..9a30c3f 100644
--- a/config.example.yaml
+++ b/config.example.yaml
@@ -65,6 +65,10 @@ clients:
   - type: cohere
     api_key: xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
 
+  # See https://docs.perplexity.ai/docs/getting-started
+  - type: perplexity
+    api_key: pplx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+
   # See https://console.groq.com/docs/quickstart
   - type: groq
     api_key: gsk_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
diff --git a/src/client/mod.rs b/src/client/mod.rs
index ff7ea42..f7d448f 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -28,6 +28,7 @@ register_client!(
     (claude, "claude", ClaudeConfig, ClaudeClient),
     (mistral, "mistral", MistralConfig, MistralClient),
     (cohere, "cohere", CohereConfig, CohereClient),
+    (perplexity, "perplexity", PerplexityConfig, PerplexityClient),
     (groq, "groq", GroqConfig, GroqClient),
     (ollama, "ollama", OllamaConfig, OllamaClient),
     (ernie, "ernie", ErnieConfig, ErnieClient),
diff --git a/src/client/perplexity.rs b/src/client/perplexity.rs
new file mode 100644
index 0000000..99b2565
--- /dev/null
+++ b/src/client/perplexity.rs
@@ -0,0 +1,19 @@
+openai_compatible_module!(
+    PerplexityConfig,
+    PerplexityClient,
+    "https://api.perplexity.ai",
+    [
+        // https://docs.perplexity.ai/docs/model-cards
+        ("sonar-small-chat", "text", 16384),
+        ("sonar-small-online", "text", 12000),
+        ("sonar-medium-chat", "text", 16384),
+        ("sonar-medium-online", "text", 12000),
+
+        ("llama-3-8b-instruct", "text", 8192),
+        ("llama-3-70b-instruct", "text", 8192),
+        ("codellama-70b-instruct", "text", 16384),
+        ("mistral-7b-instruct", "text", 16384),
+        ("mixtral-8x7b-instruct", "text", 16384),
+        ("mixtral-8x22b-instruct", "text", 16384),
+    ]
+);