Merge e77f5fe7c1 into 275655fd2e

examples/pipelines/providers/perplexity_manifold_pipeline.py

@@ -1,167 +1,184 @@
-from typing import List, Union, Generator, Iterator
+"""
-from pydantic import BaseModel
+title: Perplexity Pipeline
+author: Seyed Yahya Shirazi
+author_url: neuromechanist.github.io
+date: 2024-11-19
+version: 1.0
+license: MIT
+description: A pipeline for generating text using the Perplexity API.
+requirements: requests, sseclient-py
+environment_variables: PERPLEXITY_API_KEY
+"""
+
 import os
 import requests
+import json
+from typing import List, Union, Generator, Iterator
+from pydantic import BaseModel
+import sseclient
 
 from utils.pipelines.main import pop_system_message
 
 
 class Pipeline:
     class Valves(BaseModel):
-        PERPLEXITY_API_BASE_URL: str = "https://api.perplexity.ai"
         PERPLEXITY_API_KEY: str = ""
         pass
 
     def __init__(self):
         self.type = "manifold"
-        self.name = "Perplexity: "
+        self.id = "perplexity_manifold"
+        self.name = ""
 
         self.valves = self.Valves(
-            **{
+            **{"PERPLEXITY_API_KEY": os.getenv("PERPLEXITY_API_KEY", "your-api-key-here")}
-                "PERPLEXITY_API_KEY": os.getenv(
-                    "PERPLEXITY_API_KEY", "your-perplexity-api-key-here"
-                )
-            }
         )
+        self.url = 'https://api.perplexity.ai/chat/completions'
+        self.update_headers()
 
-        # Debugging: print the API key to ensure it's loaded
+    def update_headers(self):
-        print(f"Loaded API Key: {self.valves.PERPLEXITY_API_KEY}")
+        self.headers = {
+            'Authorization': f'Bearer {self.valves.PERPLEXITY_API_KEY}',
+            'Content-Type': 'application/json'
+        }
 
-        # List of models
+    def get_perplexity_models(self):
-        self.pipelines = [
+        return [
-            {
-                "id": "llama-3.1-sonar-large-128k-online",
-                "name": "Llama 3.1 Sonar Large 128k Online"
-            },
             {
                 "id": "llama-3.1-sonar-small-128k-online",
-                "name": "Llama 3.1 Sonar Small 128k Online"
+                "name": "Perplexity Llama 3.1 Sonar Small"
             },
             {
-                "id": "llama-3.1-sonar-large-128k-chat",
+                "id": "llama-3.1-sonar-large-128k-online",
-                "name": "Llama 3.1 Sonar Large 128k Chat"
+                "name": "Perplexity Llama 3.1 Sonar Large"
             },
             {
-                "id": "llama-3.1-sonar-small-128k-chat",
+                "id": "llama-3.1-sonar-huge-128k-online",
-                "name": "Llama 3.1 Sonar Small 128k Chat"
+                "name": "Perplexity Llama 3.1 Sonar Huge"
             },
-            {
-                "id": "llama-3.1-8b-instruct", "name": "Llama 3.1 8B Instruct"
-            },
-            {
-                "id": "llama-3.1-70b-instruct", "name": "Llama 3.1 70B Instruct"
-            }
         ]
         pass
 
     async def on_startup(self):
-        # This function is called when the server is started.
         print(f"on_startup:{__name__}")
         pass
 
     async def on_shutdown(self):
-        # This function is called when the server is stopped.
         print(f"on_shutdown:{__name__}")
         pass
 
     async def on_valves_updated(self):
-        # This function is called when the valves are updated.
+        self.update_headers()
-        print(f"on_valves_updated:{__name__}")
+
-        # No models to fetch, static setup
+    def pipelines(self) -> List[dict]:
-        pass
+        return self.get_perplexity_models()
 
     def pipe(
         self, user_message: str, model_id: str, messages: List[dict], body: dict
     ) -> Union[str, Generator, Iterator]:
-        # This is where you can add your custom pipelines like RAG.
-        print(f"pipe:{__name__}")
-
-        system_message, messages = pop_system_message(messages)
-        system_prompt = "You are a helpful assistant."
-        if system_message is not None:
-            system_prompt = system_message["content"]
-
-        print(system_prompt)
-        print(messages)
-        print(user_message)
-
-        headers = {
-            "Authorization": f"Bearer {self.valves.PERPLEXITY_API_KEY}",
-            "Content-Type": "application/json",
-            "accept": "application/json"
-        }
-
-        payload = {
-            "model": model_id,
-            "messages": [
-                {"role": "system", "content": system_prompt},
-                *messages
-            ],
-            "stream": body.get("stream", True),
-            "return_citations": True,
-            "return_images": True
-        }
-
-        if "user" in payload:
-            del payload["user"]
-        if "chat_id" in payload:
-            del payload["chat_id"]
-        if "title" in payload:
-            del payload["title"]
-
-        print(payload)
-
         try:
-            r = requests.post(
+            # Remove unnecessary keys
-                url=f"{self.valves.PERPLEXITY_API_BASE_URL}/chat/completions",
+            for key in ['user', 'chat_id', 'title']:
-                json=payload,
+                body.pop(key, None)
-                headers=headers,
-                stream=True,
-            )
 
-            r.raise_for_status()
+            system_message, messages = pop_system_message(messages)
+
+            # Add system message as the first message if present
+            processed_messages = []
+            if system_message:
+                processed_messages.append({
+                    "role": "system",
+                    "content": str(system_message)
+                })
+
+            # Process remaining messages
+            for message in messages:
+                content = (
+                    message["content"][0]["text"]
+                    if isinstance(message["content"], list)
+                    else message["content"]
+                )
+                processed_messages.append({
+                    "role": message["role"],
+                    "content": content
+                })
+
+            # Prepare the payload with Perplexity-specific parameters
+            payload = {
+                "model": model_id,
+                "messages": processed_messages,
+                "max_tokens": body.get("max_tokens", None),  # Optional in Perplexity
+                "temperature": body.get("temperature", 0.2),  # Perplexity default
+                "top_p": body.get("top_p", 0.9),  # Perplexity default
+                "top_k": body.get("top_k", 0),    # Perplexity default
+                "stream": body.get("stream", False),
+                "presence_penalty": body.get("presence_penalty", 0),
+                "frequency_penalty": body.get("frequency_penalty", 1),
+            }
+
+            # Add Perplexity-specific features if specified
+            if "search_domain_filter" in body:
+                payload["search_domain_filter"] = body["search_domain_filter"]
+            if "return_images" in body:
+                payload["return_images"] = body["return_images"]
+            if "return_related_questions" in body:
+                payload["return_related_questions"] = body["return_related_questions"]
+            if "search_recency_filter" in body:
+                payload["search_recency_filter"] = body["search_recency_filter"]
 
             if body.get("stream", False):
-                return r.iter_lines()
+                return self.stream_response(payload)
             else:
-                response = r.json()
+                return self.get_completion(payload)
-                formatted_response = {
-                    "id": response["id"],
-                    "model": response["model"],
-                    "created": response["created"],
-                    "usage": response["usage"],
-                    "object": response["object"],
-                    "choices": [
-                        {
-                            "index": choice["index"],
-                            "finish_reason": choice["finish_reason"],
-                            "message": {
-                                "role": choice["message"]["role"],
-                                "content": choice["message"]["content"]
-                            },
-                            "delta": {"role": "assistant", "content": ""}
-                        } for choice in response["choices"]
-                    ]
-                }
-                return formatted_response
         except Exception as e:
             return f"Error: {e}"
 
+    def format_response_with_citations(self, content: str, citations: List[str]) -> str:
+        """Format the response by appending citations at the end."""
+        if not citations:
+            return content
 
-if __name__ == "__main__":
+        # Content already contains [1], [2], etc. references
-    import argparse
+        formatted_response = content + "\n\nReferences:\n"
+        for i, url in enumerate(citations, 1):
+            formatted_response += f"[{i}] {url}\n"
+        return formatted_response
 
-    parser = argparse.ArgumentParser(description="Perplexity API Client")
+    def stream_response(self, payload: dict) -> Generator:
-    parser.add_argument("--api-key", type=str, required=True,
+        response = requests.post(self.url, headers=self.headers, json=payload, stream=True)
-                        help="API key for Perplexity")
+        accumulated_content = ""
-    parser.add_argument("--prompt", type=str, required=True,
-                        help="Prompt to send to the Perplexity API")
 
-    args = parser.parse_args()
+        if response.status_code == 200:
+            client = sseclient.SSEClient(response)
+            citations = None
+            for event in client.events():
+                try:
+                    data = json.loads(event.data)
+                    if "citations" in data:
+                        citations = data["citations"]
+                    if data["choices"][0]["finish_reason"] is None:
+                        content = data["choices"][0]["delta"]["content"]
+                        accumulated_content += content
+                        yield content
+                    elif data["choices"][0]["finish_reason"] == "stop" and citations:
+                        yield "\n\nReferences:\n" + "\n".join(
+                            f"[{i}] {url}" for i, url in enumerate(citations, 1)
+                        )
+                        break
+                except json.JSONDecodeError:
+                    print(f"Failed to parse JSON: {event.data}")
+                except KeyError as e:
+                    print(f"Unexpected data structure: {e}")
+                    print(f"Full data: {data}")
+        else:
+            raise Exception(f"Error: {response.status_code} - {response.text}")
 
-    pipeline = Pipeline()
+    def get_completion(self, payload: dict) -> str:
-    pipeline.valves.PERPLEXITY_API_KEY = args.api_key
+        response = requests.post(self.url, headers=self.headers, json=payload)
-    response = pipeline.pipe(
+        if response.status_code == 200:
-        user_message=args.prompt, model_id="llama-3-sonar-large-32k-online", messages=[], body={"stream": False})
+            res = response.json()
-
+            content = res["choices"][0]["message"]["content"]
-    print("Response:", response)
+            citations = res.get("citations", [])
+            return self.format_response_with_citations(content, citations)
+        else:
+            raise Exception(f"Error: {response.status_code} - {response.text}")