lionagi 0.0.209__py3-none-any.whl → 0.0.211__py3-none-any.whl

lionagi/__init__.py

@@ -19,12 +19,10 @@ from .version import __version__
 from dotenv import load_dotenv
 
 from .utils import *
-from .schema import *
+from .schema.base_schema import *
 from .structures import *
-from .loaders import *
-from .iservices import *
+from .api_service import *
 from .tools import *
-from .core import *
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)

lionagi/api_service/base_endpoint.py

@@ -0,0 +1,65 @@
+from typing import Any, Dict, NoReturn, Optional, Type, List, Union
+from .base_rate_limiter import BaseRateLimiter, SimpleRateLimiter
+
+
+class BaseEndpoint:
+    """
+    Represents an API endpoint with rate limiting capabilities.
+
+    This class encapsulates the details of an API endpoint, including its rate limiter.
+
+    Attributes:
+        endpoint (str): The API endpoint path.
+        rate_limiter_class (Type[li.BaseRateLimiter]): The class used for rate limiting requests to the endpoint.
+        max_requests (int): The maximum number of requests allowed per interval.
+        max_tokens (int): The maximum number of tokens allowed per interval.
+        interval (int): The time interval in seconds for replenishing rate limit capacities.
+        config (Dict): Configuration parameters for the endpoint.
+        rate_limiter (Optional[li.BaseRateLimiter]): The rate limiter instance for this endpoint.
+
+    Examples:
+        # Example usage of EndPoint with SimpleRateLimiter
+        endpoint = EndPoint(
+            max_requests=100,
+            max_tokens=1000,
+            interval=60,
+            endpoint_='chat/completions',
+            rate_limiter_class=li.SimpleRateLimiter,
+            config={'param1': 'value1'}
+        )
+        asyncio.run(endpoint.init_rate_limiter())
+    """
+
+    def __init__(
+        self,
+        max_requests: int = 1_000,
+        max_tokens: int = 100_000,
+        interval: int = 60,
+        endpoint_: Optional[str] = None,
+        rate_limiter_class: Type[BaseRateLimiter] = SimpleRateLimiter,
+        encode_kwargs=None,
+        token_encoding_name=None,
+        config: Dict = None,
+    ) -> None:
+        self.endpoint = endpoint_ or 'chat/completions'
+        self.rate_limiter_class = rate_limiter_class
+        self.max_requests = max_requests
+        self.max_tokens = max_tokens
+        self.interval = interval
+        self.token_encoding_name = token_encoding_name
+        self.config = config or {}
+        self.rate_limiter: Optional[BaseRateLimiter] = None
+        self._has_initialized = False
+        self.encode_kwargs = encode_kwargs or {}
+
+    async def init_rate_limiter(self) -> None:
+        """Initializes the rate limiter for the endpoint."""
+        self.rate_limiter = await self.rate_limiter_class.create(
+            self.max_requests, self.max_tokens, self.interval, self.token_encoding_name
+        )
+        self._has_initialized = True
+
+
+class Embedding(BaseEndpoint):
+    ...
+    

lionagi/api_service/base_rate_limiter.py

@@ -0,0 +1,121 @@
+import asyncio
+import logging
+from abc import ABC
+from typing import Dict, NoReturn, Optional
+
+from ..utils import APIUtil
+
+
+class BaseRateLimiter(ABC):
+    def __init__(self, max_requests: int, max_tokens: int, interval: int = 60, token_encoding_name=None) -> None:
+        self.interval: int = interval
+        self.max_requests: int = max_requests
+        self.max_tokens: int = max_tokens
+        self.available_request_capacity: int = max_requests
+        self.available_token_capacity: int = max_tokens
+        self.rate_limit_replenisher_task: Optional[asyncio.Task[NoReturn]] = None
+        self._stop_replenishing: asyncio.Event = asyncio.Event()
+        self._lock: asyncio.Lock = asyncio.Lock()
+        self.token_encoding_name = token_encoding_name
+
+    async def start_replenishing(self) -> NoReturn:
+        """Starts the replenishment of rate limit capacities at regular intervals."""
+        try:
+            while not self._stop_replenishing.is_set():
+                await asyncio.sleep(self.interval)
+                async with self._lock:
+                    self.available_request_capacity = self.max_requests
+                    self.available_token_capacity = self.max_tokens
+        except asyncio.CancelledError:
+            logging.info("Rate limit replenisher task cancelled.")
+
+        except Exception as e:
+            logging.error(f"An error occurred in the rate limit replenisher: {e}")
+
+    async def stop_replenishing(self) -> None:
+        """Stops the replenishment task."""
+        if self.rate_limit_replenisher_task:
+            self.rate_limit_replenisher_task.cancel()
+            await self.rate_limit_replenisher_task
+        self._stop_replenishing.set()
+
+    async def request_permission(self, required_tokens) -> bool:
+        """Requests permission to make an API call.
+
+        Returns True if the request can be made immediately, otherwise False.
+        """
+        async with self._lock:
+            if self.available_request_capacity > 0 and self.available_token_capacity > 0:
+                self.available_request_capacity -= 1
+                self.available_token_capacity -= required_tokens  # Assuming 1 token per request for simplicity
+                return True
+            return False
+
+    async def _call_api(
+        self, 
+        http_session, 
+        endpoint: str, 
+        base_url: str, 
+        api_key: str,
+        max_attempts: int = 3, 
+        method: str = "post", 
+        payload: Dict[str, any]=None, 
+        **kwargs,
+    ) -> Optional[Dict[str, any]]:
+        endpoint = APIUtil.api_endpoint_from_url(base_url + endpoint)
+        while True:
+            if self.available_request_capacity < 1 or self.available_token_capacity < 10:  # Minimum token count
+                await asyncio.sleep(1)  # Wait for capacity
+                continue
+            required_tokens = APIUtil.calculate_num_token(payload, endpoint, self.token_encoding_name, **kwargs)
+            
+            if await self.request_permission(required_tokens):
+                request_headers = {"Authorization": f"Bearer {api_key}"}
+                attempts_left = max_attempts
+
+                while attempts_left > 0:
+                    try:
+                        method = APIUtil.api_method(http_session, method)                         
+                        async with method(
+                            url=(base_url+endpoint), headers=request_headers, json=payload
+                        ) as response:
+                            response_json = await response.json()
+
+                            if "error" in response_json:
+                                logging.warning(
+                                    f"API call failed with error: {response_json['error']}"
+                                )
+                                attempts_left -= 1
+
+                                if "Rate limit" in response_json["error"].get("message", ""):
+                                    await asyncio.sleep(15)
+                            else:
+                                return response_json
+                    except Exception as e:
+                        logging.warning(f"API call failed with exception: {e}")
+                        attempts_left -= 1
+
+                logging.error("API call failed after all attempts.")
+                break
+            else:
+                await asyncio.sleep(1)
+
+    @classmethod
+    async def create(cls, max_requests: int, max_tokens: int, interval: int = 60, token_encoding_name = None) -> 'BaseRateLimiter':
+        instance = cls(max_requests, max_tokens, interval, token_encoding_name)
+        instance.rate_limit_replenisher_task = asyncio.create_task(
+            instance.start_replenishing()
+        )
+        return instance
+
+
+class SimpleRateLimiter(BaseRateLimiter):
+    """
+    A simple implementation of a rate limiter.
+
+    Inherits from BaseRateLimiter and provides a basic rate limiting mechanism.
+    """
+
+    def __init__(self, max_requests: int, max_tokens: int, interval: int = 60, token_encoding_name=None) -> None:
+        """Initializes the SimpleRateLimiter with the specified parameters."""
+        super().__init__(max_requests, max_tokens, interval, token_encoding_name)

lionagi/api_service/base_service.py

@@ -0,0 +1,146 @@
+import asyncio
+import logging
+import aiohttp
+from abc import ABC
+from dataclasses import dataclass
+from typing import Any, Dict, NoReturn, Optional, Type, List, Union
+
+from ..utils import nget, APIUtil, to_list, lcall
+from .base_rate_limiter import BaseRateLimiter, SimpleRateLimiter
+from .status_tracker import StatusTracker
+
+from .base_endpoint import BaseEndpoint
+
+
+class BaseService:
+    """
+    Base class for services that interact with API endpoints.
+
+    This class provides a foundation for services that need to make API calls with rate limiting.
+
+    Attributes:
+        api_key (Optional[str]): The API key used for authentication.
+        schema (Dict[str, Any]): The schema defining the service's endpoints.
+        status_tracker (StatusTracker): The object tracking the status of API calls.
+        endpoints (Dict[str, BaseEndpoint]): A dictionary of endpoint objects.
+    """
+
+    base_url: str = ''
+    available_endpoints: list = []
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        schema: Dict[str, Any] = None,
+        token_encoding_name: str = None,
+        max_tokens : int = 100_000,
+        max_requests : int = 1_000,
+        interval: int = 60
+    ) -> None:
+        self.api_key = api_key
+        self.schema = schema or {}
+        self.status_tracker = StatusTracker()
+        self.endpoints: Dict[str, BaseEndpoint] = {}
+        self.token_encoding_name = token_encoding_name
+        self.chat_config = {
+            'max_requests': max_requests,
+            'max_tokens': max_tokens,
+            'interval': interval,
+            "token_encoding_name": token_encoding_name
+        }
+
+
+    async def init_endpoint(self, endpoint_: Optional[Union[List[str], List[BaseEndpoint], str, BaseEndpoint]] = None) -> None:
+        """
+        Initializes the specified endpoint or all endpoints if none is specified.
+
+        Args:
+            endpoint_: The endpoint(s) to initialize. Can be a string, an BaseEndpoint, a list of strings, or a list of BaseEndpoints.
+        """
+
+        if endpoint_:
+            endpoint_ = to_list(endpoint_, flatten=True, dropna=True)
+            
+            for ep in endpoint_:
+                self._check_endpoints(ep)
+            
+                if ep not in self.endpoints:
+                    endpoint_config = self._get_endpoint(ep)
+
+                    if endpoint_config is not None:
+                        if ep == "chat/completions":
+                            self.endpoints[ep] = BaseEndpoint(
+                                max_requests=self.chat_config.get('max_requests', 1000),
+                                max_tokens=self.chat_config.get('max_tokens', 100000),
+                                interval=self.chat_config.get('interval', 60),
+                                endpoint_=ep,
+                                token_encoding_name=self.token_encoding_name,
+                                config=endpoint_config,
+                            )
+                        else:
+                            self.endpoints[ep] = BaseEndpoint(
+                                max_requests=endpoint_config.get('max_requests', 1000) if endpoint_config.get('max_requests', 1000) is not None else 1000,
+                                max_tokens=endpoint_config.get('max_tokens', 100000) if endpoint_config.get('max_tokens', 100000) is not None else 100000,
+                                interval=endpoint_config.get('interval', 60) if endpoint_config.get('interval', 60) is not None else 60,
+                                endpoint_=ep,
+                                token_encoding_name=self.token_encoding_name,
+                                config=endpoint_config,
+                            )
+
+                if not self.endpoints[ep]._has_initialized:
+                    await self.endpoints[ep].init_rate_limiter()
+
+        else:
+            for ep in self.available_endpoints:
+                endpoint_config = nget(self.schema, [ep, 'config'])
+                self.schema.get(ep, {})
+                if ep not in self.endpoints:
+                    self.endpoints[ep] = BaseEndpoint(
+                        max_requests=endpoint_config.get('max_requests', 1000),
+                        max_tokens=endpoint_config.get('max_tokens', 100000),
+                        interval=endpoint_config.get('interval', 60),
+                        endpoint_=ep,
+                        token_encoding_name=self.token_encoding_name,
+                        config=endpoint_config,
+                    )
+                if not self.endpoints[ep]._has_initialized:
+                    await self.endpoints[ep].init_rate_limiter()
+
+    async def call_api(self, payload, endpoint, method, **kwargs):
+        """
+        Calls the specified API endpoint with the given payload and method.
+
+        Args:
+            payload: The payload to send with the API call.
+            endpoint: The endpoint to call.
+            method: The HTTP method to use for the call.
+            kwargs are for tiktoken encoding
+
+        Returns:
+            The response from the API call.
+
+        Raises:
+            ValueError: If the endpoint has not been initialized.
+        """
+        if endpoint not in self.endpoints.keys():
+            raise ValueError(f'The endpoint {endpoint} has not initialized.')
+        async with aiohttp.ClientSession() as http_session:
+            completion = await self.endpoints[endpoint].rate_limiter._call_api(
+                http_session=http_session, endpoint=endpoint, base_url=self.base_url, api_key=self.api_key,
+                method=method, payload=payload, **kwargs)
+            return completion
+
+    def _check_endpoints(self, endpoint_):
+        f = lambda ep: ValueError (f"Endpoint {ep} not available for service {self.__class__.__name__}")
+        if not endpoint_ in self.available_endpoints:
+            raise f(endpoint_)
+        
+    def _get_endpoint(self, endpoint_):
+        if endpoint_ not in self.endpoints:
+            endpoint_config = nget(self.schema, [endpoint_, 'config'])
+            self.schema.get(endpoint_, {})
+
+            if isinstance(endpoint_, BaseEndpoint):
+                self.endpoints[endpoint_.endpoint] = endpoint_
+            return None
+        return endpoint_config

lionagi/api_service/chat_completion.py

@@ -0,0 +1,6 @@
+from .base_endpoint import BaseEndpoint
+from .payload_package import PayloadCreation
+
+class ChatCompletion(BaseEndpoint):
+    ...
+    

lionagi/api_service/embeddings.py

@@ -0,0 +1,6 @@
+from .base_endpoint import BaseEndpoint
+from .payload_package import PayloadCreation
+
+class Embeddings(BaseEndpoint):
+    ...
+    

lionagi/api_service/payload_package.py

@@ -0,0 +1,47 @@
+from lionagi.utils.api_util import APIUtil
+
+class PayloadCreation:
+    
+    @classmethod
+    def chat_completion(cls, messages, llmconfig, schema, **kwargs):
+        """
+        Creates a payload for the chat completion operation.
+
+        Args:
+            messages: The messages to include in the chat completion.
+            llmconfig: Configuration for the language model.
+            schema: The schema describing required and optional fields.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The constructed payload.
+        """
+        return APIUtil._create_payload(
+            input_=messages, 
+            config=llmconfig, 
+            required_=schema['required'], 
+            optional_=schema['optional'],
+            input_key="messages", 
+            **kwargs)
+
+    @classmethod
+    def fine_tuning(cls, training_file, llmconfig, schema, **kwargs):
+        """
+        Creates a payload for the fine-tuning operation.
+
+        Args:
+            training_file: The file containing training data.
+            llmconfig: Configuration for the language model.
+            schema: The schema describing required and optional fields.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The constructed payload.
+        """
+        return APIUtil._create_payload(
+            input_=training_file, 
+            config=llmconfig, 
+            required_=schema['required'], 
+            optional_=schema['optional'],
+            input_key="training_file", 
+            **kwargs)

lionagi/api_service/status_tracker.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class StatusTracker:
+    """
+    Keeps track of various task statuses within a system.
+
+    Attributes:
+        num_tasks_started (int): The number of tasks that have been initiated.
+        num_tasks_in_progress (int): The number of tasks currently being processed.
+        num_tasks_succeeded (int): The number of tasks that have completed successfully.
+        num_tasks_failed (int): The number of tasks that have failed.
+        num_rate_limit_errors (int): The number of tasks that failed due to rate limiting.
+        num_api_errors (int): The number of tasks that failed due to API errors.
+        num_other_errors (int): The number of tasks that failed due to other errors.
+
+    Examples:
+        >>> tracker = StatusTracker()
+        >>> tracker.num_tasks_started += 1
+        >>> tracker.num_tasks_succeeded += 1
+    """
+    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

lionagi/core/__init__.py

@@ -1,7 +1,7 @@
 from .branch import Branch
 from .session import Session
 
-__all__  = [
-    'Branch',
-    'Session'
+__all__ = [
+    "Branch",
+    "Session"
 ]

lionagi/core/branch.py

@@ -9,9 +9,9 @@ from lionagi.utils.sys_util import create_path, is_same_dtype
 from lionagi.utils import as_dict, lcall,to_df, to_list, CoreUtil
 
 
-from lionagi.iservices.base_service import BaseService, StatusTracker
-from lionagi.iservices.oai import OpenAIService
-from lionagi.iservices.openrouter import OpenRouterService
+from lionagi.services.base_service import BaseService, StatusTracker
+from lionagi.services.oai import OpenAIService
+from lionagi.services.openrouter import OpenRouterService
 from lionagi.configs.oai_configs import oai_schema
 from lionagi.configs.openrouter_configs import openrouter_schema
 from lionagi.schema import DataLogger, Tool
@@ -28,6 +28,25 @@ except:
     pass
 
 class Branch:
+    """
+    Represents a branch in a conversation with messages, instruction sets, and tool management.
+
+    A `Branch` is a subset of a conversation that contains messages, instruction sets, and tools for managing interactions
+    within the conversation. It encapsulates the state and behavior of a specific branch of conversation flow.
+
+    Attributes:
+        _cols (List[str]): A list of column names for the DataFrame containing messages.
+        messages (pd.DataFrame): A DataFrame containing messages for the branch.
+        instruction_sets (Dict[str, InstructionSet]): A dictionary of instruction sets associated with the branch.
+        tool_manager (ToolManager): The tool manager for managing tools within the branch.
+        service (Optional[BaseService]): The service associated with the branch.
+        llmconfig (Optional[Dict]): Configuration for the LLM (Large Language Model) service.
+        name (Optional[str]): The name of the branch.
+        pending_ins (Dict): Dictionary to store pending inputs for the branch.
+        pending_outs (Deque): Queue to store pending outputs for the branch.
+        logger (Optional[DataLogger]): Logger for data logging.
+        status_tracker (StatusTracker): Tracks the status of the branch.
+    """
     _cols = ["node_id", "role", "sender", "timestamp", "content"]
     
     def __init__(self, name: Optional[str] = None, messages: Optional[pd.DataFrame] = None,

lionagi/core/session.py

@@ -5,14 +5,26 @@ import pandas as pd
 from lionagi.utils.sys_util import create_path
 from lionagi.utils import to_list, to_df
 from lionagi.schema import Tool
-from lionagi.iservices.base_service import BaseService
+from lionagi.services.base_service import BaseService
 from lionagi.core.branch import Branch
 from lionagi.core.branch_manager import BranchManager
 from lionagi.core.messages import Instruction, System
 
 
 class Session:
-
+    """
+    Represents a session for managing conversations and branches.
+
+    A `Session` encapsulates the state and behavior for managing conversations and their branches.
+    It provides functionality for initializing and managing conversation sessions, including setting up default
+    branches, configuring language learning models, managing tools, and handling session data logging.
+
+    Attributes:
+        branches (Dict[str, Branch]): A dictionary of branch instances associated with the session.
+        service (Optional[BaseService]): The external service instance associated with the session.
+        branch_manager (BranchManager): The manager for handling branches within the session.
+        logger (Optional[Any]): The logger instance for session data logging.
+    """
     def __init__(
         self,
         system: Optional[Union[str, System]] = None,

lionagi/schema/__init__.py

@@ -1,11 +1,8 @@
-from .base_node import BaseNode
-from .base_tool import Tool
-from .data_logger import DataLogger
-from .data_node import DataNode
+from .base_schema import BaseNode, Tool, DataLogger, DataNode
 
 __all__ = [
-    "BaseNode", 
-    "DataNode", 
-    "Tool", 
-    "DataLogger", 
+    "BaseNode",
+    "Tool",
+    "DataLogger",
+    "DataNode"
 ]