Python Script for Perplexity API: Interactive Search Client with Streaming & History Support

General
,
1 
import os
import requests
import time
import json

class PerplexityAPI:
    def __init__(self, api_key, model="sonar"):
        self.api_key = api_key
        self.model = model
        self.api_url = "https://api.perplexity.ai/chat/completions"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        self.history = []

    def query(self, prompt, stream=False, temperature=0.7, max_tokens=1000):
        messages = [{"role": "user", "content": prompt}]
        messages = self.history[-10:] + messages
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "top_p": 0.9,
            "stream": stream,
            "return_citations": True
        }
        if stream:
            return self._stream_response(payload)
        else:
            return self._normal_response(payload)

    def _normal_response(self, payload):
        response = requests.post(self.api_url, headers=self.headers, json=payload)
        response.raise_for_status()
        data = response.json()
        self._save_history(payload["messages"][-1]["content"], data["choices"][0]["message"]["content"])
        return data

    def _stream_response(self, payload):
        response = requests.post(self.api_url, headers=self.headers, json=payload, stream=True)
        response.raise_for_status()
        full_text = ""
        for line in response.iter_lines():
            if line:
                line = line.decode("utf-8")
                if line.startswith("data: "):
                    data_str = line[6:]
                    if data_str == "[DONE]":
                        break
                    chunk = json.loads(data_str)
                    content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "")
                    print(content, end="", flush=True)
                    full_text += content
        print()
        self._save_history(payload["messages"][-1]["content"], full_text)
        return {"content": full_text}

    def _save_history(self, user_msg, assistant_msg):
        self.history.append({"role": "user", "content": user_msg})
        self.history.append({"role": "assistant", "content": assistant_msg})

if __name__ == "__main__":
    key = os.getenv("PERPLEXITY_API_KEY", "your_api_key_here")
    client = PerplexityAPI(key)
    while True:
        query = input("\nEnter your query (or 'exit' to quit): ").strip()
        if query.lower() == "exit":
            break
        try:
            print("Searching...")
            response = client.query(query)
            print("\nResponse:\n", response["choices"][0]["message"]["content"])
        except Exception as e:
            print("Error:", e)