scripts/mech_request_utils.py

# -*- coding: utf-8 -*-
# ------------------------------------------------------------------------------
#
#   Copyright 2024 Valory AG
#
#   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.
#
# ------------------------------------------------------------------------------

"""Script for retrieving mech requests and their delivers."""
import json
import time
import pickle
from random import uniform
from typing import Any, Dict, Tuple
import requests
from gql import Client, gql
from gql.transport.requests import RequestsHTTPTransport
from tools import (
    GET_CONTENTS_BATCH_SIZE,
    IRRELEVANT_TOOLS,
    create_session,
    request,
)
from tqdm import tqdm
from web3_utils import (
    FPMM_QS_CREATOR,
    FPMM_PEARL_CREATOR,
    IPFS_POLL_INTERVAL,
    SUBGRAPH_POLL_INTERVAL,
)
from concurrent.futures import ThreadPoolExecutor, as_completed
from utils import (
    DATA_DIR,
    JSON_DATA_DIR,
    MECH_SUBGRAPH_URL,
    SUBGRAPH_API_KEY,
    IPFS_ADDRESS,
)

NUM_WORKERS = 10
BLOCKS_CHUNK_SIZE = 10000
TEXT_ALIGNMENT = 30
MINIMUM_WRITE_FILE_DELAY_SECONDS = 20
MECH_FROM_BLOCK_RANGE = 50000

last_write_time = 0.0

REQUESTS_QUERY_FILTER = """
query requests_query($sender_not_in: [Bytes!], $id_gt: Bytes, $blockNumber_gte: BigInt, $blockNumber_lte: BigInt) {
  requests(where: {sender_not_in: $sender_not_in, id_gt: $id_gt, blockNumber_gte: $blockNumber_gte, blockNumber_lte: $blockNumber_lte},  orderBy: id, first: 1000) {
    blockNumber
    blockTimestamp
    id
    ipfsHash
    requestId
    sender
    transactionHash
  }
}
"""

DELIVERS_QUERY_NO_FILTER = """
query delivers_query($id_gt: Bytes, $blockNumber_gte: BigInt, $blockNumber_lte: BigInt) {
  delivers(where: {id_gt: $id_gt, blockNumber_gte: $blockNumber_gte, blockNumber_lte: $blockNumber_lte}, orderBy: id, first: 1000) {
    blockNumber
    blockTimestamp
    id
    ipfsHash
    requestId
    sender
    transactionHash
  }
}

"""
DELIVERS_QUERY = """
query delivers_query($requestId: BigInt, $blockNumber_gte: BigInt, $blockNumber_lte: BigInt) {
  delivers(where: {requestId: $requestId, blockNumber_gte: $blockNumber_gte, blockNumber_lte: $blockNumber_lte}, orderBy: blockNumber, first: 1000) {
    blockNumber
    blockTimestamp
    id
    ipfsHash
    requestId
    sender
    transactionHash
  }
}
"""

MISSING_DELIVERS_QUERY = """
query delivers_query($requestId: BigInt, $blockNumber_gte: BigInt, $blockNumber_lte: BigInt) {
  delivers(where: {requestId: $requestId, blockNumber_gte: $blockNumber_gte, blockNumber_lte: $blockNumber_lte}, orderBy: blockNumber, first: 1000) {
    blockNumber
    blockTimestamp
    id
    ipfsHash
    requestId
    sender
    transactionHash
  }
}
"""


def collect_all_mech_requests(from_block: int, to_block: int, filename: str) -> Tuple:

    print(f"Fetching all mech requests from {from_block} to {to_block}")
    mech_requests = {}
    duplicated_reqIds = []
    mech_subgraph_url = MECH_SUBGRAPH_URL.substitute(subgraph_api_key=SUBGRAPH_API_KEY)
    transport = RequestsHTTPTransport(url=mech_subgraph_url)
    client = Client(transport=transport, fetch_schema_from_transport=True)

    id_gt = "0x00"
    nr_errors = 0
    while True:
        variables = {
            "sender_not_in": [FPMM_QS_CREATOR, FPMM_PEARL_CREATOR],
            "id_gt": id_gt,
            "blockNumber_gte": str(from_block),  # str
            "blockNumber_lte": str(to_block),  # str
        }
        try:
            response = fetch_with_retry(client, REQUESTS_QUERY_FILTER, variables)

            items = response.get("requests", [])

            if not items:
                break

            for mech_request in items:
                if mech_request["id"] not in mech_requests:
                    mech_requests[mech_request["id"]] = mech_request
                else:
                    duplicated_reqIds.append(mech_request["id"])
        except Exception as e:
            # counter for errors
            nr_errors += 1
            print(f"Error while getting the response: {e}")

        id_gt = items[-1]["id"]
        time.sleep(SUBGRAPH_POLL_INTERVAL)
        print(f"New execution for id_gt = {id_gt}")
        if len(duplicated_reqIds) > 0:
            print(f"Number of duplicated req Ids = {len(duplicated_reqIds)}")
        save_json_file(mech_requests, filename)

    print(f"Number of requests = {len(mech_requests)}")
    print(f"Number of duplicated req Ids = {len(duplicated_reqIds)}")
    save_json_file(mech_requests, filename)
    return mech_requests, duplicated_reqIds, nr_errors


def fetch_with_retry(client, query, variables, max_retries=5):
    for attempt in range(max_retries):
        try:
            return client.execute(gql(query), variable_values=variables)
        except Exception as e:
            if attempt == max_retries - 1:
                raise e
            wait_time = (2**attempt) + uniform(0, 1)  # exponential backoff with jitter
            time.sleep(wait_time)


def collect_all_mech_delivers(from_block: int, to_block: int, filename: str) -> Tuple:

    print(f"Fetching all mech delivers from {from_block} to {to_block}")

    mech_delivers = {}
    duplicated_requestIds = []
    mech_subgraph_url = MECH_SUBGRAPH_URL.substitute(subgraph_api_key=SUBGRAPH_API_KEY)
    transport = RequestsHTTPTransport(url=mech_subgraph_url)
    client = Client(transport=transport, fetch_schema_from_transport=True)
    to_block = (
        to_block + MECH_FROM_BLOCK_RANGE
    )  # there is a delay between deliver and request
    id_gt = ""
    nr_errors = 0
    while True:
        variables = {
            "id_gt": id_gt,
            "blockNumber_gte": str(from_block),  # str
            "blockNumber_lte": str(to_block),  # str
        }
        try:
            response = fetch_with_retry(client, DELIVERS_QUERY_NO_FILTER, variables)
            items = response.get("delivers", [])

            if not items:
                break

            for mech_deliver in items:
                if mech_deliver["requestId"] not in mech_delivers:
                    mech_delivers[mech_deliver["requestId"]] = [mech_deliver]
                else:
                    duplicated_requestIds.append(mech_deliver["requestId"])
                    # we will handle the duplicated later
        except Exception as e:
            # counter for errors
            nr_errors += 1
            print(f"Error while getting the response: {e}")
            # return None, None

        id_gt = items[-1]["id"]
        time.sleep(SUBGRAPH_POLL_INTERVAL)
        print(f"New execution for id_gt = {id_gt}")
        if len(duplicated_requestIds) > 0:
            print(f"Number of duplicated request id = {len(duplicated_requestIds)}")
        save_json_file(mech_delivers, filename)
    print(f"Number of delivers = {len(mech_delivers)}")
    print(f"Number of duplicated request id = {len(duplicated_requestIds)}")
    save_json_file(mech_delivers, filename)
    return mech_delivers, duplicated_requestIds, nr_errors


def collect_missing_delivers(request_id: int, block_number: int) -> Dict[str, Any]:
    to_block = (
        block_number + MECH_FROM_BLOCK_RANGE
    )  # there is a delay between deliver and request
    print(f"Fetching all missing delivers from {block_number} to {to_block}")
    mech_delivers = {}
    mech_subgraph_url = MECH_SUBGRAPH_URL.substitute(subgraph_api_key=SUBGRAPH_API_KEY)
    transport = RequestsHTTPTransport(url=mech_subgraph_url)
    client = Client(transport=transport, fetch_schema_from_transport=True)

    variables = {
        "requestId": request_id,
        "blockNumber_gte": str(block_number),  # str
        "blockNumber_lte": str(to_block),  # str
    }
    try:
        response = fetch_with_retry(client, MISSING_DELIVERS_QUERY, variables)
        items = response.get("delivers", [])
        # If the user sends requests with the same values (tool, prompt, nonce) it
        # will generate the same requestId. Therefore, multiple items can be retrieved
        # at this point. We assume the most likely deliver to this request is the
        # one with the closest blockNumber among all delivers with the same requestId.
        if items:
            return items[0]
    except Exception as e:
        print(f"Error while getting the response: {e}")
        # TODO count how many mech requests without a deliver do we have

    return mech_delivers


def populate_requests_ipfs_contents(
    session: requests.Session, mech_requests: Dict[str, Any], keys_to_traverse: list
) -> dict:
    updated_dict = {}
    wrong_response_count = 0
    for k in tqdm(
        keys_to_traverse,
        desc="Fetching IPFS contents for requests",
        position=1,
        unit="results",
    ):
        mech_request = mech_requests[k]

        if "ipfsContents" not in mech_request:
            ipfs_hash = mech_request["ipfsHash"]
            url = f"{IPFS_ADDRESS}{ipfs_hash}/metadata.json"
            response = request(session, url)
            if response is None:
                tqdm.write(f"Skipping {mech_request=}. because response was None")
                wrong_response_count += 1
                continue
            try:
                contents = response.json()
                if contents["tool"] in IRRELEVANT_TOOLS:
                    continue
                mech_request["ipfsContents"] = contents
            except requests.exceptions.JSONDecodeError:
                tqdm.write(
                    f"Skipping {mech_request} because of JSONDecodeError when parsing response"
                )
                wrong_response_count += 1
                continue
        updated_dict[k] = mech_request
        time.sleep(IPFS_POLL_INTERVAL)

    return updated_dict, wrong_response_count


def populate_delivers_ipfs_contents(
    session: requests.Session, mech_requests: Dict[str, Any], keys_to_traverse: list
) -> dict:
    """Function to complete the delivers content info from ipfs"""
    updated_dict = {}
    errors = 0
    for k in tqdm(
        keys_to_traverse,
        desc="Fetching IPFS contents for delivers",
        position=1,
        unit="results",
    ):
        mech_request = mech_requests[k]
        if "deliver" not in mech_request or len(mech_request["deliver"]) == 0:
            print(f"Skipping mech request {mech_request} because of no delivers info")
            continue

        deliver = mech_request["deliver"]
        if "ipfsContents" not in deliver:
            ipfs_hash = deliver["ipfsHash"]
            request_id = deliver["requestId"]
            url = f"{IPFS_ADDRESS}{ipfs_hash}/{request_id}"
            response = request(session, url)
            if response is None:
                tqdm.write(f"Skipping {mech_request=}.")
                continue
            try:
                contents = response.json()
                metadata = contents.get("metadata", None)
                if metadata and contents["metadata"]["tool"] in IRRELEVANT_TOOLS:
                    continue
                contents.pop("cost_dict", None)
                deliver["ipfsContents"] = contents
            except requests.exceptions.JSONDecodeError:
                tqdm.write(f"Skipping {mech_request} because of JSONDecodeError")
                continue
            except Exception:
                errors += 1
                tqdm.write(
                    f"Skipping {mech_request} because of error parsing the response"
                )
                continue
        updated_dict[k] = mech_request
        time.sleep(IPFS_POLL_INTERVAL)

    return updated_dict, errors


def write_mech_events_to_file(
    mech_requests: Dict[str, Any],
    filename: str,
    force_write: bool = False,
) -> None:
    global last_write_time  # pylint: disable=global-statement
    now = time.time()

    if len(mech_requests) == 0:
        return

    filename_path = DATA_DIR / filename
    if force_write or (now - last_write_time) >= MINIMUM_WRITE_FILE_DELAY_SECONDS:
        with open(filename_path, "w", encoding="utf-8") as file:
            json.dump(mech_requests, file, indent=2)
        last_write_time = now


def save_json_file(data: Dict[str, Any], filename: str):
    """Function to save the content into a json file"""
    filename_path = JSON_DATA_DIR / filename
    with open(filename_path, "w", encoding="utf-8") as file:
        json.dump(data, file, indent=2)


def merge_json_files(old_file: str, new_file: str):
    # read old file
    with open(JSON_DATA_DIR / old_file, "r") as f:
        old_data = json.load(f)

    # read the new file
    with open(JSON_DATA_DIR / new_file, "r") as f:
        new_data = json.load(f)

    # Merge the two JSON files and remove duplicates
    old_data.update(new_data)

    # Save the merged JSON file
    print(f"{old_file} updated")
    save_json_file(old_data, old_file)


def clean_mech_delivers(requests_filename: str, delivers_filename: str) -> None:
    """Function to remove from the delivers json file the request Ids that are not in the mech requests"""
    # read mech requests
    with open(JSON_DATA_DIR / requests_filename, "r") as file:
        mech_requests = json.load(file)

    list_reqIds = [mech_requests[k].get("requestId") for k in mech_requests.keys()]

    # remove requestIds from delivers that are not in this list
    with open(JSON_DATA_DIR / delivers_filename, "r") as file:
        mech_delivers = json.load(file)

    print(f"original size of the file {len(mech_delivers)}")
    mech_delivers = {
        k: v
        for k, v in tqdm(
            mech_delivers.items(),
            total=len(mech_delivers),
            desc="Filtering delivers dictionary",
        )
        if k in set(list_reqIds)
    }

    print(f"final size of the file {len(mech_delivers)}")
    save_json_file(mech_delivers, delivers_filename)


def get_request_block_numbers(
    mech_requests: Dict[str, Any], target_req_id: int
) -> list:
    block_numbers = []

    for entry in mech_requests.values():
        if entry["requestId"] == target_req_id:
            block_numbers.append(entry["blockNumber"])

    return block_numbers


def update_block_request_map(block_request_id_map: dict) -> None:
    print("Saving block request id map info")
    with open(JSON_DATA_DIR / "block_request_id_map.pickle", "wb") as handle:
        pickle.dump(block_request_id_map, handle, protocol=pickle.HIGHEST_PROTOCOL)


def fix_duplicate_requestIds(requests_filename: str, delivers_filename: str) -> dict:
    print("Fix duplicated request Ids")
    with open(JSON_DATA_DIR / delivers_filename, "r") as file:
        data_delivers = json.load(file)

    with open(JSON_DATA_DIR / requests_filename, "r") as file:
        mech_requests = json.load(file)
    list_request_Ids = list(data_delivers.keys())

    list_duplicated_reqIds = []
    for req_Id in list_request_Ids:
        if len(data_delivers.get(req_Id)) > 1:
            list_duplicated_reqIds.append(req_Id)

    print(len(list_duplicated_reqIds))
    block_request_id_map = {}

    for req_Id in list_duplicated_reqIds:
        # get the list of mech request block numbers for that requestId
        block_nrs = get_request_block_numbers(mech_requests, req_Id)
        # get the list of mech delivers
        mech_delivers_list = data_delivers.get(req_Id)  # list of dictionaries
        if len(block_nrs) > 1:
            print("More than one block number was found")
        for block_nr in block_nrs:
            key = (block_nr, req_Id)
            min_difference_request = min(
                mech_delivers_list,
                key=lambda x: abs(int(x["blockNumber"]) - int(block_nr)),
            )
            block_request_id_map[key] = min_difference_request

    update_block_request_map(block_request_id_map)

    return block_request_id_map


def merge_requests_delivers(
    requests_filename: str, delivers_filename: str, filename: str
) -> None:
    print("Merge request delivers")
    """Function to map requests and delivers"""
    with open(JSON_DATA_DIR / delivers_filename, "r") as file:
        mech_delivers = json.load(file)

    with open(JSON_DATA_DIR / requests_filename, "r") as file:
        mech_requests = json.load(file)

    # read the block map for duplicated requestIds
    with open(JSON_DATA_DIR / "block_request_id_map.pickle", "rb") as handle:
        # key = (block_nr, req_Id) value = delivers dictionary
        block_request_id_map = pickle.load(handle)
    for _, mech_req in tqdm(
        mech_requests.items(),
        desc=f"Merging delivers data into the mech requests",
    ):
        if "deliver" in mech_req:
            continue

        block_number_req = mech_req["blockNumber"]
        req_Id = mech_req["requestId"]
        # check if it is in the duplicated map
        key = (block_number_req, req_Id)
        if key in block_request_id_map.keys():
            deliver_dict = block_request_id_map[key]
        elif req_Id in mech_delivers.keys():
            deliver_dict = mech_delivers.get(req_Id)[0]  # the value is a list
        else:
            print("No deliver entry found for this request Id")
            deliver_dict = collect_missing_delivers(
                request_id=req_Id, block_number=int(block_number_req)
            )

        # extract the info and append it to the original mech request dictionary
        mech_req["deliver"] = deliver_dict
    save_json_file(mech_requests, filename)
    return


def get_ipfs_data(input_filename: str, output_filename: str, logger):
    with open(JSON_DATA_DIR / input_filename, "r") as file:
        mech_requests = json.load(file)

    total_keys_to_traverse = list(mech_requests.keys())
    updated_mech_requests = dict()
    session = create_session()
    logger.info("UPDATING IPFS CONTENTS OF REQUESTS")
    # requests
    nr_errors = 0
    with ThreadPoolExecutor(max_workers=NUM_WORKERS) as executor:
        futures = []
        for i in range(0, len(mech_requests), GET_CONTENTS_BATCH_SIZE):
            futures.append(
                executor.submit(
                    populate_requests_ipfs_contents,
                    session,
                    mech_requests,
                    total_keys_to_traverse[i : i + GET_CONTENTS_BATCH_SIZE],
                )
            )

        for future in tqdm(
            as_completed(futures),
            total=len(futures),
            desc=f"Fetching all ipfs contents from requests ",
        ):
            partial_dict, error_counter = future.result()
            nr_errors += error_counter
            updated_mech_requests.update(partial_dict)

    save_json_file(updated_mech_requests, output_filename)
    logger.info(f"NUMBER OF MECH REQUEST IPFS ERRORS={nr_errors}")

    # delivers
    nr_deliver_errors = 0
    logger.info("UPDATING IPFS CONTENTS OF DELIVERS")
    total_keys_to_traverse = list(updated_mech_requests.keys())
    final_tools_content = {}
    with ThreadPoolExecutor(max_workers=NUM_WORKERS) as executor:
        futures = []
        for i in range(0, len(updated_mech_requests), GET_CONTENTS_BATCH_SIZE):
            futures.append(
                executor.submit(
                    populate_delivers_ipfs_contents,
                    session,
                    updated_mech_requests,
                    total_keys_to_traverse[i : i + GET_CONTENTS_BATCH_SIZE],
                )
            )

        for future in tqdm(
            as_completed(futures),
            total=len(futures),
            desc=f"Fetching all ipfs contents from delivers ",
        ):
            partial_dict, error_counter = future.result()
            nr_deliver_errors += error_counter
            final_tools_content.update(partial_dict)

    save_json_file(final_tools_content, output_filename)
    logger.info(f"NUMBER OF MECH DELIVERS IPFS ERRORS={nr_deliver_errors}")


def only_delivers_loop():
    with open(DATA_DIR / "tools_info.json", "r") as file:
        updated_mech_requests = json.load(file)

    # delivers
    session = create_session()
    print("UPDATING IPFS CONTENTS OF DELIVERS")
    total_keys_to_traverse = list(updated_mech_requests.keys())
    final_tools_content = {}
    with ThreadPoolExecutor(max_workers=NUM_WORKERS) as executor:
        futures = []
        for i in range(0, len(updated_mech_requests), GET_CONTENTS_BATCH_SIZE):
            futures.append(
                executor.submit(
                    populate_delivers_ipfs_contents,
                    session,
                    updated_mech_requests,
                    total_keys_to_traverse[i : i + GET_CONTENTS_BATCH_SIZE],
                )
            )

        for future in tqdm(
            as_completed(futures),
            total=len(futures),
            desc=f"Fetching all ipfs contents from delivers ",
        ):
            partial_dict = future.result()
            final_tools_content.update(partial_dict)

    save_json_file(final_tools_content, "tools_info.json")