lionagi 0.4.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

lionagi/integrations/perplexity_/PerplexityModel.py

@@ -0,0 +1,269 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from pathlib import Path
+
+import yaml
+from dotenv import load_dotenv
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_serializer,
+    model_validator,
+)
+
+from lionagi.integrations.perplexity_.api_endpoints.api_request import (
+    PerplexityRequest,
+)
+from lionagi.integrations.perplexity_.api_endpoints.chat_completions.request.request_body import (
+    PerplexityChatCompletionRequestBody,
+)
+from lionagi.integrations.perplexity_.api_endpoints.match_response import (
+    match_response,
+)
+from lionagi.service.rate_limiter import RateLimiter, RateLimitError
+from lionagi.service.service_util import invoke_retry
+from lionagi.service.token_calculator import TiktokenCalculator
+
+load_dotenv()
+path = Path(__file__).parent
+
+price_config_file_name = path / "perplexity_price_data.yaml"
+max_output_token_file_name = path / "perplexity_max_output_token_data.yaml"
+
+
+class PerplexityModel(BaseModel):
+    model: str = Field(description="ID of the model to use.")
+
+    request_model: PerplexityRequest = Field(description="Making requests")
+
+    rate_limiter: RateLimiter = Field(
+        description="Rate Limiter to track usage"
+    )
+
+    text_token_calculator: TiktokenCalculator = Field(
+        default=None, description="Token Calculator"
+    )
+
+    estimated_output_len: int = Field(
+        default=0, description="Expected output len before making request"
+    )
+
+    model_config = ConfigDict(extra="forbid")
+
+    @model_validator(mode="before")
+    @classmethod
+    def parse_input(cls, data: dict):
+        if not isinstance(data, dict):
+            raise ValueError("Invalid init param")
+
+        # parse request model
+        request_model_params = {
+            "api_key": data.pop("api_key", None),
+            "endpoint": data.pop("endpoint", None),
+            "method": data.pop("method", None),
+            "content_type": data.pop("content_type", None),
+        }
+        try:
+            api_key = os.getenv(request_model_params["api_key"], None)
+            if api_key:
+                request_model_params["api_key"] = api_key
+        except Exception:
+            pass
+
+        data["request_model"] = PerplexityRequest(**request_model_params)
+
+        # parse rate limiter
+        if "rate_limiter" not in data:
+            rate_limiter_params = {}
+            if limit_tokens := data.pop("limit_tokens", None):
+                rate_limiter_params["limit_tokens"] = limit_tokens
+            if limit_requests := data.pop("limit_requests", None):
+                rate_limiter_params["limit_requests"] = limit_requests
+
+            data["rate_limiter"] = RateLimiter(**rate_limiter_params)
+
+        # parse token calculator
+        try:
+            # Perplexity uses cl100k_base encoding like OpenAI
+            text_calc = TiktokenCalculator(encoding_name="cl100k_base")
+            data["text_token_calculator"] = text_calc
+        except Exception:
+            pass
+
+        return data
+
+    @field_serializer("request_model")
+    def serialize_request_model(self, value: PerplexityRequest):
+        return value.model_dump(exclude_unset=True)
+
+    @invoke_retry(max_retries=3, base_delay=1, max_delay=60)
+    async def invoke(
+        self,
+        request_body: PerplexityChatCompletionRequestBody,
+        estimated_output_len: int = 0,
+        output_file=None,
+        parse_response=True,
+    ):
+        if request_model := getattr(request_body, "model"):
+            if request_model != self.model:
+                raise ValueError(
+                    f"Request model does not match. Model is {self.model}, but request is made for {request_model}."
+                )
+
+        # check remaining rate limit
+        input_token_len = await self.get_input_token_len(request_body)
+
+        if getattr(request_body, "max_tokens", None):
+            estimated_output_len = request_body.max_tokens
+
+        invoke_viability_result = self.verify_invoke_viability(
+            input_tokens_len=input_token_len,
+            estimated_output_len=estimated_output_len,
+        )
+        if not invoke_viability_result:
+            raise RateLimitError(
+                message="Rate limit reached for requests",
+                input_token_len=input_token_len,
+                estimated_output_len=estimated_output_len,
+            )
+
+        try:
+            if getattr(request_body, "stream", None):
+                return await self.stream(
+                    request_body,
+                    output_file=output_file,
+                    parse_response=parse_response,
+                )
+
+            response_body, response_headers = await self.request_model.invoke(
+                json_data=request_body,
+                output_file=output_file,
+                with_response_header=True,
+                parse_response=False,
+            )
+
+            if response_body:
+                # Update rate limit based on usage
+                if response_body.get("usage"):
+                    total_token_usage = response_body["usage"]["total_tokens"]
+                    if date_str := response_headers.get("date"):
+                        self.rate_limiter.update_rate_limit(
+                            date_str, total_token_usage
+                        )
+                    else:
+                        self.rate_limiter.update_rate_limit(
+                            None, total_token_usage
+                        )
+                else:
+                    self.rate_limiter.update_rate_limit(None)
+
+            if parse_response:
+                return match_response(self.request_model, response_body)
+            else:
+                return response_body
+
+        except Exception as e:
+            raise e
+
+    async def stream(
+        self,
+        request_body: PerplexityChatCompletionRequestBody,
+        output_file=None,
+        parse_response=True,
+        verbose=True,
+    ):
+        response_chunks = []
+        response_headers = None
+
+        async for chunk in self.request_model.stream(
+            json_data=request_body,
+            output_file=output_file,
+            with_response_header=True,
+            verbose=verbose,
+        ):
+            if isinstance(chunk, dict):
+                if "headers" in chunk:
+                    response_headers = chunk["headers"]
+                else:
+                    response_chunks.append(chunk)
+
+        # Update rate limit if we have usage information
+        if response_chunks and response_chunks[-1].get("usage"):
+            total_token_usage = response_chunks[-1]["usage"]["total_tokens"]
+            if date_str := response_headers.get("date"):
+                self.rate_limiter.update_rate_limit(
+                    date_str, total_token_usage
+                )
+            else:
+                self.rate_limiter.update_rate_limit(None, total_token_usage)
+
+        if parse_response:
+            return match_response(self.request_model, response_chunks)
+        else:
+            return response_chunks
+
+    async def get_input_token_len(
+        self, request_body: PerplexityChatCompletionRequestBody
+    ):
+        if request_model := getattr(request_body, "model"):
+            if request_model != self.model:
+                raise ValueError(
+                    f"Request model does not match. Model is {self.model}, but request is made for {request_model}."
+                )
+
+        total_tokens = 0
+        for message in request_body.messages:
+            total_tokens += self.text_token_calculator.calculate(
+                message.content
+            )
+
+        return total_tokens
+
+    def verify_invoke_viability(
+        self, input_tokens_len: int = 0, estimated_output_len: int = 0
+    ):
+        self.rate_limiter.release_tokens()
+
+        estimated_output_len = (
+            estimated_output_len
+            if estimated_output_len != 0
+            else self.estimated_output_len
+        )
+        if estimated_output_len == 0:
+            with open(max_output_token_file_name) as file:
+                output_token_config = yaml.safe_load(file)
+                estimated_output_len = output_token_config.get(self.model, 0)
+                self.estimated_output_len = estimated_output_len
+
+        if self.rate_limiter.check_availability(
+            input_tokens_len, estimated_output_len
+        ):
+            return True
+        else:
+            return False
+
+    def estimate_text_price(
+        self,
+        input_text: str,
+        estimated_num_of_output_tokens: int = 0,
+    ):
+        if self.text_token_calculator is None:
+            raise ValueError("Token calculator not available")
+
+        num_of_input_tokens = self.text_token_calculator.calculate(input_text)
+
+        with open(price_config_file_name) as file:
+            price_config = yaml.safe_load(file)
+
+        model_price_info_dict = price_config["model"][self.model]
+        estimated_price = (
+            model_price_info_dict["input_tokens"] * num_of_input_tokens
+            + model_price_info_dict["output_tokens"]
+            * estimated_num_of_output_tokens
+        )
+
+        return estimated_price

lionagi/integrations/perplexity_/PerplexityService.py

@@ -0,0 +1,113 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import inspect
+
+from dotenv import load_dotenv
+
+from lionagi.integrations.perplexity_.api_endpoints.chat_completions.request.request_body import (
+    PerplexityChatCompletionRequestBody,
+)
+from lionagi.service import Service, register_service
+
+from .PerplexityModel import PerplexityModel
+
+load_dotenv()
+
+
+@register_service
+class PerplexityService(Service):
+    def __init__(
+        self,
+        api_key: str,
+        name: str = None,
+    ):
+        super().__setattr__("_initialized", False)
+        self.api_key = api_key
+        self.name = name
+        self.rate_limiters = {}  # model: RateLimiter
+        super().__setattr__("_initialized", True)
+
+    def __setattr__(self, key, value):
+        if getattr(self, "_initialized", False) and key in [
+            "api_key",
+        ]:
+            raise AttributeError(
+                f"Cannot modify '{key}' after initialization. "
+                f"Please set a new service object for new keys."
+            )
+        super().__setattr__(key, value)
+
+    def check_rate_limiter(
+        self,
+        perplexity_model: PerplexityModel,
+        limit_requests: int = None,
+        limit_tokens: int = None,
+    ):
+        # Map model versions to their base models for shared rate limiting
+        shared_models = {
+            "llama-3.1-sonar-small-128k-online": "llama-3.1-sonar-small",
+            "llama-3.1-sonar-medium-128k-online": "llama-3.1-sonar-medium",
+            "llama-3.1-sonar-large-128k-online": "llama-3.1-sonar-large",
+        }
+
+        if perplexity_model.model in shared_models:
+            model = shared_models[perplexity_model.model]
+        else:
+            model = perplexity_model.model
+
+        if model not in self.rate_limiters:
+            self.rate_limiters[model] = perplexity_model.rate_limiter
+        else:
+            perplexity_model.rate_limiter = self.rate_limiters[model]
+            if limit_requests:
+                perplexity_model.rate_limiter.limit_requests = limit_requests
+            if limit_tokens:
+                perplexity_model.rate_limiter.limit_tokens = limit_tokens
+
+        return perplexity_model
+
+    @staticmethod
+    def match_data_model(task_name: str) -> dict:
+        """Match task name to appropriate request and response models."""
+        if task_name == "create_chat_completion":
+            return {"request_body": PerplexityChatCompletionRequestBody}
+        raise ValueError(f"No data models found for task: {task_name}")
+
+    @classmethod
+    def list_tasks(cls):
+        methods = []
+        for name, member in inspect.getmembers(
+            cls, predicate=inspect.isfunction
+        ):
+            if name not in [
+                "__init__",
+                "__setattr__",
+                "check_rate_limiter",
+                "match_data_model",
+            ]:
+                methods.append(name)
+        return methods
+
+    # Chat Completions
+    def create_chat_completion(
+        self, model: str, limit_tokens: int = None, limit_requests: int = None
+    ):
+        model_obj = PerplexityModel(
+            model=model,
+            api_key=self.api_key,
+            endpoint="chat/completions",
+            method="POST",
+            content_type="application/json",
+            limit_tokens=limit_tokens,
+            limit_requests=limit_requests,
+        )
+
+        return self.check_rate_limiter(
+            model_obj, limit_requests=limit_requests, limit_tokens=limit_tokens
+        )
+
+    @property
+    def allowed_roles(self):
+        return ["user", "assistant", "system"]

lionagi/integrations/perplexity_/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0

lionagi/integrations/perplexity_/api_endpoints/api_request.py

@@ -0,0 +1,171 @@
+import json
+from collections.abc import AsyncGenerator
+from typing import Any, Dict, Optional, Tuple, Union
+
+import aiohttp
+from pydantic import BaseModel, ConfigDict, Field
+
+from lionagi.integrations.perplexity_.api_endpoints.data_models import (
+    PerplexityEndpointRequestBody,
+)
+
+
+class PerplexityRequest(BaseModel):
+    """Handler for making requests to the Perplexity API."""
+
+    api_key: str = Field(
+        description="API key for authentication", exclude=True
+    )
+    endpoint: str = Field(description="Endpoint for request")
+    method: str = Field(description="HTTP method")
+    content_type: str | None = "application/json"
+    base_url: str = "https://api.perplexity.ai"
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="allow",  # Allow extra attributes for mocking in tests
+    )
+
+    async def invoke(
+        self,
+        json_data: None | (
+            dict[str, Any] | PerplexityEndpointRequestBody
+        ) = None,
+        form_data: BaseModel | None = None,
+        output_file: str | None = None,
+        with_response_header: bool = False,
+        parse_response: bool = True,
+    ) -> dict[str, Any] | tuple[dict[str, Any], dict[str, str]] | bytes | None:
+        """Make a request to the Perplexity API."""
+        url = f"{self.base_url}/{self.endpoint}"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+        }
+
+        if self.content_type:
+            headers["Content-Type"] = self.content_type
+
+        if isinstance(json_data, PerplexityEndpointRequestBody):
+            json_data = json_data.model_dump(exclude_unset=True)
+
+        async with aiohttp.ClientSession() as client:
+            response = await client.request(
+                method=self.method,
+                url=url,
+                headers=headers,
+                json=json_data,
+                data=form_data.model_dump() if form_data else None,
+            )
+            async with response:
+                if response.status != 200:
+                    try:
+                        error_body = await response.json()
+                        error_msg = error_body.get("error", {}).get(
+                            "message", str(error_body)
+                        )
+                    except Exception:
+                        error_msg = await response.text()
+                    raise Exception(
+                        f"API request failed with status {response.status}: {error_msg}"
+                    )
+
+                if output_file:
+                    with open(output_file, "wb") as f:
+                        async for chunk in response.content.iter_chunked(8192):
+                            f.write(chunk)
+                    return None
+
+                if parse_response:
+                    response_body = await response.json()
+                else:
+                    response_body = await response.text()
+                    try:
+                        response_body = json.loads(response_body)
+                    except json.JSONDecodeError:
+                        pass
+
+                if with_response_header:
+                    headers = {
+                        k.lower(): v for k, v in response.headers.items()
+                    }
+                    return response_body, headers
+                return response_body
+
+    async def stream(
+        self,
+        json_data: None | (
+            dict[str, Any] | PerplexityEndpointRequestBody
+        ) = None,
+        output_file: str | None = None,
+        with_response_header: bool = False,
+        verbose: bool = True,
+    ) -> AsyncGenerator[dict[str, Any], None]:
+        """Stream responses from the Perplexity API."""
+        if isinstance(json_data, PerplexityEndpointRequestBody):
+            json_data = json_data.model_dump(exclude_unset=True)
+
+        if not json_data.get("stream", False):
+            json_data["stream"] = True
+
+        url = f"{self.base_url}/{self.endpoint}"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream",
+        }
+
+        async with aiohttp.ClientSession() as client:
+            response = await client.post(url, headers=headers, json=json_data)
+            async with response:
+                if response.status != 200:
+                    try:
+                        error_body = await response.json()
+                        error_msg = error_body.get("error", {}).get(
+                            "message", str(error_body)
+                        )
+                    except Exception:
+                        error_msg = await response.text()
+                    raise Exception(
+                        f"API request failed with status {response.status}: {error_msg}"
+                    )
+
+                if with_response_header:
+                    headers = {
+                        k.lower(): v for k, v in response.headers.items()
+                    }
+                    yield {"headers": headers}
+
+                file_handle = None
+                if output_file:
+                    try:
+                        file_handle = open(output_file, "w")
+                    except Exception as e:
+                        raise ValueError(
+                            f"Failed to open output file {output_file}: {e}"
+                        )
+
+                try:
+                    async for chunk in response.content:
+                        if chunk:
+                            chunk_str = chunk.decode("utf-8").strip()
+                            if chunk_str.startswith("data: "):
+                                chunk_str = chunk_str[
+                                    6:
+                                ]  # Remove "data: " prefix
+                            try:
+                                chunk_data = json.loads(chunk_str)
+                                if file_handle:
+                                    file_handle.write(
+                                        json.dumps(chunk_data) + "\n"
+                                    )
+                                if verbose and "choices" in chunk_data:
+                                    content = chunk_data["choices"][0][
+                                        "delta"
+                                    ]["content"]
+                                    print(content, end="", flush=True)
+                                yield chunk_data
+                            except json.JSONDecodeError:
+                                continue
+                finally:
+                    if file_handle:
+                        file_handle.close()

lionagi/integrations/perplexity_/api_endpoints/chat_completions/request/request_body.py

@@ -0,0 +1,121 @@
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from lionagi.integrations.perplexity_.api_endpoints.data_models import (
+    PerplexityEndpointRequestBody,
+)
+
+
+class Message(BaseModel):
+    """A message in the chat completion request."""
+
+    role: str = Field(description="The role of the message author.")
+    content: str = Field(description="The content of the message.")
+
+    model_config = ConfigDict(extra="forbid")
+
+    @field_validator("role")
+    @classmethod
+    def validate_role(cls, role: str) -> str:
+        valid_roles = ["system", "user", "assistant"]
+        if role not in valid_roles:
+            raise ValueError(f"Role must be one of {valid_roles}")
+        return role
+
+    @field_validator("content")
+    @classmethod
+    def validate_content(cls, content: str) -> str:
+        if not content or not content.strip():
+            raise ValueError("Content cannot be empty")
+        return content
+
+
+class PerplexityChatCompletionRequestBody(PerplexityEndpointRequestBody):
+    """Request body for chat completion requests."""
+
+    model: str = Field(description="ID of the model to use.")
+    messages: list[Message] = Field(
+        description="A list of messages comprising the conversation so far."
+    )
+    max_tokens: int | None = Field(
+        default=None,
+        description="The maximum number of tokens to generate in the completion.",
+        ge=1,
+    )
+    temperature: float | None = Field(
+        default=0.2,
+        description="What sampling temperature to use, between 0 and 2.",
+        ge=0,
+        lt=2,
+    )
+    top_p: float | None = Field(
+        default=0.9,
+        description="An alternative to sampling with temperature, called nucleus sampling.",
+        ge=0,
+        le=1,
+    )
+    search_domain_filter: list[str] | None = Field(
+        default=None,
+        description="List of domains to limit citations to.",
+    )
+    return_images: bool | None = Field(
+        default=False,
+        description="Whether to return images in the response.",
+    )
+    return_related_questions: bool | None = Field(
+        default=False,
+        description="Whether to return related questions in the response.",
+    )
+    search_recency_filter: str | None = Field(
+        default=None,
+        description="Filter for search results recency (month, week, day, hour).",
+    )
+    top_k: int | None = Field(
+        default=0,
+        description="The number of tokens to keep for highest top-k filtering.",
+        ge=0,
+        le=2048,
+    )
+    stream: bool | None = Field(
+        default=False,
+        description="Whether to stream back partial progress.",
+    )
+    presence_penalty: float | None = Field(
+        default=0,
+        description="Penalty for new tokens based on their presence in text so far.",
+        ge=-2,
+        le=2,
+    )
+    frequency_penalty: float | None = Field(
+        default=1,
+        description="Penalty for new tokens based on their frequency in text so far.",
+        gt=0,
+    )
+
+    model_config = ConfigDict(extra="forbid")
+
+    @field_validator("search_recency_filter")
+    @classmethod
+    def validate_search_recency_filter(cls, value: str | None) -> str | None:
+        if value is not None:
+            valid_filters = ["month", "week", "day", "hour"]
+            if value not in valid_filters:
+                raise ValueError(
+                    f"Search recency filter must be one of {valid_filters}"
+                )
+        return value
+
+    @field_validator("search_domain_filter")
+    @classmethod
+    def validate_search_domain_filter(
+        cls, value: list[str] | None
+    ) -> list[str] | None:
+        if value is not None and len(value) > 3:
+            raise ValueError("Search domain filter is limited to 3 domains")
+        return value
+
+    @field_validator("messages")
+    @classmethod
+    def validate_messages(cls, messages: list[Message]) -> list[Message]:
+        if not messages:
+            raise ValueError("At least one message is required")
+        return messages