Added ultra endpoints implementation

Author: cmoutafidis

jup_ag_sdk/clients/jupiter_client.py

@@ -4,8 +4,7 @@
 
 import base58
 import httpx
-from solana.rpc.api import Client
-from solders.solders import Keypair, SendTransactionResp, VersionedTransaction
+from solders.solders import Keypair, VersionedTransaction
 
 
 class JupiterClient:
@@ -17,12 +16,10 @@ class JupiterClient:
     def __init__(
         self,
         api_key: Optional[str],
-        rpc_url: Optional[str],
         private_key_env_var: str,
         timeout: int,
     ):
         self.api_key = api_key
-        self.rpc = Client(rpc_url) if rpc_url else None
         self.base_url = (
             "https://api.jup.ag" if api_key else "https://lite-api.jup.ag"
         )
@@ -56,13 +53,6 @@ def _get_public_key(self) -> str:
         )
         return str(wallet.pubkey())
 
-    def _send_transaction(
-        self, transaction: VersionedTransaction
-    ) -> SendTransactionResp:
-        if not self.rpc:
-            raise ValueError("Client was initialized without RPC URL.")
-        return self.rpc.send_transaction(transaction)
-
     def _sign_base64_transaction(
         self, transaction_base64: str
     ) -> VersionedTransaction:

jup_ag_sdk/clients/swap_api_client.py

@@ -11,7 +11,7 @@
 
 class UltraApiClient(JupiterClient):
     """
-    A client for interacting with the Jupiter Swap API.
+    A client for interacting with the Jupiter Ultra API.
     Inherits from JupiterClient.
     """
 
@@ -23,7 +23,6 @@ def __init__(
     ):
         super().__init__(
             api_key=api_key,
-            rpc_url=None,
             private_key_env_var=private_key_env_var,
             timeout=timeout,
         )
@@ -93,3 +92,41 @@ def order_and_execute(self, request: UltraOrderRequest) -> Dict[str, Any]:
         )
 
         return self.execute(execute_request)
+
+    def balances(self, address: str) -> Dict[str, Any]:
+        """
+        Get token balances of an account from the Jupiter Ultra API.
+
+        Args:
+            address (str): The public key of the account to get balances for.
+
+        Returns:
+            dict: The dict api response.
+        """
+        url = f"{self.base_url}/ultra/v1/balances/{address}"
+        response = self.client.get(url, headers=self._get_headers())
+        response.raise_for_status()
+
+        return response.json()  # type: ignore
+
+    def shield(self, mints: list[str]) -> Dict[str, Any]:
+        """
+        Get token information and warnings for specific mints
+        from the Jupiter Ultra API.
+
+        Args:
+            mints (list[str]): List of token mint addresses
+            to get information for.
+
+        Returns:
+            dict: The dict api response with warnings information.
+        """
+        params = {"mints": ",".join(mints)}
+
+        url = f"{self.base_url}/ultra/v1/shield"
+        response = self.client.get(
+            url, params=params, headers=self._get_headers()
+        )
+        response.raise_for_status()
+
+        return response.json()  # type: ignore

jup_ag_sdk/clients/ultra_api_client.py

@@ -9,6 +9,8 @@ class UltraOrderRequest(BaseModel):
     output_mint: str
     amount: int
     taker: Optional[str] = None
+    referral_account: Optional[str] = None
+    referral_fee: Optional[int] = None
 
     def to_dict(self) -> Dict[str, Any]:
         params = self.model_dump(exclude_none=True)

jup_ag_sdk/models/swap_api/__init__.py

@@ -0,0 +1,33 @@
+from dotenv import load_dotenv
+
+from jup_ag_sdk.clients.ultra_api_client import UltraApiClient
+
+
+def test_ultra_get_balances() -> None:
+    """
+    Test the UltraApiClient balances method.
+    """
+    load_dotenv()
+    client = UltraApiClient()
+
+    address = client._get_public_key()
+
+    try:
+        balances_response = client.balances(str(address))
+        assert (
+            "SOL" in balances_response
+        ), "Response does not contain 'SOL' key."
+
+        print()
+        print("Balances API Response:")
+        for token, details in balances_response.items():
+            print(f"Token: {token}")
+            print(f"  - Amount: {details['amount']}")
+            print(f"  - UI Amount: {details['uiAmount']}")
+            print(f"  - Slot: {details['slot']}")
+            print(f"  - Is Frozen: {details['isFrozen']}")
+
+    except Exception as e:
+        print("Error occurred while fetching balances:", str(e))
+    finally:
+        client.close()

jup_ag_sdk/models/swap_api/quote_request_model.py

@@ -21,15 +21,16 @@ def test_ultra_get_order_and_execute() -> None:
     )
 
     try:
-        rpc_response = client.order_and_execute(order_request)
-        signature = str(rpc_response["signature"])
+        client_response = client.order_and_execute(order_request)
+        signature = str(client_response["signature"])
         assert (
             signature is not None
         ), "Transaction signature is missing or invalid."
-        print(
-            f"Transaction sent successfully!"
-            f"View transaction on Solscan: https://solscan.io/tx/{signature}"
-        )
+
+        print()
+        print("Order and Execute API Response:")
+        print(f"  - Transaction Signature: {signature}")
+        print(f"  - View on Solscan: https://solscan.io/tx/{signature}")
 
     except Exception as e:
         print("Error occurred while processing the order:", str(e))