Spaces:
Running
Running
File size: 3,468 Bytes
a2cbcac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | import os
from typing import Dict, Any, Optional
import aiohttp
from .utils import ToolResult
from ..config import config
PERPLEXITY_BASE_URL = "https://api.perplexity.ai/chat/completions"
async def research_with_perplexity(
symbol: Optional[str] = None,
query: Optional[str] = None,
model: str = "llama-3.1-sonar-small-128k-online",
max_tokens: int = 1000
) -> ToolResult:
"""
Research with Perplexity - backward compatibility function.
Args:
symbol: Stock symbol to research (optional)
query: Custom query string (optional)
model: Perplexity model to use
max_tokens: Maximum tokens in response
Returns:
ToolResult with research data
"""
if query:
research_text = query
elif symbol:
research_text = f"Latest developments and financial performance analysis for {symbol} stock"
else:
return ToolResult(success=False, error="Either symbol or query must be provided")
return await research_query(research_text, model, max_tokens)
async def research_query(
query: str,
model: str = "llama-3.1-sonar-small-128k-online",
max_tokens: int = 1000
) -> ToolResult:
"""
Perform research query using Perplexity API.
Args:
query: Research question or query
model: Perplexity model to use
max_tokens: Maximum tokens in response
Returns:
ToolResult with research response
"""
perplexity_key = config.get_api_key('perplexity')
if not perplexity_key:
return ToolResult(
success=False,
error="PERPLEXITY_API_KEY not found in environment variables"
)
try:
headers = {
"Authorization": f"Bearer {perplexity_key}",
"Content-Type": "application/json"
}
payload = {
"model": model,
"messages": [
{"role": "user", "content": query}
],
"max_tokens": max_tokens,
"temperature": 0.2,
"return_citations": True,
"return_images": False
}
async with aiohttp.ClientSession() as session:
async with session.post(PERPLEXITY_BASE_URL, json=payload, headers=headers) as response:
if response.status == 200:
data = await response.json()
# Extract response content
content = data.get('choices', [{}])[0].get('message', {}).get('content', '')
citations = data.get('citations', [])
return ToolResult(
success=True,
data={
'query': query,
'response': content,
'citations': citations,
'model': model,
'usage': data.get('usage', {})
}
)
else:
error_text = await response.text()
return ToolResult(
success=False,
error=f"HTTP {response.status}: {error_text}"
)
except Exception as e:
return ToolResult(
success=False,
error=f"Research query failed: {str(e)}"
)
|