Allow chain/chain_id to take str/int and convert to int (#12)

* Ensure that chain_id is always an int

* Allow str types, and call normalize function

* allow str | int for any chain param, and handle with normalization

* move normalize_chain_id to common/utils.py

* remove normalize_chain_ids and ensure types are correct

* allow str chain_id/chain and use noramlize function to convert to int

* feat: add regex to normalize_chain_id to extract digits

* feat: add initial tests for utils.py

* docs: add information on running pytest

* build: add pytest to deps

* fix: rebase to main

* docs: adjust uv commands

Author: cjber

python/thirdweb-ai/README.md

@@ -146,10 +146,46 @@ def adapt_to_my_framework(tools: list[Tool]):
 
 This project is licensed under the Apache License 2.0 - see the LICENSE file for details.
 
-## Contributing
+## Development and Testing
 
-Contributions are welcome! Please feel free to submit a Pull Request.
+### Setting up development environment
 
-## Support
+```bash
+# Clone the repository
+git clone https://github.com/thirdweb-dev/ai.git
+cd ai/python/thirdweb-ai
+
+# Install dependencies with UV
+uv sync
+```
+
+### Running tests
+
+We use pytest for testing. You can run the tests with:
+
+```bash
+# Run all tests
+uv run pytest
+
+# Run tests with verbose output
+uv run pytest -v
 
-If you need help or have questions, please contact [[email protected]](mailto:[email protected]).
+# Run specific test file
+uv run pytest tests/common/test_utils.py
+
+# Run tests with coverage report
+uv run pytest --cov=thirdweb_ai
+
+# Run tests and generate HTML coverage report
+uv run pytest --cov=thirdweb_ai --cov-report=html
+```
+
+### Linting and Type Checking
+
+```bash
+# Run the ruff linter
+uv run ruff check .
+
+# Run type checking with pyright
+uv run pyright
+```

python/thirdweb-ai/src/thirdweb_ai/common/utils.py

@@ -0,0 +1,31 @@
+import re
+
+
+def extract_digits(value: int | str) -> int:
+    value_str = str(value).strip("\"'")
+    digit_match = re.search(r"\d+", value_str)
+
+    if not digit_match:
+        raise ValueError(f"Chain ID '{value}' does not contain any digits")
+
+    extracted_digits = digit_match.group()
+
+    if not extracted_digits.isdigit():
+        raise ValueError(
+            f"Extracted value '{extracted_digits}' is not a valid digit string"
+        )
+    return int(extracted_digits)
+
+
+def normalize_chain_id(
+    chain_id: int | str | list[int | str] | None,
+) -> int | list[int] | None:
+    """Normalize chain IDs to integers."""
+
+    if chain_id is None:
+        return None
+
+    if isinstance(chain_id, list):
+        return [extract_digits(c) for c in chain_id]
+
+    return extract_digits(chain_id)

python/thirdweb-ai/src/thirdweb_ai/services/engine.py

@@ -1,5 +1,6 @@
 from typing import Annotated, Any
 
+from thirdweb_ai.common.utils import normalize_chain_id
 from thirdweb_ai.services.service import Service
 from thirdweb_ai.tools.tool import tool
 
@@ -9,15 +10,15 @@ def __init__(
         self,
         engine_url: str,
         engine_auth_jwt: str,
-        chain_id: int | None = None,
+        chain_id: int | str | None = None,
         backend_wallet_address: str | None = None,
         secret_key: str = "",
     ):
         super().__init__(base_url=engine_url, secret_key=secret_key)
         self.engine_url = engine_url
         self.engine_auth_jwt = engine_auth_jwt
         self.backend_wallet_address = backend_wallet_address
-        self.chain_id = str(chain_id) if chain_id else None
+        self.chain_id = normalize_chain_id(chain_id)
 
     def _make_headers(self):
         headers = super()._make_headers()
@@ -75,7 +76,7 @@ def get_all_backend_wallet(
     def get_wallet_balance(
         self,
         chain_id: Annotated[
-            str,
+            str | int,
             "The numeric blockchain network ID to query (e.g., '1' for Ethereum mainnet, '137' for Polygon). If not provided, uses the default chain ID configured in the Engine instance.",
         ],
         backend_wallet_address: Annotated[
@@ -84,9 +85,9 @@ def get_wallet_balance(
         ] = None,
     ) -> dict[str, Any]:
         """Get wallet balance for native or ERC20 tokens."""
-        chain_id = chain_id or self.chain_id
+        normalized_chain = normalize_chain_id(chain_id) or self.chain_id
         backend_wallet_address = backend_wallet_address or self.backend_wallet_address
-        return self._get(f"backend-wallet/{chain_id}/{backend_wallet_address}/get-balance")
+        return self._get(f"backend-wallet/{normalized_chain}/{backend_wallet_address}/get-balance")
 
     @tool(
         description="Send an on-chain transaction. This powerful function can transfer native currency (ETH, MATIC), ERC20 tokens, or execute any arbitrary contract interaction. The transaction is signed and broadcast to the blockchain automatically."
@@ -122,10 +123,10 @@ def send_transaction(
             "data": data or "0x",
         }
 
-        chain_id = chain_id or self.chain_id
+        normalized_chain = normalize_chain_id(chain_id) or self.chain_id
         backend_wallet_address = backend_wallet_address or self.backend_wallet_address
         return self._post(
-            f"backend-wallet/{chain_id}/send-transaction",
+            f"backend-wallet/{normalized_chain}/send-transaction",
             payload,
             headers={"X-Backend-Wallet-Address": backend_wallet_address},
         )
@@ -170,8 +171,8 @@ def read_contract(
             "functionName": function_name,
             "args": function_args or [],
         }
-        chain_id = chain_id or self.chain_id
-        return self._get(f"contract/{chain_id!s}/{contract_address}/read", payload)
+        normalized_chain = normalize_chain_id(chain_id) or self.chain_id
+        return self._get(f"contract/{normalized_chain}/{contract_address}/read", payload)
 
     @tool(
         description="Execute a state-changing function on a smart contract by sending a transaction. This allows you to modify on-chain data, such as transferring tokens, minting NFTs, or updating contract configuration. The transaction is automatically signed by your backend wallet and submitted to the blockchain."
@@ -208,9 +209,9 @@ def write_contract(
         if value and value != "0":
             payload["txOverrides"] = {"value": value}
 
-        chain_id = chain_id or self.chain_id
+        normalized_chain = normalize_chain_id(chain_id) or self.chain_id
         return self._post(
-            f"contract/{chain_id!s}/{contract_address}/write",
+            f"contract/{normalized_chain}/{contract_address}/write",
             payload,
             headers={"X-Backend-Wallet-Address": self.backend_wallet_address},
         )