aiverify-moonshot 0.4.10__py3-none-any.whl → 0.5.0__py3-none-any.whl

moonshot/src/connectors/connector.py

@@ -3,16 +3,32 @@ from __future__ import annotations
 import asyncio
 import time
 from abc import abstractmethod
-from asyncio import sleep
-from functools import wraps
+from functools import partial, wraps
 from pathlib import Path
 from typing import Callable
 
+from tenacity import RetryCallState, retry, stop_after_attempt, wait_random_exponential
+
 from moonshot.src.configs.env_variables import EnvVariables
 from moonshot.src.connectors.connector_prompt_arguments import ConnectorPromptArguments
+from moonshot.src.connectors.connector_response import ConnectorResponse
 from moonshot.src.connectors_endpoints.connector_endpoint_arguments import (
     ConnectorEndpointArguments,
 )
+from moonshot.src.messages_constants import (
+    CONNECTOR_CREATE_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR,
+    CONNECTOR_CREATE_ERROR,
+    CONNECTOR_GET_AVAILABLE_ITEMS_ERROR,
+    CONNECTOR_GET_PREDICTION_ARGUMENTS_CONNECTOR_VALIDATION_ERROR,
+    CONNECTOR_GET_PREDICTION_ARGUMENTS_GENERATED_PROMPT_VALIDATION_ERROR,
+    CONNECTOR_GET_PREDICTION_ERROR,
+    CONNECTOR_GET_PREDICTION_INFO,
+    CONNECTOR_GET_PREDICTION_TIME_TAKEN_INFO,
+    CONNECTOR_LOAD_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR,
+    CONNECTOR_LOAD_CONNECTOR_INSTANCE_RUNTIME_ERROR,
+    CONNECTOR_PERFORM_RETRY_CALLBACK_ERROR,
+    CONNECTOR_SET_SYSTEM_PROMPT_VALIDATION_ERROR,
+)
 from moonshot.src.storage.storage import Storage
 from moonshot.src.utils.import_modules import get_instance
 from moonshot.src.utils.log import configure_logger
@@ -21,43 +37,56 @@ from moonshot.src.utils.log import configure_logger
 logger = configure_logger(__name__)
 
 
+def perform_retry_callback(connector_id: str, retry_state: RetryCallState) -> None:
+    """
+    Callback function to log retry attempts with detailed information.
+
+    This function is called by the tenacity library before each retry attempt.
+    It logs the retry attempt number, the sleep time before the next attempt,
+    the error message that caused the retry, and the connector ID.
+
+    Args:
+        connector_id (str): The ID of the connector.
+        retry_state (RetryCallState): The state of the retry call, which includes
+            information about the current attempt, the exception raised, and the next action.
+    """
+    sleep_time = retry_state.idle_for if retry_state else 0
+    exception = (
+        retry_state.outcome.exception() if retry_state.outcome else "Unknown exception"
+    )
+    logger.error(
+        CONNECTOR_PERFORM_RETRY_CALLBACK_ERROR.format(
+            connector_id=connector_id,
+            attempt_no=retry_state.attempt_number,
+            sleep=f"{sleep_time:.2f}",
+            message=str(exception),
+        )
+    )
+
+
 def perform_retry(func):
     """
-    A decorator to perform retries on a function.
+    A decorator to perform retries on a function using tenacity.
 
-    This decorator wraps a function to enable retrying the function call
-    if it fails. The number of retries and the delay between retries
-    are determined by the `retries_times` and `allow_retries`
-    attributes of the class instance.
+    This decorator wraps an asynchronous function to enable retrying the function call
+    if it fails. The number of attempts and the delay between attempts
+    are determined by the `max_attempts` attribute of the class instance.
 
-    Parameters:
-    - func (Callable): The function to be wrapped and retried.
+    Args:
+        func (Callable): The asynchronous function to be wrapped and retried.
 
     Returns:
-    - Callable: A wrapper function that includes retry logic.
+        Callable: A wrapper function that includes retry logic.
     """
 
     async def wrapper(self, *args, **kwargs):
-        if self.allow_retries:
-            retry_count = 0
-            base_delay = 1
-            while retry_count <= self.retries_times:
-                # Perform the request
-                try:
-                    return await func(self, *args, **kwargs)
-                except Exception as exc:
-                    logger.warning(f"Operation failed. {str(exc)} - Retrying...")
-
-                # Perform retry
-                retry_count += 1
-                if retry_count <= self.retries_times:
-                    delay = base_delay * (2**retry_count)
-                    logger.warning(
-                        f"Attempt {retry_count}, Retrying in {delay} seconds..."
-                    )
-                    await sleep(delay)
-            # Raise an exception
-            raise ConnectionError("Failed to get response.")
+        retry_decorator = retry(
+            wait=wait_random_exponential(min=1, max=60),
+            stop=stop_after_attempt(self.max_attempts),
+            after=partial(perform_retry_callback, self.id),
+            reraise=True,
+        )
+        return await retry_decorator(func)(self, *args, **kwargs)
 
     return wrapper
 
@@ -70,6 +99,7 @@ class Connector:
         self.token = ep_args.token
         self.max_concurrency = ep_args.max_concurrency
         self.max_calls_per_second = ep_args.max_calls_per_second
+        self.model = ep_args.model
         self.params = ep_args.params
 
         # Rate limiting
@@ -79,21 +109,19 @@ class Connector:
         self.updated_at = time.time()
         self.semaphore = asyncio.Semaphore(ep_args.max_concurrency)
 
-        # Set Prompts if they exists
+        # Set Prompts if they exist
         self.pre_prompt = ep_args.params.get("pre_prompt", "")
         self.post_prompt = ep_args.params.get("post_prompt", "")
         self.system_prompt = ep_args.params.get("system_prompt", "")
 
         # Connection timeout
         self.timeout = ep_args.params.get("timeout", 600)
-        self.allow_retries = ep_args.params.get("allow_retries", True)
-        self.retries_times = ep_args.params.get("num_of_retries", 3)
+        self.max_attempts = ep_args.params.get("max_attempts", 3)
 
         # Optional params
         excluded_keys = {
-            "allow_retries",
             "timeout",
-            "num_of_retries",
+            "max_attempts",
             "pre_prompt",
             "post_prompt",
             "system_prompt",
@@ -113,10 +141,6 @@ class Connector:
 
         The method updates the `updated_at` attribute to the current time after tokens are added, ensuring that
         the next token addition will be calculated based on the correct elapsed time.
-
-        Usage:
-            # Assume `self` is an instance of a class with `rate_limiter`, `tokens`, and `updated_at` attributes.
-            await self.add_tokens()
         """
         now = time.time()
         elapsed = now - self.updated_at
@@ -143,11 +167,6 @@ class Connector:
 
         Returns:
             Callable: The decorated function wrapped with rate limiting logic.
-
-        Usage:
-            @rate_limited
-            async def some_async_function(*args, **kwargs):
-                # Function implementation
         """
 
         @wraps(func)
@@ -168,9 +187,9 @@ class Connector:
         return wrapper
 
     @abstractmethod
-    async def get_response(self, prompt: str) -> str:
+    async def get_response(self, prompt: str) -> ConnectorResponse:
         """
-        Abstract method to be implemented by subclasses to get a response from the connector.
+        Abstract method to be implemented by subclasses to obtain a response from the connector.
 
         This method should asynchronously send a prompt to the connector's API and return the response.
 
@@ -178,14 +197,15 @@ class Connector:
             prompt (str): The input prompt to be sent to the connector.
 
         Returns:
-            str: The response received from the connector.
+            ConnectorResponse: An instance containing the response received from the connector and any
+            additional metadata.
         """
         pass
 
     @classmethod
     def load(cls, ep_args: ConnectorEndpointArguments) -> Connector:
         """
-        This method dynamically loads a connector instance based on the provided endpoint arguments.
+        Dynamically loads a connector instance based on the provided endpoint arguments.
 
         The connector type specified in the `ep_args` is used to dynamically load the corresponding
         connector class. The connector is then instantiated with the provided endpoint arguments. If the
@@ -201,17 +221,24 @@ class Connector:
         Raises:
             RuntimeError: If the specified connector type does not match any available connector classes.
         """
+        if ep_args is None or not isinstance(ep_args, ConnectorEndpointArguments):
+            raise ValueError(
+                CONNECTOR_LOAD_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR
+            )
+
         connector_instance = get_instance(
             ep_args.connector_type,
             Storage.get_filepath(
                 EnvVariables.CONNECTORS.name, ep_args.connector_type, "py"
             ),
         )
-        if connector_instance:
+        if connector_instance and isinstance(connector_instance, Callable):
             return connector_instance(ep_args)
         else:
             raise RuntimeError(
-                f"Unable to get defined connector instance - {ep_args.connector_type}"
+                CONNECTOR_LOAD_CONNECTOR_INSTANCE_RUNTIME_ERROR.format(
+                    message=ep_args.connector_type
+                )
             )
 
     @staticmethod
@@ -221,7 +248,7 @@ class Connector:
 
         This method takes a ConnectorEndpointArguments object, which contains the necessary information
         to initialize and return a Connector object. The Connector object is created by calling the
-        `load_connector` method, which dynamically loads and initializes the connector based on the
+        `load` method, which dynamically loads and initializes the connector based on the
         endpoint arguments provided.
 
         Args:
@@ -229,12 +256,20 @@ class Connector:
 
         Returns:
             Connector: An initialized Connector object based on the provided endpoint arguments.
+
+        Raises:
+            ValueError: If the provided endpoint arguments are invalid.
+            Exception: If there is an error during the creation of the connector.
         """
         try:
+            if ep_args is None or not isinstance(ep_args, ConnectorEndpointArguments):
+                raise ValueError(
+                    CONNECTOR_CREATE_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR
+                )
             return Connector.load(ep_args)
 
         except Exception as e:
-            logger.error(f"Failed to create connector: {str(e)}")
+            logger.error(CONNECTOR_CREATE_ERROR.format(message=str(e)))
             raise e
 
     @staticmethod
@@ -243,7 +278,7 @@ class Connector:
         Fetches a list of all available connector types.
 
         This method employs the `get_connectors` method to locate all Python files in the directory
-        defined by the `EnvironmentVars.CONNECTORS` environment variable. It subsequently excludes any files that are
+        defined by the `EnvVariables.CONNECTORS` environment variable. It subsequently excludes any files that are
         not intended to be exposed as connectors (those containing "__" in their names). The method yields a list of the
         names of these connector types.
 
@@ -261,7 +296,7 @@ class Connector:
             ]
 
         except Exception as e:
-            logger.error(f"Failed to get available connectors: {str(e)}")
+            logger.error(CONNECTOR_GET_AVAILABLE_ITEMS_ERROR.format(message=str(e)))
             raise e
 
     @staticmethod
@@ -277,7 +312,7 @@ class Connector:
         object, which is used to generate the prediction. The method also optionally takes a `prompt_callback` function,
         which is called after the prediction is generated.
 
-        The method first prints a message indicating that it is predicting the prompt. It then records the start time
+        The method logs a message indicating that it is predicting the prompt. It then records the start time
         and uses the `connector` to generate a prediction for the `generated_prompt`. The duration of the prediction
         is calculated and stored in the `generated_prompt`.
 
@@ -297,9 +332,24 @@ class Connector:
         Raises:
             Exception: If there is an error during prediction.
         """
+        if generated_prompt is None or not isinstance(
+            generated_prompt, ConnectorPromptArguments
+        ):
+            raise ValueError(
+                CONNECTOR_GET_PREDICTION_ARGUMENTS_GENERATED_PROMPT_VALIDATION_ERROR
+            )
+
+        if connector is None or not isinstance(connector, Connector):
+            raise ValueError(
+                CONNECTOR_GET_PREDICTION_ARGUMENTS_CONNECTOR_VALIDATION_ERROR
+            )
+
         try:
             logger.info(
-                f"Predicting prompt {generated_prompt.prompt_index} [{connector.id}]"
+                CONNECTOR_GET_PREDICTION_INFO.format(
+                    connector_id=connector.id,
+                    prompt_index=generated_prompt.prompt_index,
+                )
             )
 
             start_time = time.perf_counter()
@@ -308,7 +358,11 @@ class Connector:
             )
             generated_prompt.duration = time.perf_counter() - start_time
             logger.debug(
-                f"[Prompt {generated_prompt.prompt_index}] took {generated_prompt.duration:.4f}s"
+                CONNECTOR_GET_PREDICTION_TIME_TAKEN_INFO.format(
+                    connector_id=connector.id,
+                    prompt_index=generated_prompt.prompt_index,
+                    prompt_duration=f"{generated_prompt.duration:.4f}",
+                )
             )
 
             # Call prompt callback
@@ -319,17 +373,28 @@ class Connector:
             return generated_prompt
 
         except Exception as e:
-            logger.error(f"Failed to get prediction: {str(e)}")
+            logger.error(
+                CONNECTOR_GET_PREDICTION_ERROR.format(
+                    connector_id=connector.id,
+                    prompt_index=generated_prompt.prompt_index,
+                    message=str(e),
+                )
+            )
             raise e
 
     def set_system_prompt(self, system_prompt: str) -> None:
         """
-        Assigns a new system prompt to this connector instance.
+        Sets a new system prompt for this connector instance.
 
-        The system prompt serves as a preconfigured command or message that the connector can use to initiate
-        interactions or execute specific operations.
+        The system prompt is a predefined message or command that the connector can use to start interactions
+        or perform specific tasks.
 
-        Parameters:
+        Args:
             system_prompt (str): The new system prompt to set for this connector.
+
+        Raises:
+            ValueError: If the provided system prompt is not a string or is None.
         """
+        if system_prompt is None or not isinstance(system_prompt, str):
+            raise ValueError(CONNECTOR_SET_SYSTEM_PROMPT_VALIDATION_ERROR)
         self.system_prompt = system_prompt

moonshot/src/connectors/connector_prompt_arguments.py

@@ -1,17 +1,27 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Annotated, Any
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+
+from moonshot.src.connectors.connector_response import ConnectorResponse
 
 
 class ConnectorPromptArguments(BaseModel):
-    prompt_index: int  # The index of the prompt in the dataset
+    class Config:
+        arbitrary_types_allowed = True
+
+    prompt_index: Annotated[
+        int, Field(strict=True, ge=0)
+    ]  # The index of the prompt in the dataset
 
     prompt: str  # The actual prompt text
 
     target: Any  # The target response for the prompt
 
-    predicted_results: Any = ""  # The predicted results, default is an empty string
+    predicted_results: ConnectorResponse | None = (
+        None  # The predicted results, default is None
+    )
 
-    duration: float = 0.0  # The duration it took to get the results, default is 0.0
+    # The duration it took to get the results, must be a positive float
+    duration: Annotated[float, Field(strict=True, ge=0.0)] = 0.0

moonshot/src/connectors/connector_response.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel
+
+
+class ConnectorResponse(BaseModel):
+    response: str = ""
+    context: list = []
+
+    def to_dict(self) -> dict:
+        """
+        Converts the ConnectorResponse instance to a dictionary.
+
+        Returns:
+            dict: A dictionary representation of the ConnectorResponse instance.
+        """
+        return {"response": self.response, "context": self.context}

moonshot/src/connectors_endpoints/connector_endpoint.py

@@ -1,12 +1,20 @@
 from pathlib import Path
 
-from pydantic import validate_call
+from pydantic import constr, validate_call
 from slugify import slugify
 
 from moonshot.src.configs.env_variables import EnvVariables
 from moonshot.src.connectors_endpoints.connector_endpoint_arguments import (
     ConnectorEndpointArguments,
 )
+from moonshot.src.messages_constants import (
+    CONNECTOR_ENDPOINT_CREATE_ERROR,
+    CONNECTOR_ENDPOINT_DELETE_ERROR,
+    CONNECTOR_ENDPOINT_GET_AVAILABLE_ITEMS_ERROR,
+    CONNECTOR_ENDPOINT_READ_ERROR,
+    CONNECTOR_ENDPOINT_READ_INVALID,
+    CONNECTOR_ENDPOINT_UPDATE_ERROR,
+)
 from moonshot.src.storage.storage import Storage
 from moonshot.src.utils.log import configure_logger
 
@@ -16,6 +24,7 @@ logger = configure_logger(__name__)
 
 class ConnectorEndpoint:
     @staticmethod
+    @validate_call
     def create(ep_args: ConnectorEndpointArguments) -> str:
         """
         Creates a new connector endpoint and stores its details as a JSON object.
@@ -47,6 +56,7 @@ class ConnectorEndpoint:
                 "token": ep_args.token,
                 "max_calls_per_second": ep_args.max_calls_per_second,
                 "max_concurrency": ep_args.max_concurrency,
+                "model": ep_args.model,
                 "params": ep_args.params,
             }
 
@@ -57,12 +67,12 @@ class ConnectorEndpoint:
             return ep_id
 
         except Exception as e:
-            logger.error(f"Failed to create endpoint: {str(e)}")
+            logger.error(CONNECTOR_ENDPOINT_CREATE_ERROR.format(message=str(e)))
             raise e
 
     @staticmethod
     @validate_call
-    def read(ep_id: str) -> ConnectorEndpointArguments:
+    def read(ep_id: constr(min_length=1)) -> ConnectorEndpointArguments:
         """
         Retrieves the details of a specified endpoint by its ID.
 
@@ -72,27 +82,24 @@ class ConnectorEndpoint:
         any other error occurs, an exception is raised with an appropriate error message.
 
         Args:
-            ep_id (str): The unique identifier of the endpoint whose details are to be retrieved.
+            ep_id (constr(min_length=1)): The unique identifier of the endpoint whose details are to be retrieved.
 
         Returns:
             ConnectorEndpointArguments: An instance filled with the endpoint's details.
 
         Raises:
-            RuntimeError: If the endpoint ID is empty or the specified endpoint does not exist.
+            RuntimeError: If the specified endpoint does not exist.
             Exception: For any issues encountered during the file reading or data parsing process.
         """
         try:
-            if not ep_id:
-                raise RuntimeError("Connector Endpoint ID is empty.")
-
             endpoint_details = ConnectorEndpoint._read_endpoint(ep_id)
             if not endpoint_details:
-                raise RuntimeError(f"Endpoint with ID '{ep_id}' does not exist.")
+                raise RuntimeError(CONNECTOR_ENDPOINT_READ_INVALID.format(ep_id=ep_id))
 
             return ConnectorEndpointArguments(**endpoint_details)
 
         except Exception as e:
-            logger.error(f"Failed to read endpoint: {str(e)}")
+            logger.error(CONNECTOR_ENDPOINT_READ_ERROR.format(message=str(e)))
             raise e
 
     @staticmethod
@@ -124,6 +131,7 @@ class ConnectorEndpoint:
         return connector_endpoint_info
 
     @staticmethod
+    @validate_call
     def update(ep_args: ConnectorEndpointArguments) -> bool:
         """
         Updates the endpoint information in the storage based on the provided ConnectorEndpointArguments object.
@@ -138,11 +146,10 @@ class ConnectorEndpoint:
             ep_args (ConnectorEndpointArguments): The object encapsulating the updated attributes of the endpoint.
 
         Returns:
-            bool: Indicates whether the update operation was successful. Returns True if the update was successfully
-            persisted to the storage; otherwise, an exception is raised.
+            bool: True if the update was successfully persisted to the storage; otherwise, an exception is raised.
 
         Raises:
-            Exception: Signifies a failure in the update process, potentially due to issues with data serialization or
+            Exception: If the update process encounters an error, potentially due to issues with data serialization or
             storage access.
         """
         try:
@@ -160,24 +167,24 @@ class ConnectorEndpoint:
             return True
 
         except Exception as e:
-            logger.error(f"Failed to update endpoint: {str(e)}")
+            logger.error(CONNECTOR_ENDPOINT_UPDATE_ERROR.format(message=str(e)))
             raise e
 
     @staticmethod
     @validate_call
-    def delete(ep_id: str) -> bool:
+    def delete(ep_id: constr(min_length=1)) -> bool:
         """
         Deletes the endpoint with the specified ID.
 
         This method attempts to delete the endpoint corresponding to the given ID from the storage.
-        If the deletion is successful, it returns True. If an error occurs, it prints an error message
+        If the deletion is successful, it returns True. If an error occurs, it logs an error message
         and re-raises the exception.
 
         Args:
-            ep_id (str): The unique identifier of the endpoint to be deleted.
+            ep_id (constr(min_length=1)): The unique identifier of the endpoint to be deleted
 
         Returns:
-            bool: True if the endpoint was successfully deleted.
+            bool: True if the endpoint was successfully deleted; otherwise, an exception is raised.
 
         Raises:
             Exception: If the deletion process encounters an error.
@@ -187,7 +194,7 @@ class ConnectorEndpoint:
             return True
 
         except Exception as e:
-            logger.error(f"Failed to delete endpoint: {str(e)}")
+            logger.error(CONNECTOR_ENDPOINT_DELETE_ERROR.format(message=str(e)))
             raise e
 
     @staticmethod
@@ -204,6 +211,9 @@ class ConnectorEndpoint:
         Returns:
             tuple[list[str], list[ConnectorEndpointArguments]]: A tuple containing a list of endpoint IDs and a list of
             ConnectorEndpointArguments objects with endpoint details.
+
+        Raises:
+            Exception: If the process of fetching available items encounters an error.
         """
         try:
             retn_eps = []
@@ -223,5 +233,7 @@ class ConnectorEndpoint:
             return retn_eps_ids, retn_eps
 
         except Exception as e:
-            logger.error(f"Failed to get available endpoints: {str(e)}")
+            logger.error(
+                CONNECTOR_ENDPOINT_GET_AVAILABLE_ITEMS_ERROR.format(message=str(e))
+            )
             raise e

moonshot/src/connectors_endpoints/connector_endpoint_arguments.py

@@ -22,6 +22,8 @@ class ConnectorEndpointArguments(BaseModel):
         gt=0
     )  # max_concurrency (int): The number of concurrent api calls
 
+    model: str  # model (str): The model identifier for the LLM connector.
+
     params: dict  # params (dict): A dictionary that contains connection specified parameters
 
     # created_date (str): The date and time the endpoint was created in isoformat without 'T'.
@@ -34,7 +36,7 @@ class ConnectorEndpointArguments(BaseModel):
 
         This method takes all the attributes of the ConnectorEndpointArguments instance and constructs a dictionary
         with attribute names as keys and their corresponding values. This includes the id, name, connector_type, uri,
-        token, max_calls_per_second, max_concurrency, params, and created_date. This dictionary can be used for
+        token, max_calls_per_second, max_concurrency, model, params, and created_date. This dictionary can be used for
         serialization purposes, such as storing the endpoint information in a JSON file or sending it over a network.
 
         Returns:
@@ -49,6 +51,7 @@ class ConnectorEndpointArguments(BaseModel):
             "token": self.token,
             "max_calls_per_second": self.max_calls_per_second,
             "max_concurrency": self.max_concurrency,
+            "model": self.model,
             "params": self.params,
             "created_date": self.created_date,
         }