hackagent 0.3.1__py3-none-any.whl

hackagent/client.py

@@ -0,0 +1,286 @@
+import ssl
+from typing import Any, Optional, Union
+
+import httpx
+from attrs import define, evolve, field
+
+
+@define
+class Client:
+    """A class for keeping track of data related to the API
+
+    The following are accepted as keyword arguments and will be used to construct httpx Clients internally:
+
+        ``base_url``: The base URL for the API, all requests are made to a relative path to this URL
+
+        ``cookies``: A dictionary of cookies to be sent with every request
+
+        ``headers``: A dictionary of headers to be sent with every request
+
+        ``timeout``: The maximum amount of a time a request can take. API functions will raise
+        httpx.TimeoutException if this is exceeded.
+
+        ``verify_ssl``: Whether or not to verify the SSL certificate of the API server. This should be True in production,
+        but can be set to False for testing purposes.
+
+        ``follow_redirects``: Whether or not to follow redirects. Default value is False.
+
+        ``httpx_args``: A dictionary of additional arguments to be passed to the ``httpx.Client`` and ``httpx.AsyncClient`` constructor.
+
+
+    Attributes:
+        raise_on_unexpected_status: Whether or not to raise an errors.UnexpectedStatus if the API returns a
+            status code that was not documented in the source OpenAPI document. Can also be provided as a keyword
+            argument to the constructor.
+    """
+
+    raise_on_unexpected_status: bool = field(default=False, kw_only=True)
+    _base_url: str = field(alias="base_url")
+    _cookies: dict[str, str] = field(factory=dict, kw_only=True, alias="cookies")
+    _headers: dict[str, str] = field(factory=dict, kw_only=True, alias="headers")
+    _timeout: Optional[httpx.Timeout] = field(
+        default=None, kw_only=True, alias="timeout"
+    )
+    _verify_ssl: Union[str, bool, ssl.SSLContext] = field(
+        default=True, kw_only=True, alias="verify_ssl"
+    )
+    _follow_redirects: bool = field(
+        default=False, kw_only=True, alias="follow_redirects"
+    )
+    _httpx_args: dict[str, Any] = field(factory=dict, kw_only=True, alias="httpx_args")
+    _client: Optional[httpx.Client] = field(default=None, init=False)
+    _async_client: Optional[httpx.AsyncClient] = field(default=None, init=False)
+
+    def with_headers(self, headers: dict[str, str]) -> "Client":
+        """Get a new client matching this one with additional headers"""
+        if self._client is not None:
+            self._client.headers.update(headers)
+        if self._async_client is not None:
+            self._async_client.headers.update(headers)
+        return evolve(self, headers={**self._headers, **headers})
+
+    def with_cookies(self, cookies: dict[str, str]) -> "Client":
+        """Get a new client matching this one with additional cookies"""
+        if self._client is not None:
+            self._client.cookies.update(cookies)
+        if self._async_client is not None:
+            self._async_client.cookies.update(cookies)
+        return evolve(self, cookies={**self._cookies, **cookies})
+
+    def with_timeout(self, timeout: httpx.Timeout) -> "Client":
+        """Get a new client matching this one with a new timeout (in seconds)"""
+        if self._client is not None:
+            self._client.timeout = timeout
+        if self._async_client is not None:
+            self._async_client.timeout = timeout
+        return evolve(self, timeout=timeout)
+
+    def set_httpx_client(self, client: httpx.Client) -> "Client":
+        """Manually set the underlying httpx.Client
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._client = client
+        return self
+
+    def get_httpx_client(self) -> httpx.Client:
+        """Get the underlying httpx.Client, constructing a new one if not previously set"""
+        if self._client is None:
+            self._client = httpx.Client(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._client
+
+    def __enter__(self) -> "Client":
+        """Enter a context manager for self.client—you cannot enter twice (see httpx docs)"""
+        self.get_httpx_client().__enter__()
+        return self
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for internal httpx.Client (see httpx docs)"""
+        self.get_httpx_client().__exit__(*args, **kwargs)
+
+    def set_async_httpx_client(self, async_client: httpx.AsyncClient) -> "Client":
+        """Manually the underlying httpx.AsyncClient
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._async_client = async_client
+        return self
+
+    def get_async_httpx_client(self) -> httpx.AsyncClient:
+        """Get the underlying httpx.AsyncClient, constructing a new one if not previously set"""
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._async_client
+
+    async def __aenter__(self) -> "Client":
+        """Enter a context manager for underlying httpx.AsyncClient—you cannot enter twice (see httpx docs)"""
+        await self.get_async_httpx_client().__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for underlying httpx.AsyncClient (see httpx docs)"""
+        await self.get_async_httpx_client().__aexit__(*args, **kwargs)
+
+
+@define
+class AuthenticatedClient:
+    """A Client which has been authenticated for use on secured endpoints
+
+    The following are accepted as keyword arguments and will be used to construct httpx Clients internally:
+
+        ``base_url``: The base URL for the API, all requests are made to a relative path to this URL
+
+        ``cookies``: A dictionary of cookies to be sent with every request
+
+        ``headers``: A dictionary of headers to be sent with every request
+
+        ``timeout``: The maximum amount of a time a request can take. API functions will raise
+        httpx.TimeoutException if this is exceeded.
+
+        ``verify_ssl``: Whether or not to verify the SSL certificate of the API server. This should be True in production,
+        but can be set to False for testing purposes.
+
+        ``follow_redirects``: Whether or not to follow redirects. Default value is False.
+
+        ``httpx_args``: A dictionary of additional arguments to be passed to the ``httpx.Client`` and ``httpx.AsyncClient`` constructor.
+
+
+    Attributes:
+        raise_on_unexpected_status: Whether or not to raise an errors.UnexpectedStatus if the API returns a
+            status code that was not documented in the source OpenAPI document. Can also be provided as a keyword
+            argument to the constructor.
+        token: The token to use for authentication
+        prefix: The prefix to use for the Authorization header
+        auth_header_name: The name of the Authorization header
+    """
+
+    raise_on_unexpected_status: bool = field(default=False, kw_only=True)
+    _base_url: str = field(alias="base_url")
+    _cookies: dict[str, str] = field(factory=dict, kw_only=True, alias="cookies")
+    _headers: dict[str, str] = field(factory=dict, kw_only=True, alias="headers")
+    _timeout: Optional[httpx.Timeout] = field(
+        default=None, kw_only=True, alias="timeout"
+    )
+    _verify_ssl: Union[str, bool, ssl.SSLContext] = field(
+        default=True, kw_only=True, alias="verify_ssl"
+    )
+    _follow_redirects: bool = field(
+        default=False, kw_only=True, alias="follow_redirects"
+    )
+    _httpx_args: dict[str, Any] = field(factory=dict, kw_only=True, alias="httpx_args")
+    _client: Optional[httpx.Client] = field(default=None, init=False)
+    _async_client: Optional[httpx.AsyncClient] = field(default=None, init=False)
+
+    token: str
+    prefix: str = "Bearer"
+    auth_header_name: str = "Authorization"
+
+    def with_headers(self, headers: dict[str, str]) -> "AuthenticatedClient":
+        """Get a new client matching this one with additional headers"""
+        if self._client is not None:
+            self._client.headers.update(headers)
+        if self._async_client is not None:
+            self._async_client.headers.update(headers)
+        return evolve(self, headers={**self._headers, **headers})
+
+    def with_cookies(self, cookies: dict[str, str]) -> "AuthenticatedClient":
+        """Get a new client matching this one with additional cookies"""
+        if self._client is not None:
+            self._client.cookies.update(cookies)
+        if self._async_client is not None:
+            self._async_client.cookies.update(cookies)
+        return evolve(self, cookies={**self._cookies, **cookies})
+
+    def with_timeout(self, timeout: httpx.Timeout) -> "AuthenticatedClient":
+        """Get a new client matching this one with a new timeout (in seconds)"""
+        if self._client is not None:
+            self._client.timeout = timeout
+        if self._async_client is not None:
+            self._async_client.timeout = timeout
+        return evolve(self, timeout=timeout)
+
+    def set_httpx_client(self, client: httpx.Client) -> "AuthenticatedClient":
+        """Manually set the underlying httpx.Client
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._client = client
+        return self
+
+    def get_httpx_client(self) -> httpx.Client:
+        """Get the underlying httpx.Client, constructing a new one if not previously set"""
+        if self._client is None:
+            self._headers[self.auth_header_name] = (
+                f"{self.prefix} {self.token}" if self.prefix else self.token
+            )
+            self._client = httpx.Client(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._client
+
+    def __enter__(self) -> "AuthenticatedClient":
+        """Enter a context manager for self.client—you cannot enter twice (see httpx docs)"""
+        self.get_httpx_client().__enter__()
+        return self
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for internal httpx.Client (see httpx docs)"""
+        self.get_httpx_client().__exit__(*args, **kwargs)
+
+    def set_async_httpx_client(
+        self, async_client: httpx.AsyncClient
+    ) -> "AuthenticatedClient":
+        """Manually the underlying httpx.AsyncClient
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._async_client = async_client
+        return self
+
+    def get_async_httpx_client(self) -> httpx.AsyncClient:
+        """Get the underlying httpx.AsyncClient, constructing a new one if not previously set"""
+        if self._async_client is None:
+            self._headers[self.auth_header_name] = (
+                f"{self.prefix} {self.token}" if self.prefix else self.token
+            )
+            self._async_client = httpx.AsyncClient(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._async_client
+
+    async def __aenter__(self) -> "AuthenticatedClient":
+        """Enter a context manager for underlying httpx.AsyncClient—you cannot enter twice (see httpx docs)"""
+        await self.get_async_httpx_client().__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for underlying httpx.AsyncClient (see httpx docs)"""
+        await self.get_async_httpx_client().__aexit__(*args, **kwargs)

hackagent/errors.py

@@ -0,0 +1,37 @@
+"""Contains shared errors types that can be raised from API functions"""
+
+
+class UnexpectedStatus(Exception):
+    """Raised by api functions when the response status an undocumented status and Client.raise_on_unexpected_status is True"""
+
+    def __init__(self, status_code: int, content: bytes):
+        self.status_code = status_code
+        self.content = content
+
+        super().__init__(
+            f"Unexpected status code: {status_code}\n\nResponse content:\n{content.decode(errors='ignore')}"
+        )
+
+
+class HackAgentError(Exception):
+    """Base exception class for HackAgent errors"""
+
+    pass
+
+
+class ApiError(HackAgentError):
+    """Raised when an API call fails"""
+
+    def __init__(
+        self, message: str, status_code: int | None = None, response: dict | None = None
+    ):
+        self.message = message
+        self.status_code = status_code
+        self.response = response
+        super().__init__(message)
+
+
+# Alias for backward compatibility with tests
+UnexpectedStatusError = UnexpectedStatus
+
+__all__ = ["UnexpectedStatus", "UnexpectedStatusError", "HackAgentError", "ApiError"]

hackagent/logger.py

@@ -0,0 +1,83 @@
+# Copyright 2025 - AI4I. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import os
+
+from rich.logging import RichHandler
+
+_rich_handler_configured_for_package = False
+
+
+def setup_package_logging(
+    logger_name: str = "hackagent", default_level_str: str = "WARNING"
+):
+    """Configures RichHandler for the specified logger if not already set."""
+    global _rich_handler_configured_for_package
+
+    package_logger = logging.getLogger(logger_name)
+
+    if logger_name == "hackagent" and _rich_handler_configured_for_package:
+        return
+
+    has_console_handler = any(
+        isinstance(h, logging.StreamHandler) for h in package_logger.handlers
+    )
+
+    if not has_console_handler:
+        log_level_env = os.getenv(
+            f"{logger_name.upper()}_LOG_LEVEL", default_level_str
+        ).upper()
+        level = getattr(logging, log_level_env, logging.WARNING)
+        package_logger.setLevel(level)
+
+        rich_handler = RichHandler(
+            show_time=True,
+            show_level=True,
+            show_path=False,
+            markup=True,
+            rich_tracebacks=True,
+            tracebacks_show_locals=False,
+        )
+        package_logger.addHandler(rich_handler)
+        package_logger.propagate = False  # Avoid duplicate logs with root logger
+
+        if logger_name == "hackagent":
+            _rich_handler_configured_for_package = True
+            # Set default levels for common noisy libraries
+            logging.getLogger("httpx").setLevel(logging.WARNING)
+            logging.getLogger("litellm").setLevel(logging.WARNING)
+            # Add other libraries here if needed, e.g.:
+            # logging.getLogger("another_library").setLevel(logging.WARNING)
+
+        # package_logger.debug(f"RichHandler configured for '{logger_name}' logger at level {level}.")
+
+    elif any(isinstance(h, RichHandler) for h in package_logger.handlers):
+        if logger_name == "hackagent":
+            _rich_handler_configured_for_package = True
+
+    return package_logger
+
+
+def get_logger(name: str) -> logging.Logger:
+    """
+    Retrieves a logger instance.
+    If the logger is 'hackagent' or starts with 'hackagent.',
+    it ensures the package logging is set up.
+    """
+    if name == "hackagent" or name.startswith("hackagent."):
+        # Ensure base "hackagent" logger is configured first
+        setup_package_logging(logger_name="hackagent")
+    return logging.getLogger(name)

hackagent/models/__init__.py

@@ -0,0 +1,109 @@
+"""Contains all the data models used in inputs/outputs"""
+
+from .agent import Agent
+from .agent_request import AgentRequest
+from .api_token_log import APITokenLog
+from .attack import Attack
+from .attack_request import AttackRequest
+from .checkout_session_request_request import CheckoutSessionRequestRequest
+from .checkout_session_response import CheckoutSessionResponse
+from .choice import Choice
+from .choice_message import ChoiceMessage
+from .evaluation_status_enum import EvaluationStatusEnum
+from .generate_error_response import GenerateErrorResponse
+from .generate_request_request import GenerateRequestRequest
+from .generate_success_response import GenerateSuccessResponse
+from .generic_error_response import GenericErrorResponse
+from .message_request import MessageRequest
+from .organization import Organization
+from .organization_minimal import OrganizationMinimal
+from .organization_request import OrganizationRequest
+from .paginated_agent_list import PaginatedAgentList
+from .paginated_api_token_log_list import PaginatedAPITokenLogList
+from .paginated_attack_list import PaginatedAttackList
+from .paginated_organization_list import PaginatedOrganizationList
+from .paginated_prompt_list import PaginatedPromptList
+from .paginated_result_list import PaginatedResultList
+from .paginated_run_list import PaginatedRunList
+from .paginated_user_api_key_list import PaginatedUserAPIKeyList
+from .paginated_user_profile_list import PaginatedUserProfileList
+from .patched_agent_request import PatchedAgentRequest
+from .patched_attack_request import PatchedAttackRequest
+from .patched_organization_request import PatchedOrganizationRequest
+from .patched_prompt_request import PatchedPromptRequest
+from .patched_result_request import PatchedResultRequest
+from .patched_run_request import PatchedRunRequest
+from .patched_user_profile_request import PatchedUserProfileRequest
+from .prompt import Prompt
+from .prompt_request import PromptRequest
+from .result import Result
+from .result_list_evaluation_status import ResultListEvaluationStatus
+from .result_request import ResultRequest
+from .run import Run
+from .run_list_status import RunListStatus
+from .run_request import RunRequest
+from .status_enum import StatusEnum
+from .step_type_enum import StepTypeEnum
+from .trace import Trace
+from .trace_request import TraceRequest
+from .usage import Usage
+from .user_api_key import UserAPIKey
+from .user_api_key_request import UserAPIKeyRequest
+from .user_profile import UserProfile
+from .user_profile_minimal import UserProfileMinimal
+from .user_profile_request import UserProfileRequest
+
+__all__ = (
+    "Agent",
+    "AgentRequest",
+    "APITokenLog",
+    "Attack",
+    "AttackRequest",
+    "CheckoutSessionRequestRequest",
+    "CheckoutSessionResponse",
+    "Choice",
+    "ChoiceMessage",
+    "EvaluationStatusEnum",
+    "GenerateErrorResponse",
+    "GenerateRequestRequest",
+    "GenerateSuccessResponse",
+    "GenericErrorResponse",
+    "MessageRequest",
+    "Organization",
+    "OrganizationMinimal",
+    "OrganizationRequest",
+    "PaginatedAgentList",
+    "PaginatedAPITokenLogList",
+    "PaginatedAttackList",
+    "PaginatedOrganizationList",
+    "PaginatedPromptList",
+    "PaginatedResultList",
+    "PaginatedRunList",
+    "PaginatedUserAPIKeyList",
+    "PaginatedUserProfileList",
+    "PatchedAgentRequest",
+    "PatchedAttackRequest",
+    "PatchedOrganizationRequest",
+    "PatchedPromptRequest",
+    "PatchedResultRequest",
+    "PatchedRunRequest",
+    "PatchedUserProfileRequest",
+    "Prompt",
+    "PromptRequest",
+    "Result",
+    "ResultListEvaluationStatus",
+    "ResultRequest",
+    "Run",
+    "RunListStatus",
+    "RunRequest",
+    "StatusEnum",
+    "StepTypeEnum",
+    "Trace",
+    "TraceRequest",
+    "Usage",
+    "UserAPIKey",
+    "UserAPIKeyRequest",
+    "UserProfile",
+    "UserProfileMinimal",
+    "UserProfileRequest",
+)