airtrain 0.1.58__py3-none-any.whl → 0.1.62__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- airtrain/__init__.py +72 -44
- airtrain/__pycache__/__init__.cpython-313.pyc +0 -0
- airtrain/core/__pycache__/__init__.cpython-313.pyc +0 -0
- airtrain/core/__pycache__/schemas.cpython-313.pyc +0 -0
- airtrain/core/__pycache__/skills.cpython-313.pyc +0 -0
- airtrain/core/credentials.py +59 -13
- airtrain/integrations/__init__.py +21 -2
- airtrain/integrations/combined/list_models_factory.py +80 -41
- airtrain/integrations/perplexity/__init__.py +49 -0
- airtrain/integrations/perplexity/credentials.py +43 -0
- airtrain/integrations/perplexity/list_models.py +112 -0
- airtrain/integrations/perplexity/models_config.py +128 -0
- airtrain/integrations/perplexity/skills.py +279 -0
- airtrain/integrations/search/__init__.py +21 -0
- airtrain/integrations/search/exa/__init__.py +23 -0
- airtrain/integrations/search/exa/credentials.py +30 -0
- airtrain/integrations/search/exa/schemas.py +114 -0
- airtrain/integrations/search/exa/skills.py +114 -0
- airtrain/tools/__init__.py +9 -5
- airtrain/tools/command.py +248 -61
- airtrain/tools/search.py +450 -0
- airtrain/tools/testing.py +135 -0
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/METADATA +1 -1
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/RECORD +27 -15
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/WHEEL +1 -1
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/entry_points.txt +0 -0
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,43 @@
|
|
1
|
+
from pydantic import Field, SecretStr
|
2
|
+
from airtrain.core.credentials import BaseCredentials, CredentialValidationError
|
3
|
+
import requests
|
4
|
+
|
5
|
+
|
6
|
+
class PerplexityCredentials(BaseCredentials):
|
7
|
+
"""Perplexity AI API credentials"""
|
8
|
+
|
9
|
+
perplexity_api_key: SecretStr = Field(..., description="Perplexity AI API key")
|
10
|
+
|
11
|
+
_required_credentials = {"perplexity_api_key"}
|
12
|
+
|
13
|
+
async def validate_credentials(self) -> bool:
|
14
|
+
"""Validate Perplexity AI credentials by making a test API call"""
|
15
|
+
try:
|
16
|
+
headers = {
|
17
|
+
"Authorization": f"Bearer {self.perplexity_api_key.get_secret_value()}",
|
18
|
+
"Content-Type": "application/json",
|
19
|
+
}
|
20
|
+
|
21
|
+
# Small API call to check if credentials are valid
|
22
|
+
data = {
|
23
|
+
"model": "sonar-pro",
|
24
|
+
"messages": [{"role": "user", "content": "Test"}],
|
25
|
+
"max_tokens": 1,
|
26
|
+
}
|
27
|
+
|
28
|
+
# Make a synchronous request for validation
|
29
|
+
response = requests.post(
|
30
|
+
"https://api.perplexity.ai/chat/completions", headers=headers, json=data
|
31
|
+
)
|
32
|
+
|
33
|
+
if response.status_code == 200:
|
34
|
+
return True
|
35
|
+
else:
|
36
|
+
raise CredentialValidationError(
|
37
|
+
f"Invalid Perplexity AI credentials: {response.status_code} - {response.text}"
|
38
|
+
)
|
39
|
+
|
40
|
+
except Exception as e:
|
41
|
+
raise CredentialValidationError(
|
42
|
+
f"Invalid Perplexity AI credentials: {str(e)}"
|
43
|
+
)
|
@@ -0,0 +1,112 @@
|
|
1
|
+
from typing import Dict, Any, List, Optional
|
2
|
+
import requests
|
3
|
+
|
4
|
+
from airtrain.core.skills import Skill, ProcessingError
|
5
|
+
from airtrain.integrations.combined.list_models_factory import (
|
6
|
+
BaseListModelsSkill,
|
7
|
+
GenericListModelsInput,
|
8
|
+
GenericListModelsOutput,
|
9
|
+
)
|
10
|
+
from airtrain.core.schemas import InputSchema, OutputSchema
|
11
|
+
from .credentials import PerplexityCredentials
|
12
|
+
from .models_config import PERPLEXITY_MODELS_CONFIG
|
13
|
+
|
14
|
+
|
15
|
+
class PerplexityListModelsInput(InputSchema):
|
16
|
+
"""Schema for listing Perplexity AI models"""
|
17
|
+
|
18
|
+
api_models_only: bool = False
|
19
|
+
|
20
|
+
|
21
|
+
class PerplexityListModelsOutput(OutputSchema):
|
22
|
+
"""Schema for Perplexity AI models listing output"""
|
23
|
+
|
24
|
+
models: List[Dict[str, Any]]
|
25
|
+
provider: str = "perplexity"
|
26
|
+
|
27
|
+
|
28
|
+
class PerplexityListModelsSkill(BaseListModelsSkill):
|
29
|
+
"""Skill for listing Perplexity AI models"""
|
30
|
+
|
31
|
+
def __init__(self, credentials: Optional[PerplexityCredentials] = None):
|
32
|
+
"""Initialize the skill with optional credentials"""
|
33
|
+
super().__init__(provider="perplexity", credentials=credentials)
|
34
|
+
self.credentials = credentials
|
35
|
+
|
36
|
+
def get_models(self) -> List[Dict[str, Any]]:
|
37
|
+
"""Return list of Perplexity AI models."""
|
38
|
+
models = []
|
39
|
+
|
40
|
+
# Add models from the configuration
|
41
|
+
for model_id, config in PERPLEXITY_MODELS_CONFIG.items():
|
42
|
+
models.append(
|
43
|
+
{
|
44
|
+
"id": model_id,
|
45
|
+
"display_name": config["name"],
|
46
|
+
"description": config.get("description", ""),
|
47
|
+
"category": config.get("category", "unknown"),
|
48
|
+
"capabilities": {
|
49
|
+
"citations": config.get("citations", False),
|
50
|
+
"search": config.get("search", False),
|
51
|
+
"context_window": config.get("context_window", 8192),
|
52
|
+
"max_completion_tokens": config.get(
|
53
|
+
"max_completion_tokens", 4096
|
54
|
+
),
|
55
|
+
},
|
56
|
+
}
|
57
|
+
)
|
58
|
+
|
59
|
+
return models
|
60
|
+
|
61
|
+
def process(self, input_data: GenericListModelsInput) -> GenericListModelsOutput:
|
62
|
+
"""Process the input and return a list of models."""
|
63
|
+
try:
|
64
|
+
models = self.get_models()
|
65
|
+
return GenericListModelsOutput(models=models, provider="perplexity")
|
66
|
+
except Exception as e:
|
67
|
+
raise ProcessingError(f"Failed to list Perplexity AI models: {str(e)}")
|
68
|
+
|
69
|
+
|
70
|
+
# Standalone version directly using the Perplexity-specific schemas
|
71
|
+
class StandalonePerplexityListModelsSkill(
|
72
|
+
Skill[PerplexityListModelsInput, PerplexityListModelsOutput]
|
73
|
+
):
|
74
|
+
"""Standalone skill for listing Perplexity AI models"""
|
75
|
+
|
76
|
+
input_schema = PerplexityListModelsInput
|
77
|
+
output_schema = PerplexityListModelsOutput
|
78
|
+
|
79
|
+
def __init__(self, credentials: Optional[PerplexityCredentials] = None):
|
80
|
+
"""Initialize the skill with optional credentials"""
|
81
|
+
super().__init__()
|
82
|
+
self.credentials = credentials
|
83
|
+
|
84
|
+
def process(
|
85
|
+
self, input_data: PerplexityListModelsInput
|
86
|
+
) -> PerplexityListModelsOutput:
|
87
|
+
"""Process the input and return a list of models."""
|
88
|
+
try:
|
89
|
+
models = []
|
90
|
+
|
91
|
+
# Add models from the configuration
|
92
|
+
for model_id, config in PERPLEXITY_MODELS_CONFIG.items():
|
93
|
+
models.append(
|
94
|
+
{
|
95
|
+
"id": model_id,
|
96
|
+
"display_name": config["name"],
|
97
|
+
"description": config.get("description", ""),
|
98
|
+
"category": config.get("category", "unknown"),
|
99
|
+
"capabilities": {
|
100
|
+
"citations": config.get("citations", False),
|
101
|
+
"search": config.get("search", False),
|
102
|
+
"context_window": config.get("context_window", 8192),
|
103
|
+
"max_completion_tokens": config.get(
|
104
|
+
"max_completion_tokens", 4096
|
105
|
+
),
|
106
|
+
},
|
107
|
+
}
|
108
|
+
)
|
109
|
+
|
110
|
+
return PerplexityListModelsOutput(models=models, provider="perplexity")
|
111
|
+
except Exception as e:
|
112
|
+
raise ProcessingError(f"Failed to list Perplexity AI models: {str(e)}")
|
@@ -0,0 +1,128 @@
|
|
1
|
+
"""Configuration of Perplexity AI model capabilities."""
|
2
|
+
|
3
|
+
from typing import Dict, Any
|
4
|
+
|
5
|
+
|
6
|
+
# Model configuration with capabilities for each Perplexity AI model
|
7
|
+
PERPLEXITY_MODELS_CONFIG = {
|
8
|
+
# Search Models
|
9
|
+
"sonar-pro": {
|
10
|
+
"name": "Sonar Pro",
|
11
|
+
"description": "Advanced search offering with grounding, supporting complex queries and follow-ups.",
|
12
|
+
"category": "search",
|
13
|
+
"context_window": 8192,
|
14
|
+
"max_completion_tokens": 4096,
|
15
|
+
"citations": True,
|
16
|
+
"search": True,
|
17
|
+
},
|
18
|
+
"sonar": {
|
19
|
+
"name": "Sonar",
|
20
|
+
"description": "Lightweight, cost-effective search model with grounding.",
|
21
|
+
"category": "search",
|
22
|
+
"context_window": 8192,
|
23
|
+
"max_completion_tokens": 4096,
|
24
|
+
"citations": True,
|
25
|
+
"search": True,
|
26
|
+
},
|
27
|
+
# Research Models
|
28
|
+
"sonar-deep-research": {
|
29
|
+
"name": "Sonar Deep Research",
|
30
|
+
"description": "Expert-level research model conducting exhaustive searches and generating comprehensive reports.",
|
31
|
+
"category": "research",
|
32
|
+
"context_window": 8192,
|
33
|
+
"max_completion_tokens": 4096,
|
34
|
+
"citations": True,
|
35
|
+
"search": True,
|
36
|
+
},
|
37
|
+
# Reasoning Models
|
38
|
+
"sonar-reasoning-pro": {
|
39
|
+
"name": "Sonar Reasoning Pro",
|
40
|
+
"description": "Premier reasoning offering powered by DeepSeek R1 with Chain of Thought (CoT).",
|
41
|
+
"category": "reasoning",
|
42
|
+
"context_window": 8192,
|
43
|
+
"max_completion_tokens": 4096,
|
44
|
+
"citations": True,
|
45
|
+
"search": True,
|
46
|
+
"chain_of_thought": True,
|
47
|
+
},
|
48
|
+
"sonar-reasoning": {
|
49
|
+
"name": "Sonar Reasoning",
|
50
|
+
"description": "Fast, real-time reasoning model designed for quick problem-solving with search.",
|
51
|
+
"category": "reasoning",
|
52
|
+
"context_window": 8192,
|
53
|
+
"max_completion_tokens": 4096,
|
54
|
+
"citations": True,
|
55
|
+
"search": True,
|
56
|
+
"chain_of_thought": True,
|
57
|
+
},
|
58
|
+
# Offline Models
|
59
|
+
"r1-1776": {
|
60
|
+
"name": "R1-1776",
|
61
|
+
"description": "A version of DeepSeek R1 post-trained for uncensored, unbiased, and factual information.",
|
62
|
+
"category": "offline",
|
63
|
+
"context_window": 8192,
|
64
|
+
"max_completion_tokens": 4096,
|
65
|
+
"citations": False,
|
66
|
+
"search": False,
|
67
|
+
},
|
68
|
+
}
|
69
|
+
|
70
|
+
|
71
|
+
def get_model_config(model_id: str) -> Dict[str, Any]:
|
72
|
+
"""
|
73
|
+
Get the configuration for a specific Perplexity AI model.
|
74
|
+
|
75
|
+
Args:
|
76
|
+
model_id: The model ID to get configuration for
|
77
|
+
|
78
|
+
Returns:
|
79
|
+
Dict with model configuration
|
80
|
+
|
81
|
+
Raises:
|
82
|
+
ValueError: If model_id is not found in configuration
|
83
|
+
"""
|
84
|
+
if model_id in PERPLEXITY_MODELS_CONFIG:
|
85
|
+
return PERPLEXITY_MODELS_CONFIG[model_id]
|
86
|
+
|
87
|
+
# Try to find a match with different format or case
|
88
|
+
normalized_id = model_id.lower().replace("-", "").replace("_", "")
|
89
|
+
for config_id, config in PERPLEXITY_MODELS_CONFIG.items():
|
90
|
+
if normalized_id == config_id.lower().replace("-", "").replace("_", ""):
|
91
|
+
return config
|
92
|
+
|
93
|
+
# If model not found, raise an error
|
94
|
+
raise ValueError(
|
95
|
+
f"Model '{model_id}' not found in Perplexity AI models configuration"
|
96
|
+
)
|
97
|
+
|
98
|
+
|
99
|
+
def get_default_model() -> str:
|
100
|
+
"""Get the default model ID for Perplexity AI."""
|
101
|
+
return "sonar-pro"
|
102
|
+
|
103
|
+
|
104
|
+
def supports_citations(model_id: str) -> bool:
|
105
|
+
"""Check if a model supports citations."""
|
106
|
+
return get_model_config(model_id).get("citations", False)
|
107
|
+
|
108
|
+
|
109
|
+
def supports_search(model_id: str) -> bool:
|
110
|
+
"""Check if a model uses search capabilities."""
|
111
|
+
return get_model_config(model_id).get("search", False)
|
112
|
+
|
113
|
+
|
114
|
+
def get_models_by_category(category: str) -> Dict[str, Dict[str, Any]]:
|
115
|
+
"""
|
116
|
+
Get all models belonging to a specific category.
|
117
|
+
|
118
|
+
Args:
|
119
|
+
category: Category to filter by ('search', 'research', 'reasoning', 'offline')
|
120
|
+
|
121
|
+
Returns:
|
122
|
+
Dict of model IDs and their configurations that match the category
|
123
|
+
"""
|
124
|
+
return {
|
125
|
+
model_id: config
|
126
|
+
for model_id, config in PERPLEXITY_MODELS_CONFIG.items()
|
127
|
+
if config.get("category") == category
|
128
|
+
}
|
@@ -0,0 +1,279 @@
|
|
1
|
+
from typing import Dict, Any, List, Optional, Generator, Union
|
2
|
+
from pydantic import Field, validator
|
3
|
+
import requests
|
4
|
+
|
5
|
+
from airtrain.core.skills import Skill, ProcessingError
|
6
|
+
from airtrain.core.schemas import InputSchema, OutputSchema
|
7
|
+
from .credentials import PerplexityCredentials
|
8
|
+
from .models_config import get_model_config, get_default_model
|
9
|
+
|
10
|
+
|
11
|
+
class PerplexityInput(InputSchema):
|
12
|
+
"""Schema for Perplexity AI chat input"""
|
13
|
+
|
14
|
+
user_input: str = Field(..., description="User's input text")
|
15
|
+
system_prompt: Optional[str] = Field(
|
16
|
+
default=None,
|
17
|
+
description="System prompt to guide the model's behavior",
|
18
|
+
)
|
19
|
+
conversation_history: List[Dict[str, str]] = Field(
|
20
|
+
default_factory=list,
|
21
|
+
description="List of previous conversation messages in [{'role': 'user|assistant', 'content': 'message'}] format",
|
22
|
+
)
|
23
|
+
model: str = Field(
|
24
|
+
default="sonar-pro",
|
25
|
+
description="Perplexity AI model to use",
|
26
|
+
)
|
27
|
+
temperature: Optional[float] = Field(
|
28
|
+
default=0.7, description="Temperature for response generation", ge=0, le=1
|
29
|
+
)
|
30
|
+
max_tokens: Optional[int] = Field(
|
31
|
+
default=500, description="Maximum tokens in response"
|
32
|
+
)
|
33
|
+
top_p: Optional[float] = Field(
|
34
|
+
default=1.0, description="Top-p (nucleus) sampling parameter", ge=0, le=1
|
35
|
+
)
|
36
|
+
top_k: Optional[int] = Field(
|
37
|
+
default=None,
|
38
|
+
description="Top-k sampling parameter",
|
39
|
+
)
|
40
|
+
presence_penalty: Optional[float] = Field(
|
41
|
+
default=None,
|
42
|
+
description="Presence penalty parameter",
|
43
|
+
)
|
44
|
+
frequency_penalty: Optional[float] = Field(
|
45
|
+
default=None,
|
46
|
+
description="Frequency penalty parameter",
|
47
|
+
)
|
48
|
+
|
49
|
+
@validator("model")
|
50
|
+
def validate_model(cls, v):
|
51
|
+
"""Validate that the model is supported by Perplexity AI."""
|
52
|
+
try:
|
53
|
+
get_model_config(v)
|
54
|
+
return v
|
55
|
+
except ValueError as e:
|
56
|
+
raise ValueError(f"Invalid Perplexity AI model: {v}. {str(e)}")
|
57
|
+
|
58
|
+
|
59
|
+
class PerplexityCitation(OutputSchema):
|
60
|
+
"""Schema for Perplexity AI citation information"""
|
61
|
+
|
62
|
+
url: str = Field(..., description="URL of the citation source")
|
63
|
+
title: Optional[str] = Field(None, description="Title of the cited source")
|
64
|
+
snippet: Optional[str] = Field(None, description="Text snippet from the citation")
|
65
|
+
|
66
|
+
|
67
|
+
class PerplexityOutput(OutputSchema):
|
68
|
+
"""Schema for Perplexity AI chat output"""
|
69
|
+
|
70
|
+
response: str = Field(..., description="Model's response text")
|
71
|
+
used_model: str = Field(..., description="Model used for generation")
|
72
|
+
usage: Dict[str, int] = Field(..., description="Usage statistics from the API")
|
73
|
+
citations: Optional[List[PerplexityCitation]] = Field(
|
74
|
+
default=None, description="Citations used in the response, if available"
|
75
|
+
)
|
76
|
+
search_queries: Optional[List[str]] = Field(
|
77
|
+
default=None, description="Search queries used, if available"
|
78
|
+
)
|
79
|
+
|
80
|
+
|
81
|
+
class PerplexityChatSkill(Skill[PerplexityInput, PerplexityOutput]):
|
82
|
+
"""Skill for interacting with Perplexity AI models"""
|
83
|
+
|
84
|
+
input_schema = PerplexityInput
|
85
|
+
output_schema = PerplexityOutput
|
86
|
+
|
87
|
+
def __init__(self, credentials: Optional[PerplexityCredentials] = None):
|
88
|
+
"""Initialize the skill with optional credentials"""
|
89
|
+
super().__init__()
|
90
|
+
self.credentials = credentials or PerplexityCredentials.from_env()
|
91
|
+
self.api_url = "https://api.perplexity.ai/chat/completions"
|
92
|
+
|
93
|
+
def _build_messages(self, input_data: PerplexityInput) -> List[Dict[str, str]]:
|
94
|
+
"""Build messages list from input data including conversation history."""
|
95
|
+
messages = []
|
96
|
+
|
97
|
+
# Add system prompt if provided
|
98
|
+
if input_data.system_prompt:
|
99
|
+
messages.append({"role": "system", "content": input_data.system_prompt})
|
100
|
+
|
101
|
+
# Add conversation history
|
102
|
+
if input_data.conversation_history:
|
103
|
+
messages.extend(input_data.conversation_history)
|
104
|
+
|
105
|
+
# Add current user input
|
106
|
+
messages.append({"role": "user", "content": input_data.user_input})
|
107
|
+
|
108
|
+
return messages
|
109
|
+
|
110
|
+
def _prepare_api_parameters(self, input_data: PerplexityInput) -> Dict[str, Any]:
|
111
|
+
"""Prepare parameters for the API request."""
|
112
|
+
parameters = {
|
113
|
+
"model": input_data.model,
|
114
|
+
"messages": self._build_messages(input_data),
|
115
|
+
"max_tokens": input_data.max_tokens,
|
116
|
+
}
|
117
|
+
|
118
|
+
# Add optional parameters if provided
|
119
|
+
if input_data.temperature is not None:
|
120
|
+
parameters["temperature"] = input_data.temperature
|
121
|
+
|
122
|
+
if input_data.top_p is not None:
|
123
|
+
parameters["top_p"] = input_data.top_p
|
124
|
+
|
125
|
+
if input_data.top_k is not None:
|
126
|
+
parameters["top_k"] = input_data.top_k
|
127
|
+
|
128
|
+
if input_data.presence_penalty is not None:
|
129
|
+
parameters["presence_penalty"] = input_data.presence_penalty
|
130
|
+
|
131
|
+
if input_data.frequency_penalty is not None:
|
132
|
+
parameters["frequency_penalty"] = input_data.frequency_penalty
|
133
|
+
|
134
|
+
return parameters
|
135
|
+
|
136
|
+
def process(self, input_data: PerplexityInput) -> PerplexityOutput:
|
137
|
+
"""Process the input and return the complete response."""
|
138
|
+
try:
|
139
|
+
# Prepare headers with API key
|
140
|
+
headers = {
|
141
|
+
"Authorization": f"Bearer {self.credentials.perplexity_api_key.get_secret_value()}",
|
142
|
+
"Content-Type": "application/json",
|
143
|
+
}
|
144
|
+
|
145
|
+
# Prepare parameters for the API request
|
146
|
+
data = self._prepare_api_parameters(input_data)
|
147
|
+
|
148
|
+
# Make the API request
|
149
|
+
response = requests.post(self.api_url, headers=headers, json=data)
|
150
|
+
|
151
|
+
# Check if request was successful
|
152
|
+
if response.status_code != 200:
|
153
|
+
raise ProcessingError(
|
154
|
+
f"Perplexity AI API error: {response.status_code} - {response.text}"
|
155
|
+
)
|
156
|
+
|
157
|
+
# Parse the response
|
158
|
+
result = response.json()
|
159
|
+
|
160
|
+
# Extract content from the completion
|
161
|
+
content = result["choices"][0]["message"]["content"]
|
162
|
+
|
163
|
+
# Extract and process citations if available
|
164
|
+
citations = None
|
165
|
+
if "citations" in result:
|
166
|
+
citations = [
|
167
|
+
PerplexityCitation(
|
168
|
+
url=citation.get("url", ""),
|
169
|
+
title=citation.get("title"),
|
170
|
+
snippet=citation.get("snippet"),
|
171
|
+
)
|
172
|
+
for citation in result.get("citations", [])
|
173
|
+
]
|
174
|
+
|
175
|
+
# Extract search queries if available
|
176
|
+
search_queries = None
|
177
|
+
if "usage" in result and "num_search_queries" in result["usage"]:
|
178
|
+
search_queries = result.get("search_queries", [])
|
179
|
+
|
180
|
+
# Create and return output
|
181
|
+
return PerplexityOutput(
|
182
|
+
response=content,
|
183
|
+
used_model=input_data.model,
|
184
|
+
usage=result.get("usage", {}),
|
185
|
+
citations=citations,
|
186
|
+
search_queries=search_queries,
|
187
|
+
)
|
188
|
+
|
189
|
+
except Exception as e:
|
190
|
+
if isinstance(e, ProcessingError):
|
191
|
+
raise e
|
192
|
+
raise ProcessingError(f"Perplexity AI processing failed: {str(e)}")
|
193
|
+
|
194
|
+
|
195
|
+
class PerplexityProcessStreamError(Exception):
|
196
|
+
"""Error raised during stream processing"""
|
197
|
+
|
198
|
+
pass
|
199
|
+
|
200
|
+
|
201
|
+
class PerplexityStreamOutput(OutputSchema):
|
202
|
+
"""Schema for streaming output tokens"""
|
203
|
+
|
204
|
+
token: str = Field(..., description="Text token")
|
205
|
+
finish_reason: Optional[str] = Field(
|
206
|
+
None, description="Why the completion finished"
|
207
|
+
)
|
208
|
+
|
209
|
+
|
210
|
+
class PerplexityStreamingChatSkill(PerplexityChatSkill):
|
211
|
+
"""Extension of PerplexityChatSkill that supports streaming responses"""
|
212
|
+
|
213
|
+
def process_stream(
|
214
|
+
self, input_data: PerplexityInput
|
215
|
+
) -> Generator[PerplexityStreamOutput, None, None]:
|
216
|
+
"""
|
217
|
+
Process the input and stream the response tokens.
|
218
|
+
|
219
|
+
Note: Perplexity AI API may not support true streaming. In that case, this
|
220
|
+
method will make a regular API call and yield the entire response at once.
|
221
|
+
"""
|
222
|
+
try:
|
223
|
+
# Prepare headers with API key
|
224
|
+
headers = {
|
225
|
+
"Authorization": f"Bearer {self.credentials.perplexity_api_key.get_secret_value()}",
|
226
|
+
"Content-Type": "application/json",
|
227
|
+
}
|
228
|
+
|
229
|
+
# Prepare parameters for the API request, including stream=true if possible
|
230
|
+
data = self._prepare_api_parameters(input_data)
|
231
|
+
data["stream"] = True
|
232
|
+
|
233
|
+
# Make the API request
|
234
|
+
response = requests.post(
|
235
|
+
self.api_url, headers=headers, json=data, stream=True
|
236
|
+
)
|
237
|
+
|
238
|
+
# Check if request was successful
|
239
|
+
if response.status_code != 200:
|
240
|
+
raise PerplexityProcessStreamError(
|
241
|
+
f"Perplexity AI API error: {response.status_code} - {response.text}"
|
242
|
+
)
|
243
|
+
|
244
|
+
# Process the streaming response if supported
|
245
|
+
for line in response.iter_lines():
|
246
|
+
if line:
|
247
|
+
# Parse the response line
|
248
|
+
try:
|
249
|
+
# Remove 'data: ' prefix if present
|
250
|
+
if line.startswith(b"data: "):
|
251
|
+
line = line[6:]
|
252
|
+
|
253
|
+
# Parse JSON
|
254
|
+
import json
|
255
|
+
|
256
|
+
chunk = json.loads(line)
|
257
|
+
|
258
|
+
# Extract content
|
259
|
+
if "choices" in chunk and len(chunk["choices"]) > 0:
|
260
|
+
choice = chunk["choices"][0]
|
261
|
+
if "delta" in choice and "content" in choice["delta"]:
|
262
|
+
content = choice["delta"]["content"]
|
263
|
+
if content:
|
264
|
+
yield PerplexityStreamOutput(
|
265
|
+
token=content,
|
266
|
+
finish_reason=choice.get("finish_reason"),
|
267
|
+
)
|
268
|
+
except json.JSONDecodeError:
|
269
|
+
# Skip non-JSON lines
|
270
|
+
continue
|
271
|
+
except Exception as e:
|
272
|
+
raise PerplexityProcessStreamError(
|
273
|
+
f"Error processing stream chunk: {str(e)}"
|
274
|
+
)
|
275
|
+
|
276
|
+
except Exception as e:
|
277
|
+
if isinstance(e, PerplexityProcessStreamError):
|
278
|
+
raise ProcessingError(str(e))
|
279
|
+
raise ProcessingError(f"Perplexity AI streaming failed: {str(e)}")
|
@@ -0,0 +1,21 @@
|
|
1
|
+
"""
|
2
|
+
Search integrations for AirTrain.
|
3
|
+
|
4
|
+
This package provides integrations with various search providers.
|
5
|
+
"""
|
6
|
+
|
7
|
+
# Import specific search integrations as needed
|
8
|
+
from .exa import (
|
9
|
+
ExaCredentials,
|
10
|
+
ExaSearchInputSchema,
|
11
|
+
ExaSearchOutputSchema,
|
12
|
+
ExaSearchSkill,
|
13
|
+
)
|
14
|
+
|
15
|
+
__all__ = [
|
16
|
+
# Exa Search
|
17
|
+
"ExaCredentials",
|
18
|
+
"ExaSearchInputSchema",
|
19
|
+
"ExaSearchOutputSchema",
|
20
|
+
"ExaSearchSkill",
|
21
|
+
]
|
@@ -0,0 +1,23 @@
|
|
1
|
+
"""
|
2
|
+
Exa Search API integration.
|
3
|
+
|
4
|
+
This module provides integration with the Exa search API for web searching capabilities.
|
5
|
+
"""
|
6
|
+
|
7
|
+
from .credentials import ExaCredentials
|
8
|
+
from .schemas import (
|
9
|
+
ExaSearchInputSchema,
|
10
|
+
ExaSearchOutputSchema,
|
11
|
+
ExaContentConfig,
|
12
|
+
ExaSearchResult,
|
13
|
+
)
|
14
|
+
from .skills import ExaSearchSkill
|
15
|
+
|
16
|
+
__all__ = [
|
17
|
+
"ExaCredentials",
|
18
|
+
"ExaSearchInputSchema",
|
19
|
+
"ExaSearchOutputSchema",
|
20
|
+
"ExaContentConfig",
|
21
|
+
"ExaSearchResult",
|
22
|
+
"ExaSearchSkill",
|
23
|
+
]
|
@@ -0,0 +1,30 @@
|
|
1
|
+
"""
|
2
|
+
Credentials for Exa Search API.
|
3
|
+
|
4
|
+
This module provides credential management for the Exa search API.
|
5
|
+
"""
|
6
|
+
|
7
|
+
from typing import Optional
|
8
|
+
from pydantic import Field, SecretStr
|
9
|
+
|
10
|
+
from airtrain.core.credentials import BaseCredentials
|
11
|
+
|
12
|
+
|
13
|
+
class ExaCredentials(BaseCredentials):
|
14
|
+
"""Credentials for accessing the Exa search API."""
|
15
|
+
|
16
|
+
exa_api_key: SecretStr = Field(
|
17
|
+
description="Exa search API key",
|
18
|
+
)
|
19
|
+
|
20
|
+
_required_credentials = {"exa_api_key"}
|
21
|
+
|
22
|
+
async def validate_credentials(self) -> bool:
|
23
|
+
"""Validate that the required credentials are present and valid."""
|
24
|
+
# First check that required credentials are present
|
25
|
+
await super().validate_credentials()
|
26
|
+
|
27
|
+
# In a production environment, we might want to make a test API call here
|
28
|
+
# to verify the API key is actually valid, but for now we'll just check
|
29
|
+
# that it's present
|
30
|
+
return True
|