Merge pull request #42 from galadriel-ai/dalle-updates

DalleNft.sol updates:

contracts/contracts/ChatOracle.sol

@@ -9,13 +9,13 @@ interface IOracleTypes {
     struct OpenAiRequest {
         // "gpt-4-turbo-preview" or "gpt-3.5-turbo-1106"
         string model;
-        // int -20 - 20? Mapped to float -2.0 - 2.0? If bigger than 20 then null?
+        // int -20 - 20, Mapped to float -2.0 - 2.0. If bigger than 20 then null
         int8 frequencyPenalty;
         // JSON string or empty string
         string logitBias;
         // 0 for null
         uint32 maxTokens;
-        // int -20 - 20? Mapped to float -2.0 - 2.0? If bigger than 20 then null?
+        // int -20 - 20, Mapped to float -2.0 - 2.0. If bigger than 20 then null
         int8 presencePenalty;
         // JSON string or empty string
         string responseFormat;
@@ -56,13 +56,13 @@ interface IOracleTypes {
     struct GroqRequest {
         // "llama2-70b-4096", "mixtral-8x7b-32768" or "gemma-7b-it"
         string model;
-        // int -20 - 20? Mapped to float -2.0 - 2.0? If bigger than 20 then null?
+        // int -20 - 20, Mapped to float -2.0 - 2.0. If bigger than 20 then null
         int8 frequencyPenalty;
         // JSON string or empty string
         string logitBias;
         // 0 for null
         uint32 maxTokens;
-        // int -20 - 20? Mapped to float -2.0 - 2.0? If bigger than 20 then null?
+        // int -20 - 20, Mapped to float -2.0 - 2.0. If bigger than 20 then null
         int8 presencePenalty;
         // JSON string or empty string
         string responseFormat;

contracts/contracts/DalleNft.sol

@@ -4,7 +4,7 @@ pragma solidity ^0.8.20;
 // Uncomment this line to use console.log
 // import "hardhat/console.sol";
 import {ERC721URIStorage, ERC721} from "@openzeppelin/contracts/token/ERC721/extensions/ERC721URIStorage.sol";
-
+import {ERC721Enumerable} from "@openzeppelin/contracts/token/ERC721/extensions/ERC721Enumerable.sol";
 
 interface IOracle {
     function createFunctionCall(
@@ -14,7 +14,7 @@ interface IOracle {
     ) external returns (uint i);
 }
 
-contract DalleNft is ERC721URIStorage {
+contract DalleNft is ERC721, ERC721Enumerable, ERC721URIStorage {
     uint256 private _nextTokenId;
 
     struct MintInput {
@@ -24,7 +24,6 @@ contract DalleNft is ERC721URIStorage {
     }
 
     mapping(uint => MintInput) public mintInputs;
-    uint private mintsCount;
 
     event MintInputCreated(address indexed owner, uint indexed chatId);
 
@@ -66,14 +65,14 @@ contract DalleNft is ERC721URIStorage {
     }
 
     function initializeMint(string memory message) public returns (uint i) {
-        MintInput storage mintInput = mintInputs[mintsCount];
+        uint256 currentId = _nextTokenId++;
+        MintInput storage mintInput = mintInputs[currentId];
 
         mintInput.owner = msg.sender;
         mintInput.prompt = message;
+
         mintInput.isMinted = false;
 
-        uint currentId = mintsCount;
-        mintsCount = currentId + 1;
 
         string memory fullPrompt = prompt;
         fullPrompt = string.concat(fullPrompt, message);
@@ -98,9 +97,8 @@ contract DalleNft is ERC721URIStorage {
 
         mintInput.isMinted = true;
 
-        uint256 tokenId = _nextTokenId++;
-        _mint(mintInput.owner, tokenId);
-        _setTokenURI(tokenId, response);
+        _mint(mintInput.owner, runId);
+        _setTokenURI(runId, response);
     }
 
     function getMessageHistoryContents(uint chatId) public view returns (string[] memory) {
@@ -117,4 +115,37 @@ contract DalleNft is ERC721URIStorage {
         rolesArray[0] = "user";
         return rolesArray;
     }
+
+    function _update(address to, uint256 tokenId, address auth)
+        internal
+        override(ERC721, ERC721Enumerable)
+        returns (address)
+    {
+        return super._update(to, tokenId, auth);
+    }
+
+    function _increaseBalance(address account, uint128 value)
+        internal
+        override(ERC721, ERC721Enumerable)
+    {
+        super._increaseBalance(account, value);
+    }
+
+    function tokenURI(uint256 tokenId)
+        public
+        view
+        override(ERC721, ERC721URIStorage)
+        returns (string memory)
+    {
+        return super.tokenURI(tokenId);
+    }
+
+    function supportsInterface(bytes4 interfaceId)
+        public
+        view
+        override(ERC721, ERC721Enumerable, ERC721URIStorage)
+        returns (bool)
+    {
+        return super.supportsInterface(interfaceId);
+    }
 }

contracts/test/DalleNft.ts

@@ -54,6 +54,42 @@ describe("DalleNft", function () {
       const tokenUri = await dalleNft.tokenURI(0)
       expect(tokenUri).to.equal("ipfs://CID")
     });
+    it("TotalSupply works", async () => {
+      const {
+        dalleNft,
+        oracle,
+        allSigners
+      } = await loadFixture(deploy);
+      const oracleAccount = allSigners[6];
+      await dalleNft.setOracleAddress(oracle.target);
+      await oracle.updateWhitelist(oracleAccount, true);
+
+      await dalleNft.initializeMint("funky gorilla");
+
+      await oracle.connect(oracleAccount).addFunctionResponse(0, 0, "ipfs://CID", "");
+
+      const totalSupply = await dalleNft.totalSupply()
+      expect(totalSupply).to.equal(1)
+    });
+    it("Can get user NFT token ID by index", async () => {
+      const {
+        dalleNft,
+        oracle,
+        allSigners
+      } = await loadFixture(deploy);
+      const oracleAccount = allSigners[6];
+      await dalleNft.setOracleAddress(oracle.target);
+      await oracle.updateWhitelist(oracleAccount, true);
+
+      await dalleNft.initializeMint("funky gorilla");
+      await oracle.connect(oracleAccount).addFunctionResponse(0, 0, "ipfs://CID", "");
+
+      await dalleNft.initializeMint("funky gorilla 2");
+      await oracle.connect(oracleAccount).addFunctionResponse(1, 1, "ipfs://CID", "");
+
+      const tokenId = await dalleNft.tokenOfOwnerByIndex(allSigners[0].address, 1)
+      expect(tokenId).to.equal(1)
+    });
   });
   describe("Error handling", function () {
     it("Cannot mint same prompt twice", async () => {