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.
@@ -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