airtrain 0.1.58__py3-none-any.whl → 0.1.62__py3-none-any.whl

airtrain/integrations/search/exa/schemas.py

@@ -0,0 +1,114 @@
+"""
+Schemas for Exa Search API.
+
+This module defines the input and output schemas for the Exa search API.
+"""
+
+from typing import Dict, List, Optional, Any, Union
+from pydantic import BaseModel, Field, HttpUrl
+
+from airtrain.core.schemas import InputSchema, OutputSchema
+
+
+class ExaContentConfig(BaseModel):
+    """Configuration for the content to be returned by Exa search."""
+
+    text: bool = Field(default=True, description="Whether to return text content.")
+    extractedText: Optional[bool] = Field(
+        default=None, description="Whether to return extracted text content."
+    )
+    embedded: Optional[bool] = Field(
+        default=None, description="Whether to return embedded content."
+    )
+    links: Optional[bool] = Field(
+        default=None, description="Whether to return links from the content."
+    )
+    screenshot: Optional[bool] = Field(
+        default=None, description="Whether to return screenshots of the content."
+    )
+    highlighted: Optional[bool] = Field(
+        default=None, description="Whether to return highlighted text."
+    )
+
+
+class ExaSearchInputSchema(InputSchema):
+    """Input schema for Exa search API."""
+
+    query: str = Field(description="The search query to execute.")
+    numResults: Optional[int] = Field(
+        default=None, description="Number of results to return."
+    )
+    contents: Optional[ExaContentConfig] = Field(
+        default_factory=ExaContentConfig,
+        description="Configuration for the content to be returned.",
+    )
+    highlights: Optional[dict] = Field(
+        default=None, description="Highlighting configuration for search results."
+    )
+    useAutoprompt: Optional[bool] = Field(
+        default=None, description="Whether to use autoprompt for the search."
+    )
+    type: Optional[str] = Field(default=None, description="Type of search to perform.")
+    includeDomains: Optional[List[str]] = Field(
+        default=None, description="List of domains to include in the search."
+    )
+    excludeDomains: Optional[List[str]] = Field(
+        default=None, description="List of domains to exclude from the search."
+    )
+
+
+class ExaModerationConfig(BaseModel):
+    """Moderation configuration returned in search results."""
+
+    llamaguardS1: Optional[bool] = None
+    llamaguardS3: Optional[bool] = None
+    llamaguardS4: Optional[bool] = None
+    llamaguardS12: Optional[bool] = None
+    domainBlacklisted: Optional[bool] = None
+    domainBlacklistedMedia: Optional[bool] = None
+
+
+class ExaHighlight(BaseModel):
+    """Highlight information for a search result."""
+
+    text: str
+    score: float
+
+
+class ExaSearchResult(BaseModel):
+    """Individual search result from Exa."""
+
+    id: str
+    url: str
+    title: Optional[str] = None
+    text: Optional[str] = None
+    extractedText: Optional[str] = None
+    embedded: Optional[Dict[str, Any]] = None
+    score: float
+    published: Optional[str] = None
+    author: Optional[str] = None
+    highlights: Optional[List[ExaHighlight]] = None
+    robotsAllowed: Optional[bool] = None
+    moderationConfig: Optional[ExaModerationConfig] = None
+    urls: Optional[List[str]] = None
+
+
+class ExaCostDetails(BaseModel):
+    """Cost details for an Exa search request."""
+
+    total: float
+    search: Dict[str, float]
+    contents: Dict[str, float]
+
+
+class ExaSearchOutputSchema(OutputSchema):
+    """Output schema for Exa search API."""
+
+    results: List[ExaSearchResult] = Field(description="List of search results.")
+    query: str = Field(description="The original search query.")
+    autopromptString: Optional[str] = Field(
+        default=None, description="Autoprompt string used for the search if enabled."
+    )
+    costDollars: Optional[ExaCostDetails] = Field(
+        default=None, description="Cost details for the search request."
+    )

airtrain/integrations/search/exa/skills.py

@@ -0,0 +1,114 @@
+"""
+Skills for Exa Search API.
+
+This module provides skills for using the Exa search API.
+"""
+
+import json
+import logging
+import httpx
+from typing import Optional, Dict, Any, List, cast
+
+from pydantic import ValidationError
+
+from airtrain.core.skills import Skill, ProcessingError
+from .credentials import ExaCredentials
+from .schemas import ExaSearchInputSchema, ExaSearchOutputSchema, ExaSearchResult
+
+
+logger = logging.getLogger(__name__)
+
+
+class ExaSearchSkill(Skill[ExaSearchInputSchema, ExaSearchOutputSchema]):
+    """Skill for searching the web using the Exa search API."""
+
+    input_schema = ExaSearchInputSchema
+    output_schema = ExaSearchOutputSchema
+
+    EXA_API_ENDPOINT = "https://api.exa.ai/search"
+
+    def __init__(
+        self,
+        credentials: ExaCredentials,
+        timeout: float = 60.0,
+        max_retries: int = 3,
+        **kwargs,
+    ):
+        """
+        Initialize the Exa search skill.
+
+        Args:
+            credentials: Credentials for accessing the Exa API
+            timeout: Timeout for API requests in seconds
+            max_retries: Maximum number of retries for failed requests
+        """
+        super().__init__(**kwargs)
+        self.credentials = credentials
+        self.timeout = timeout
+        self.max_retries = max_retries
+
+    async def process(self, input_data: ExaSearchInputSchema) -> ExaSearchOutputSchema:
+        """
+        Process a search request using the Exa API.
+
+        Args:
+            input_data: Search input parameters
+
+        Returns:
+            Search results from Exa
+
+        Raises:
+            ProcessingError: If there's an issue with the API request
+        """
+        try:
+            # Prepare request payload
+            payload = input_data.model_dump(exclude_none=True)
+
+            # Build request headers
+            headers = {
+                "content-type": "application/json",
+                "Authorization": f"Bearer {self.credentials.api_key.get_secret_value()}",
+            }
+
+            # Make the API request
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    self.EXA_API_ENDPOINT,
+                    headers=headers,
+                    json=payload,
+                    timeout=self.timeout,
+                )
+
+                # Check for successful response
+                if response.status_code == 200:
+                    result_data = response.json()
+
+                    # Construct the output schema
+                    output = ExaSearchOutputSchema(
+                        results=result_data.get("results", []),
+                        query=input_data.query,
+                        autopromptString=result_data.get("autopromptString"),
+                        costDollars=result_data.get("costDollars"),
+                    )
+
+                    return output
+                else:
+                    # Handle error responses
+                    error_message = f"Exa API returned status code {response.status_code}: {response.text}"
+                    logger.error(error_message)
+                    raise ProcessingError(error_message)
+
+        except httpx.TimeoutException:
+            error_message = f"Timeout while querying Exa API (timeout={self.timeout}s)"
+            logger.error(error_message)
+            raise ProcessingError(error_message)
+
+        except ValidationError as e:
+            error_message = f"Schema validation error: {str(e)}"
+            logger.error(error_message)
+            raise ProcessingError(error_message)
+
+        except Exception as e:
+            error_message = f"Unexpected error while querying Exa API: {str(e)}"
+            logger.error(error_message)
+            raise ProcessingError(error_message)

airtrain/tools/__init__.py

@@ -12,30 +12,34 @@ from .registry import (
     ToolFactory,
     ToolValidationError,
     register_tool,
-    execute_tool_call
+    execute_tool_call,
 )
 
 # Import standard tools
 from .filesystem import ListDirectoryTool, DirectoryTreeTool
 from .network import ApiCallTool
-from .command import ExecuteCommandTool, FindFilesTool
+from .command import ExecuteCommandTool, FindFilesTool, TerminalNavigationTool
+from .search import SearchTermTool, WebSearchTool
+from .testing import RunPytestTool
 
 __all__ = [
     # Base classes
     "BaseTool",
     "StatelessTool",
     "StatefulTool",
-    
     # Registry components
     "ToolFactory",
     "ToolValidationError",
     "register_tool",
     "execute_tool_call",
-    
     # Standard tools
     "ListDirectoryTool",
     "DirectoryTreeTool",
     "ApiCallTool",
     "ExecuteCommandTool",
     "FindFilesTool",
-] 
+    "TerminalNavigationTool",
+    "SearchTermTool",
+    "WebSearchTool",
+    "RunPytestTool",
+]