lionagi 0.5.5__py3-none-any.whl → 0.6.1__py3-none-any.whl

lionagi/service/imodel.py

@@ -1,130 +1,297 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
 import warnings
 
-from lionagi.service.service import Service
-from lionagi.service.service_match_util import (
-    match_parameters,
-    match_service,
-    match_task_method,
+from .endpoints.base import APICalling, EndPoint
+from .endpoints.match_endpoint import match_endpoint
+from .endpoints.rate_limited_processor import RateLimitedAPIExecutor
+
+warnings.filterwarnings(
+    "ignore",
+    message=".*Valid config keys have changed in V2.*",
+    category=UserWarning,
+    module="pydantic._internal._config",
 )
 
 
 class iModel:
+    """Manages API calls for a specific provider with optional rate-limiting.
+
+    The iModel class encapsulates a specific endpoint configuration (e.g.,
+    chat or completion endpoints). It determines and sets the necessary
+    API key based on the provider and uses a RateLimitedAPIExecutor to
+    handle queuing and throttling requests.
+
+    Attributes:
+        endpoint (EndPoint):
+            The chosen endpoint object (constructed via `match_endpoint` if
+            none is provided).
+        should_invoke_endpoint (bool):
+            If True, the endpoint is called for real. If False, calls might
+            be mocked or cached.
+        kwargs (dict):
+            Any additional keyword arguments passed to initialize and
+            configure the iModel (e.g., `model`, `api_key`).
+        executor (RateLimitedAPIExecutor):
+            The rate-limited executor that queues and runs API calls in a
+            controlled fashion.
+    """
 
     def __init__(
         self,
-        provider: str | Service,
-        task: str = "chat",
-        model: str = None,
+        provider: str = None,
+        base_url: str = None,
+        endpoint: str | EndPoint = "chat",
+        endpoint_params: list[str] | None = None,
         api_key: str = None,
-        api_key_schema: str = None,
-        interval_tokens: int = None,
-        interval_requests: int = None,
+        queue_capacity: int = 100,
+        capacity_refresh_time: float = 60,
+        interval: float | None = None,
+        limit_requests: int = None,
+        limit_tokens: int = None,
+        invoke_with_endpoint: bool = True,
         **kwargs,
-    ):
-        if api_key is not None:
-            api_key = api_key
-        elif api_key_schema is not None:
-            api_key = api_key_schema
+    ) -> None:
+        """Initializes the iModel instance.
 
-        if task == "chat":
+        Args:
+            provider (str, optional):
+                Name of the provider (e.g., 'openai', 'anthropic').
+            base_url (str, optional):
+                Base URL for the API (if a custom endpoint is needed).
+            endpoint (str | EndPoint, optional):
+                Either a string representing the endpoint type (e.g., 'chat')
+                or an `EndPoint` instance.
+            endpoint_params (list[str] | None, optional):
+                Additional parameters for the endpoint (e.g., 'v1' or other).
+            api_key (str, optional):
+                An explicit API key. If not given, tries to load one from
+                environment variables based on the provider.
+            queue_capacity (int, optional):
+                Maximum number of requests allowed in the queue before
+                executing them.
+            capacity_refresh_time (float, optional):
+                Time interval (in seconds) after which the queue capacity
+                is refreshed.
+            interval (float | None, optional):
+                Interval in seconds to check or process requests in
+                the queue. If None, defaults to capacity_refresh_time.
+            limit_requests (int | None, optional):
+                Maximum number of requests allowed per cycle, if any.
+            limit_tokens (int | None, optional):
+                Maximum number of tokens allowed per cycle, if any.
+            invoke_with_endpoint (bool, optional):
+                If True, the endpoint is actually invoked. If False,
+                calls might be mocked or cached.
+            **kwargs:
+                Additional keyword arguments, such as `model`, or any other
+                provider-specific fields.
+        """
+        if api_key is None:
             match provider:
                 case "openai":
-                    task = "create_chat_completion"
+                    api_key = "OPENAI_API_KEY"
                 case "anthropic":
-                    task = "create_message"
-                case "groq":
-                    task = "create_chat_completion"
+                    api_key = "ANTHROPIC_API_KEY"
+                case "openrouter":
+                    api_key = "OPENROUTER_API_KEY"
                 case "perplexity":
-                    task = "create_chat_completion"
-
-        if isinstance(provider, str):
-            if api_key is None:
-                match provider:
-                    case "openai":
-                        api_key = "OPENAI_API_KEY"
-                    case "anthropic":
-                        api_key = "ANTHROPIC_API_KEY"
-                    case "groq":
-                        api_key = "GROQ_API_KEY"
-                    case "perplexity":
-                        api_key = "PERPLEXIY_API_KEY"
-
-            self.service = match_service(provider, api_key=api_key, **kwargs)
-        elif isinstance(provider, Service):
-            self.service = provider
-            if api_key:
-                warnings.warn(
-                    "A Service instance was provided along with api key info."
-                    "The the separately provided api_key or api_key_schema will be ignored."
-                )
+                    api_key = "PERPLEXITY_API_KEY"
+                case "groq":
+                    api_key = "GROQ_API_KEY"
+
+        if os.getenv(api_key, None) is not None:
+            self.api_key_scheme = api_key
+            api_key = os.getenv(api_key)
+
+        kwargs["api_key"] = api_key
+        model = kwargs.get("model", None)
+        if model:
+            if not provider:
+                if "/" in model:
+                    provider = model.split("/")[0]
+                    model = model.replace(provider + "/", "")
+                    kwargs["model"] = model
+                else:
+                    raise ValueError("Provider must be provided")
+
+        if isinstance(endpoint, EndPoint):
+            self.endpoint = endpoint
         else:
-            raise ValueError(
-                "Invalid provider. Please provide a valid provider name or valid service object."
+            self.endpoint = match_endpoint(
+                provider=provider,
+                base_url=base_url,
+                endpoint=endpoint,
+                endpoint_params=endpoint_params,
             )
+        if provider:
+            self.endpoint.config.provider = provider
+        if base_url:
+            self.endpoint.config.base_url = base_url
 
-        task_method_list = match_task_method(task, self.service)
-        if len(task_method_list) == 0:
-            raise ValueError(
-                "No matching task found. "
-                "Please refer to the service provider's API to provide a valid task."
-            )
-        if len(task_method_list) > 1:
-            raise ValueError(
-                f"Multiple possible tasks found. Please specify: {task_method_list}"
-            )
-        self.task = task_method_list[0]
-        task_method = getattr(self.service, task_method_list[0])
-        task_params = match_parameters(
-            task_method, model, interval_tokens, interval_requests
+        self.should_invoke_endpoint = invoke_with_endpoint
+        self.kwargs = kwargs
+        self.executor = RateLimitedAPIExecutor(
+            queue_capacity=queue_capacity,
+            capacity_refresh_time=capacity_refresh_time,
+            interval=interval,
+            limit_requests=limit_requests,
+            limit_tokens=limit_tokens,
+        )
+
+    def create_api_calling(self, **kwargs) -> APICalling:
+        """Constructs an `APICalling` object from endpoint-specific payload.
+
+        Args:
+            **kwargs:
+                Additional arguments used to generate the payload (merged
+                with self.kwargs).
+
+        Returns:
+            APICalling:
+                An `APICalling` instance with the constructed payload,
+                headers, and the selected endpoint.
+        """
+        kwargs.update(self.kwargs)
+        payload = self.endpoint.create_payload(**kwargs)
+        return APICalling(
+            payload=payload["payload"],
+            headers=payload["headers"],
+            endpoint=self.endpoint,
+            is_cached=payload.get("is_cached", False),
+            should_invoke_endpoint=self.should_invoke_endpoint,
         )
+
+    async def process_chunk(self, chunk) -> None:
+        """Processes a chunk of streaming data.
+
+        Override this method in subclasses if you need custom handling
+        of streaming responses from the API.
+
+        Args:
+            chunk:
+                A portion of the streamed data returned by the API.
+        """
+        pass
+
+    async def stream(self, **kwargs) -> APICalling | None:
+        """Performs a streaming API call with the given arguments.
+
+        Args:
+            **kwargs:
+                Arguments for the request, merged with self.kwargs.
+
+        Returns:
+            `APICalling` | None:
+                An APICalling instance upon success, or None if something
+                goes wrong.
+        """
         try:
-            self.request_model = task_method(**task_params)
+            kwargs["stream"] = True
+            api_call = self.create_api_calling(**kwargs)
+            async for i in api_call.stream():
+                await self.process_chunk(i)
+            return api_call
         except Exception as e:
-            raise ValueError(
-                f"{self.task} requires the following to be provided as input: {e}."
-            )
+            raise ValueError(f"Failed to stream API call: {e}")
 
-        self.model = model
-        self.configs = kwargs
+    async def invoke(self, **kwargs) -> APICalling | None:
+        """Invokes a rate-limited API call with the given arguments.
 
-        self.data_model = self.service.match_data_model(self.task)
+        Args:
+            **kwargs:
+                Arguments for the request, merged with self.kwargs.
 
-    def parse_to_data_model(self, **kwargs):
+        Returns:
+            APICalling | None:
+                The `APICalling` object if successfully invoked and
+                completed; otherwise None.
 
-        if kwargs.get("model") and self.model:
-            if kwargs.get("model") != self.model:
-                raise ValueError(
-                    f"Models are inconsistent. This iModel is for {self.model}"
-                )
+        Raises:
+            ValueError:
+                If the call fails or if an error occurs during invocation.
+        """
+        try:
+            kwargs.pop("stream", None)
+            api_call = self.create_api_calling(**kwargs)
+            if (
+                self.executor.processor is None
+                or self.executor.processor.is_stopped()
+            ):
+                await self.executor.start()
 
-        output = {}
-        for invoke_param, data_model in self.data_model.items():
-            data_model_dict = {}
-            if "model" in data_model.model_fields:
-                data_model_dict["model"] = self.model
-            for key in self.configs:
-                if key in data_model.model_fields:
-                    data_model_dict[key] = self.configs[key]
-            for key in kwargs:
-                if key in data_model.model_fields:
-                    data_model_dict[key] = kwargs[key]
-            output[invoke_param] = data_model(**data_model_dict)
-        return output
+            await self.executor.append(api_call)
+            await self.executor.forward()
+            if api_call.id in self.executor.completed_events:
+                return self.executor.pile.pop(api_call.id)
+        except Exception as e:
+            raise ValueError(f"Failed to invoke API call: {e}")
 
-    async def invoke(self, **kwargs):
-        return await self.request_model.invoke(**kwargs)
+    @property
+    def allowed_roles(self) -> set[str]:
+        """list[str]: Roles that are permissible for this endpoint.
 
-    def list_tasks(self):
-        return self.service.list_tasks()
+        Returns:
+            If the endpoint has an `allowed_roles` attribute, returns that;
+            otherwise, defaults to `{"system", "user", "assistant"}`.
+        """
+        if hasattr(self.endpoint, "allowed_roles"):
+            return self.endpoint.allowed_roles
+        return {"system", "user", "assistant"}
 
     @property
-    def allowed_roles(self):
-        return self.service.allowed_roles
+    def sequential_exchange(self) -> bool:
+        """bool: Indicates whether requests must occur in a strict sequence.
 
-    @property
-    def sequential_exchange(self):
-        """whether the service requires user/assistant exchange"""
-        return self.service.sequential_exchange
+        Returns:
+            True if the endpoint requires sequential handling of
+            messages; False otherwise.
+        """
+        return self.endpoint.sequential_exchange
+
+    def to_dict(self):
+        kwargs = self.kwargs
+        if "kwargs" in self.kwargs:
+            kwargs = self.kwargs["kwargs"]
+        return {
+            "provider": self.endpoint.config.provider,
+            "endpoint": self.endpoint.config.model_dump(),
+            "api_key": (
+                self.api_key_scheme
+                if hasattr(self, "api_key_scheme")
+                else None
+            ),
+            "processor_config": self.executor.config,
+            "invoke_with_endpoint": self.should_invoke_endpoint,
+            **{k: v for k, v in kwargs.items() if k != "api_key"},
+        }
 
+    @classmethod
+    def from_dict(cls, data: dict):
+        provider = data.pop("provider", None)
+        base_url = data.pop("base_url", None)
+        api_key = data.pop("api_key", None)
+        processor_config = data.pop("processor_config", {})
 
-__all__ = ["iModel"]
+        endpoint_config_params = data.pop("endpoint", {})
+        endpoint_ = endpoint_config_params.pop("endpoint", None)
+        endpoint_params = endpoint_config_params.get("endpoint_params", None)
+
+        endpoint = match_endpoint(
+            provider=provider,
+            base_url=base_url,
+            endpoint=endpoint_,
+            endpoint_params=endpoint_params,
+        )
+        endpoint.update_config(**endpoint_config_params)
+        return cls(
+            provider=provider,
+            base_url=base_url,
+            endpoint=endpoint,
+            api_key=api_key,
+            **data,
+            **processor_config,
+        )

lionagi/service/manager.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from lionagi.protocols._concepts import Manager
+from lionagi.utils import is_same_dtype
+
+from .endpoints.chat_completion import ChatCompletionEndPoint
+from .imodel import iModel
+
+
+class iModelManager(Manager):
+
+    def __init__(self, *args: iModel, **kwargs):
+        super().__init__()
+
+        self.registry: dict[str, iModel] = {}
+        if args:
+            if not is_same_dtype(args, iModel):
+                raise TypeError("Input models are not instances of iModel")
+            for model in args:
+                self.register_imodel(model.endpoint.endpoint, model)
+
+        if kwargs:
+            for name, model in kwargs.items():
+                self.register_imodel(name, model)
+
+    @property
+    def chat(self) -> iModel | None:
+        return self.registry.get("chat", None)
+
+    @property
+    def parse(self) -> iModel | None:
+        return self.registry.get("parse", None)
+
+    def register_imodel(self, name: str, model: iModel):
+        if isinstance(model.endpoint, ChatCompletionEndPoint):
+            if name != "parse":
+                self.registry["chat"] = model
+            else:
+                self.registry["parse"] = model
+        elif isinstance(model, iModel):
+            self.registry[name] = model
+        else:
+            raise TypeError("Input model is not an instance of iModel")

lionagi/service/providers/anthropic_/messages.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from lionagi.service.endpoints.chat_completion import ChatCompletionEndPoint
+
+CHAT_COMPLETION_CONFIG = {
+    "provider": "anthropic",
+    "endpoint": "messages",
+    "method": "post",
+    "requires_tokens": True,
+    "openai_compatible": False,
+    "is_invokeable": True,
+    "is_streamable": True,
+    "base_url": "https://api.anthropic.com/v1",
+    "api_version": "2023-06-01",
+    "required_kwargs": {
+        "messages",
+        "model",
+        "max_tokens",
+    },
+    "optional_kwargs": {
+        "metadata",
+        "stop_sequences",
+        "stream",
+        "system",
+        "temperature",
+        "tool_choice",
+        "tools",
+        "top_p",
+        "top_k",
+    },
+    "allowed_roles": ["user", "assistant"],
+}
+
+
+class AnthropicChatCompletionEndPoint(ChatCompletionEndPoint):
+    """
+    Documentation: https://docs.anthropic.com/en/api/
+    """
+
+    def __init__(self, config: dict = CHAT_COMPLETION_CONFIG):
+        super().__init__(config)
+
+    def create_payload(self, **kwargs) -> dict:
+        payload = {}
+        cached = kwargs.get("cached", False)
+        headers = kwargs.get("headers", {})
+        for k, v in kwargs.items():
+            if k in self.acceptable_kwargs:
+                payload[k] = v
+        if "api_key" in kwargs:
+            headers["x-api-key"] = kwargs["api_key"]
+        headers["anthropic-version"] = kwargs.pop(
+            "api_version", self.config.api_version
+        )
+        if "content-type" not in kwargs:
+            headers["content-type"] = "application/json"
+
+        return {
+            "payload": payload,
+            "headers": headers,
+            "cached": cached,
+        }

lionagi/service/providers/groq_/chat_completions.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from lionagi.service.endpoints.chat_completion import ChatCompletionEndPoint
+
+CHAT_COMPLETION_CONFIG = {
+    "provider": "groq",
+    "base_url": "https://api.groq.com/openai/v1",
+    "endpoint": "chat/completions",
+    "method": "post",
+    "openai_compatible": True,
+    "is_invokeable": True,
+    "requires_tokens": True,
+    "is_streamable": True,
+    "required_kwargs": {
+        "messages",
+        "model",
+    },
+    "deprecated_kwargs": {
+        "max_tokens",
+        "function_call",
+        "functions",
+    },
+    "optional_kwargs": {
+        "store",
+        "reasoning_effort",
+        "metadata",
+        "frequency_penalty",
+        "max_completion_tokens",
+        "modalities",
+        "prediction",
+        "audio",
+        "presence_penalty",
+        "seed",
+        "service_tier",
+        "stop",
+        "stream",
+        "stream_options",
+        "temperature",
+        "top_p",
+        "tools",
+        "tool_choice",
+        "parallel_tool_calls",
+        "user",
+    },
+}
+
+
+class GroqChatCompletionEndPoint(ChatCompletionEndPoint):
+    """
+    Documentation: https://console.groq.com/docs/overview
+    """
+
+    def __init__(self, config: dict = CHAT_COMPLETION_CONFIG):
+        super().__init__(config)

lionagi/service/providers/openai_/chat_completions.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from lionagi.service.endpoints.chat_completion import ChatCompletionEndPoint
+
+CHAT_COMPLETION_CONFIG = {
+    "provider": "openai",
+    "base_url": "https://api.openai.com/v1",
+    "endpoint": "chat/completions",
+    "method": "post",
+    "openai_compatible": True,
+    "is_invokeable": True,
+    "requires_tokens": True,
+    "is_streamable": True,
+    "required_kwargs": {
+        "messages",
+        "model",
+    },
+    "deprecated_kwargs": {
+        "max_tokens",
+        "function_call",
+        "functions",
+    },
+    "optional_kwargs": {
+        "store",
+        "reasoning_effort",
+        "metadata",
+        "frequency_penalty",
+        "logit_bias",
+        "logprobs",
+        "top_logprobs",
+        "max_completion_tokens",
+        "n",
+        "modalities",
+        "prediction",
+        "audio",
+        "presence_penalty",
+        "response_format",
+        "seed",
+        "service_tier",
+        "stop",
+        "stream",
+        "stream_options",
+        "temperature",
+        "top_p",
+        "tools",
+        "tool_choice",
+        "parallel_tool_calls",
+        "user",
+    },
+    "allowed_roles": ["user", "assistant", "system"],
+}
+
+
+class OpenAIChatCompletionEndPoint(ChatCompletionEndPoint):
+    """
+    Documentation: https://platform.openai.com/docs/api-reference/chat/create
+    """
+
+    def __init__(self, config: dict = CHAT_COMPLETION_CONFIG):
+        super().__init__(config)