client.py•5.15 kB
"""Perplexity API client for MCP integration."""
import logging
import re
import sys
from typing import Dict, List, Optional, Any
import httpx
from config import PERPLEXITY_BASE_URL, PERPLEXITY_TIMEOUT, get_api_key
# Configure logging to stderr
logger = logging.getLogger(__name__)
class PerplexityClient:
"""Client for interacting with Perplexity API."""
def __init__(self):
"""Initialize the Perplexity client."""
self.api_key = get_api_key()
self.base_url = PERPLEXITY_BASE_URL
self.timeout = PERPLEXITY_TIMEOUT
self.headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
def chat_completion(self,messages: List[Dict[str, str]],model: str,**kwargs) -> Dict[str, Any]:
"""
Send a chat completion request to Perplexity API.
Args:
messages: List of message objects with role and content
model: Model name (e.g., "sonar-pro", "sonar-reasoning-pro", "sonar-deep-research")
**kwargs: Additional parameters for the API request
Returns:
Dict containing the API response
"""
try:
# Prepare request payload
payload = {
"model": model,
"messages": messages,
**kwargs
}
# Log request details to stderr
logger.info(f"Making request to Perplexity API with model: {model}")
logger.debug(f"Request payload: {payload}")
# Make the API request
with httpx.Client(timeout=self.timeout) as client:
response = client.post(
f"{self.base_url}/chat/completions",
headers=self.headers,
json=payload
)
# Check for HTTP errors
response.raise_for_status()
# Parse response
result = response.json()
# Log response details to stderr
logger.info(f"Received response from Perplexity API")
if "usage" in result:
usage = result["usage"]
logger.info(f"Token usage - Prompt: {usage.get('prompt_tokens', 0)}, "
f"Completion: {usage.get('completion_tokens', 0)}, "
f"Total: {usage.get('total_tokens', 0)}")
return result
except httpx.HTTPStatusError as e:
logger.error(f"HTTP error from Perplexity API: {e.response.status_code}")
logger.error(f"Response content: {e.response.text}")
return {
"error": f"HTTP {e.response.status_code}",
"message": f"API request failed: {e.response.text}"
}
except httpx.TimeoutException:
logger.error(f"Request timeout after {self.timeout} seconds")
return {
"error": "timeout",
"message": f"Request timed out after {self.timeout} seconds"
}
except Exception as e:
logger.exception("Unexpected error in chat_completion")
return {
"error": "unexpected_error",
"message": f"Unexpected error: {str(e)}"
}
def format_response(self, api_response: Dict[str, Any]) -> Dict[str, Any]:
"""
Format API response for MCP tool return.
Args:
api_response: Raw API response from Perplexity
Returns:
Formatted response for MCP tool with only content and citations
"""
# Handle error responses
if "error" in api_response:
return api_response
try:
# Extract main content
content = ""
if "choices" in api_response and api_response["choices"]:
content = api_response["choices"][0]["message"]["content"]
# Remove <think>...</think> sections for reasoning models
# This removes the thinking tokens that appear in medium/large responses
content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL)
# Clean up any extra whitespace left after removing think tags
content = re.sub(r'\n\s*\n\s*\n', '\n\n', content)
content = content.strip()
# Format response - only include content and citations
formatted = {
"content": content,
"citations": api_response.get("citations", [])
}
return formatted
except Exception as e:
logger.exception("Error formatting API response")
return {
"error": "format_error",
"message": f"Failed to format response: {str(e)}"
}