lionagi 0.0.111__py3-none-any.whl → 0.0.113__py3-none-any.whl

lionagi/utils/api_util.py

@@ -1,439 +1,86 @@
-import asyncio
 import logging
 import re
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import Any, Callable, Dict, Generator, NoReturn, Optional
-
-from .sys_util import append_to_jsonl
-
-
-@dataclass
-class StatusTracker:
-    """
-    Class for keeping track of various task statuses.
-    
-    This class serves as a simple way to monitor different types of task
-    outcomes and errors within a system. It uses dataclasses for easy
-    creation and management of state.
-    
-    Attributes:
-        num_tasks_started:
-            The number of tasks that have been initiated.
-        num_tasks_in_progress:
-            The number of tasks currently being processed.
-        num_tasks_succeeded:
-            The number of tasks that have completed successfully.
-        num_tasks_failed:
-            The number of tasks that have failed.
-        num_rate_limit_errors:
-            The number of tasks that failed due to rate limiting.
-        num_api_errors:
-            The number of tasks that failed due to API errors.
-        num_other_errors:
-            The number of tasks that failed due to other errors.
+from typing import Callable
+        
+        
+def api_method(http_session, method: str = "post") -> Callable:
     """
-    num_tasks_started: int = 0
-    num_tasks_in_progress: int = 0
-    num_tasks_succeeded: int = 0
-    num_tasks_failed: int = 0
-    num_rate_limit_errors: int = 0
-    num_api_errors: int = 0
-    num_other_errors: int = 0
-
+    Retrieves the appropriate HTTP method from an HTTP session object.
 
-class AsyncQueue:
-    """
-    A queue class that handles asynchronous operations using asyncio.
+    Parameters:
+        http_session: The HTTP session object from which to retrieve the method.
+        method (str): The HTTP method to retrieve. Defaults to 'post'.
 
-    This class provides an asynchronous queue that can enqueue items, process them
-    asynchronously, and support graceful shutdowns. It is designed to facilitate
-    concurrent task processing in an orderly and controlled manner.
+    Returns:
+        Callable: The HTTP method function from the session object.
 
-    Attributes:
-        queue (asyncio.Queue):
-            A queue to hold items for asynchronous processing.
-        _stop_event (asyncio.Event):
-            An event to signal when the queue should stop processing.
+    Raises:
+        ValueError: If the provided method is not one of ['post', 'delete', 'head', 'options', 'patch'].
 
-    Methods:
-        enqueue(item):
-            Add an item to the queue for processing.
-        dequeue():
-            Remove and return an item from the queue.
-        join():
-            Wait until all items in the queue have been processed.
-        stop():
-            Signal to stop processing new items in the queue.
-        stopped():
-            Check if the queue has been signaled to stop.
-        process_requests(func):
-            Process items using a provided function.
+    Examples:
+        api_method = api_methods(session, "post") # Retrieves the 'post' method from the session
     """
-
-    def __init__(self) -> None:
-        """
-        Initializes an AsyncQueue object with an empty asyncio Queue and a stop event.
-        """
-        self.queue = asyncio.Queue()
-        self._stop_event = asyncio.Event()
-
-    async def enqueue(self, item: Any) -> None:
-        """
-        Asynchronously add an item to the queue for processing.
-
-        Parameters:
-            item (Any): The item to be added to the queue.
-
-        Example:
-            >>> async_queue = AsyncQueue()
-            >>> asyncio.run(async_queue.enqueue('Task 1'))
-        """
-        await self.queue.put(item)
-
-    async def dequeue(self) -> Any:
-        """
-        Asynchronously remove and return an item from the queue.
-
-        If the queue is empty, this method will wait until an item is available.
-
-        Returns:
-            Any: The next item from the queue.
-
-        Example:
-            >>> async_queue = AsyncQueue()
-            >>> asyncio.run(async_queue.enqueue('Task 1'))
-            >>> asyncio.run(async_queue.dequeue())
-            'Task 1'
-        """
-        return await self.queue.get()
-
-    async def join(self) -> None:
-        """
-        Asynchronously wait until all items in the queue have been processed.
-
-        This method blocks until every item that has been enqueued is processed, 
-        ensuring that all tasks are completed.
-
-        Example:
-            >>> async_queue = AsyncQueue()
-            >>> asyncio.run(async_queue.enqueue('Task 1'))
-            >>> asyncio.run(async_queue.join())  # This will block until 'Task 1' is processed.
-        """
-        await self.queue.join()
-
-    async def stop(self) -> None:
-        """
-        Signal the queue to stop processing new items.
-
-        Once called, the queue will not process any new items after the current ones
-        are completed, allowing for a graceful shutdown.
-
-        Example:
-            >>> async_queue = AsyncQueue()
-            >>> asyncio.run(async_queue.stop())  # This signals the queue to stop processing.
-        """
-        self._stop_event.set()
-
-    def stopped(self) -> bool:
-        """
-        Check if the queue has been signaled to stop processing.
-
-        Returns:
-            bool: True if a stop has been signaled, False otherwise.
-
-        Example:
-            >>> async_queue = AsyncQueue()
-            >>> asyncio.run(async_queue.stop())
-            >>> async_queue.stopped()
-            True
-        """
-        return self._stop_event.is_set()
-
-    async def process_requests(self, func: Callable[[Any], Any]) -> None:
-        """
-        Asynchronously process items from the queue using the provided function.
-
-        Continuously dequeues items and applies the given function to each. 
-        The processing stops when the queue is signaled to stop or a sentinel value (`None`) is dequeued.
-
-        Parameters:
-            func (Callable[[Any], Any]): A coroutine function to process items from the queue.
-
-        Example:
-            >>> async def sample_processing(task):
-            ...     print("Processing:", task)
-            >>> async_queue = AsyncQueue()
-            >>> asyncio.run(async_queue.enqueue('Task 1'))
-            >>> asyncio.run(async_queue.process_requests(sample_processing))
-            Processing: Task 1
-        """
-        while not self.stopped():
-            item = await self.dequeue()
-            if item is None:  # Using `None` as a sentinel value to cease processing.
-                await self.stop()
-                break
-            await func(item)
-
-
-class RateLimiter(ABC):
+    if method not in ["post", "delete", "head", "options", "patch"]:
+        raise ValueError("Invalid request, method must be in ['post', 'delete', 'head', 'options', 'patch']")
+    elif method == "post":
+        return http_session.post
+    elif method == "delete":
+        return http_session.delete
+    elif method == "head":
+        return http_session.head
+    elif method == "options":
+        return http_session.options
+    elif method == "patch":
+        return http_session.patch
+
+def api_endpoint_from_url(request_url: str) -> str:
     """
-    An abstract base class for rate limiting mechanisms.
+    Extracts the API endpoint from a given URL.
 
-    This class defines a structure for rate limiters, which are used to control the frequency
-    of requests sent to or received from a network interface controller or an API.
+    Parameters:
+        request_url (str): The URL from which to extract the API endpoint.
 
-    Attributes:
-        max_requests_per_minute (int):
-            Maximum number of requests permitted per minute.
-        max_tokens_per_minute (int):
-            Maximum number of tokens that can accumulate per minute.
-        available_request_capacity (int):
-            Current number of available request slots.
-        available_token_capacity (int):
-            Current number of available tokens.
+    Returns:
+        str: The extracted API endpoint, or an empty string if no match is found.
 
-    Methods:
-        rate_limit_replenisher:
-            Coroutine to replenish rate limits over time.
-        calculate_num_token:
-            Method to calculate required tokens for a request.
+    Examples:
+        endpoint = api_endpoint_from_url("https://api.example.com/v1/users")
+        # endpoint will be 'users'
     """
+    match = re.search(r"^https://[^/]+/v\d+/(.+)$", request_url)
+    return match.group(1) if match else ""
     
-    def __init__(self, max_requests_per_minute: int, max_tokens_per_minute: int) -> None:
-        """
-        Initializes the RateLimiter with specified maximum request and token limits.
-
-        Parameters:
-            max_requests_per_minute (int): Maximum requests allowed per minute.
-
-            max_tokens_per_minute (int): Maximum tokens allowed to accumulate per minute.
-
-        Example:
-            >>> class MyRateLimiter(RateLimiter):
-            ...     async def rate_limit_replenisher(self) -> NoReturn:
-            ...         # Implementation for rate replenishment.
-            ...     def calculate_num_token(self, payload: Dict[str, Any], api_endpoint: str) -> int:
-            ...         # Implementation for token calculation.
-            ...
-            >>> limiter = MyRateLimiter(100, 200)
-        """
-        self.max_requests_per_minute = max_requests_per_minute
-        self.max_tokens_per_minute = max_tokens_per_minute
-        self.available_request_capacity = max_requests_per_minute
-        self.available_token_capacity = max_tokens_per_minute
-    
-    @abstractmethod
-    async def rate_limit_replenisher(self) -> NoReturn:
-        """
-        Asynchronously replenishes rate limit capacities.
-
-        This coroutine should be implemented to periodically restore `available_request_capacity`
-        and `available_token_capacity` according to specific rules defined in subclasses.
-
-        Example:
-            >>> class MyRateLimiter(RateLimiter):
-            ...     async def rate_limit_replenisher(self) -> NoReturn:
-            ...         while True:
-            ...             # Replenishment logic here
-            ...
-            >>> limiter = MyRateLimiter(100, 200)
-        """
-        
-        ...
-    
-    @abstractmethod
-    def calculate_num_token(self, payload: Dict[str, Any], api_endpoint: str) -> int:
-        """
-        Calculates required tokens for a request.
-
-        Subclasses should implement this method to determine the number of tokens needed based
-        on the request payload and target endpoint.
-
-        Parameters:
-            payload (Dict[str, Any]): Payload of the request.
-
-            api_endpoint (str): Target API endpoint for the request.
-
-        Returns:
-            int: Calculated number of tokens required for the request.
-
-        Example:
-            >>> class MyRateLimiter(RateLimiter):
-            ...     def calculate_num_token(self, payload: Dict[str, Any], api_endpoint: str) -> int:
-            ...         return len(payload.get('data', '')) // 10
-            ...
-            >>> limiter = MyRateLimiter(100, 200)
-            >>> limiter.calculate_num_token({'data': '12345'}, 'api/send')
-            0
-        """
-        
-        ...
-
-
-class BaseAPIService(ABC):
+def api_error(response_json: dict) -> bool:
     """
-    Abstract base class for API services requiring asynchronous operations.
+    Logs a warning and returns True if an error is found in the API response.
 
-    This class provides foundational attributes and methods for error handling,
-    logging, rate-limiting, and re-attempting API calls. It is designed to be
-    subclassed for concrete implementations of specific API service interactions.
+    Parameters:
+        response_json (dict): The JSON response from the API call.
 
-    Attributes:
-        api_key (str):
-            The API key used for authenticating with the API service.
-        token_encoding_name (str):
-            The encoding for the API token.
-        max_attempts (int):
-            The maximum number of retry attempts for API calls.
-        status_tracker (StatusTracker):
-            Tracker for API call statuses.
-        queue (AsyncQueue):
-            Queue for managing API call tasks.
-        rate_limiter (RateLimiter):
-            Limiter to control the rate of API calls.
-        append_to_jsonl (callable):
-            Callable for appending data to a file in JSONL format.
+    Returns:
+        bool: True if an error is present in the response, False otherwise.
 
-    Methods:
-        call_api:
-            Abstract method to define API call mechanism in subclasses.
-        handle_error:
-            Handle errors by logging and saving details to a JSONL file.
-        api_endpoint_from_url:
-            Extract the API endpoint from a URL.
-        task_id_generator_function:
-            Generate a sequence of unique task IDs.
+    Examples:
+        if api_error(response):
+            # Handle the error
     """
+    if "error" in response_json:
+        logging.warning(f"API call failed with error: {response_json['error']}")
+        return True
+    return False
     
-    def __init__(
-        self,
-        api_key: str,
-        token_encoding_name: str,
-        max_attempts: int,
-        max_requests_per_minute: int,
-        max_tokens_per_minute: int,
-        ratelimiter,
-        status_tracker: Optional[StatusTracker] = None,
-        queue: Optional[AsyncQueue] = None,
-    ) -> None:
-        """
-        Initializes the BaseAPIService with necessary configuration.
-
-        Parameters:
-            api_key (str): The API key for authentication.
-
-            token_encoding_name (str): Encoding name for the API token.
-
-            max_attempts (int): Maximum number of attempts for an API call.
-
-            status_tracker (Optional[StatusTracker]): Tracker for API call statuses.
-
-            ratelimiter: Limiter for API call rates.
-
-            queue (Optional[AsyncQueue]): Queue for managing API tasks.
-
-        Example:
-            >>> class MyAPIService(BaseAPIService):
-            ...     # Implementation details here
-            ...
-            >>> service = MyAPIService(api_key="12345", token_encoding_name="utf-8",
-            ...                       max_attempts=3, status_tracker=None,
-            ...                       rate_limiter=ratelimiter, queue=None)
-        """
-        self.api_key = api_key
-        self.token_encoding_name = token_encoding_name
-        self.max_attempts = max_attempts
-        self.status_tracker = status_tracker or StatusTracker()
-        self.queue = queue or AsyncQueue()
-        self.rate_limiter = ratelimiter(max_requests_per_minute, max_tokens_per_minute)
-        self.append_to_jsonl = append_to_jsonl
-        
-    @abstractmethod
-    async def call_api(self) -> Any:
-        """
-        Abstract method to be implemented for making specific API calls.
-
-        This method should define the logic for interacting with an API endpoint
-        and must be implemented in subclasses.
-
-        Example:
-            >>> class MyAPIService(BaseAPIService):
-            ...     async def call_api(self):
-            ...         # Implementation details here
-            ...
-        """
-        
-        ...
-
-    def handle_error(
-        self,
-        error: Exception,
-        payload: Any,
-        metadata: Any,
-        save_filepath: str
-    ) -> None:
-        """
-        Handles exceptions that occur during the API call process.
-        
-        Updates the status tracker to indicate the error and saves details to a JSONL file.
-        
-        Parameters:
-            error (Exception): The exception that was raised during the API call.
-
-            payload (Any): The data payload that was used for the API call.
-
-            metadata (Any): Additional metadata related to the API call.
-
-            save_filepath (str): The file path where error details should be saved.
-        """
-        self.status_tracker.num_tasks_in_progress -= 1
-        self.status_tracker.num_tasks_failed += 1
-        data = (
-            [payload, [str(error)], metadata] 
-            if metadata 
-            else [payload, [str(error)]]
-        )
-        self.append_to_jsonl(data, save_filepath)
-        logging.error(f"Request failed after all attempts. Saving errors: {data}")
-
-    @staticmethod
-    def api_endpoint_from_url(request_url: str) -> str:
-        """
-        Extracts the endpoint from an API request URL.
-        
-        Parameters:
-            request_url (str): The URL from which to extract the API endpoint.
-        
-        Returns:
-            str: The extracted API endpoint.
+def api_rate_limit_error(response_json: dict) -> bool:
+    """
+    Checks if the API response indicates a rate limit error.
 
-        Example:
-            >>> BaseAPIService.api_endpoint_from_url("https://api.example.com/v1/test_endpoint")
-            'test_endpoint'
-        """
-        match = re.search(r"^https://[^/]+/v\d+/(.+)$", request_url)
-        if match:
-            return match.group(1)
-        else:
-            return ""
+    Parameters:
+        response_json (dict): The JSON response from the API call.
 
-    @staticmethod
-    def task_id_generator_function() -> Generator[int, None, None]:
-        """
-        Generates a continuous sequence of integers for task IDs.
-        
-        Yields:
-            int: The next task ID in the sequence (0, 1, 2, ...).
+    Returns:
+        bool: True if the response contains a rate limit error message, False otherwise.
 
-        Example:
-            >>> id_gen = BaseAPIService.task_id_generator_function()
-            >>> next(id_gen)
-            0
-            >>> next(id_gen)
-            1
-        """
-        task_id = 0
-        while True:
-            yield task_id
-            task_id += 1
+    Examples:
+        if rate_limit_error(response):
+            # Handle the rate limit error
+    """
+    return "Rate limit" in response_json["error"].get("message", "")