@skynetxbt/venice-ai-plugin/src/index.ts

Venice AI Plugin for Skynet Framework - Staking Contract Integration

/* eslint-disable @typescript-eslint/no-explicit-any */
import { IPlugin, log } from "@skynetxbt/core";
import { EventEmitter } from "events";
import { ethers } from "ethers";
import { StakingContract } from "./StakingContract";
import { StakingContractConfig, StakingOperationResult, StakingViewResult } from "./types";
import { TokenOperationResult, VeniceToken } from "./VeniceToken";
import { STAKING_CONTRACT_ADDRESS } from "./constants";
import axios, { AxiosInstance } from "axios";

interface RateLimit {
	amount: number;
	type: "RPM" | "RPD" | "TPM";
}

interface ModelRateLimit {
	apiModelId: string;
	rateLimits: RateLimit[];
}

interface VeniceBalances {
	VCU: number;
	USD: number;
}

interface ApiTier {
	id: string;
	isCharged: boolean;
}

interface RateLimitResponse {
	data: {
		apiTier: ApiTier;
		accessPermitted: boolean;
		balances: VeniceBalances;
		rateLimits: ModelRateLimit[];
	};
}

export interface VenicePluginConfig {
	rpcUrl?: string;
	privateKey?: string;
	stakingContractAddress?: string;
	veniceTokenAddress?: string;
	veniceApiKey?: string;
	veniceBaseUrl?: string;
}

export class VenicePlugin extends EventEmitter implements IPlugin {
	readonly name: string = "venice-plugin";
	readonly description: string = "Interactive plugin for Venice AI staking";

	private stakingContract: StakingContract;
	private veniceToken: VeniceToken;
	private provider: ethers.Provider | null = null;
	private signer: ethers.Signer | null = null;
	private apiClient: AxiosInstance | null = null;

	constructor(private config: VenicePluginConfig = {}) {
		super();

		const stakingConfig: StakingContractConfig = {
			contractAddress: config.stakingContractAddress,
			privateKey: config.privateKey,
			rpcUrl: config.rpcUrl
		};

		this.stakingContract = new StakingContract(stakingConfig);
		this.veniceToken = new VeniceToken(config.veniceTokenAddress);

		if (config.veniceApiKey && config.veniceBaseUrl) {
			this.apiClient = axios.create({
				baseURL: config.veniceBaseUrl,
				headers: {
					'Authorization': `Bearer ${config.veniceApiKey}`,
					'Content-Type': 'application/json'
				}
			});
		}
	}

	async initialize(): Promise<void> {
		try {
			if (!this.config.rpcUrl) {
				throw new Error("RPC URL is required");
			}

			const provider = new ethers.JsonRpcProvider(this.config.rpcUrl);
			let signer: ethers.Signer | undefined;

			if (this.config.privateKey) {
				signer = new ethers.Wallet(this.config.privateKey, provider);
			}

			log.info("Initializing with config:", {
				rpcUrl: this.config.rpcUrl,
				stakingContractAddress: this.config.stakingContractAddress || STAKING_CONTRACT_ADDRESS,
				hasSigner: !!signer
			});

			await this.stakingContract.initialize(provider, signer);
			await this.veniceToken.initialize(provider, signer);

			this.provider = provider;
			this.signer = signer || null;
			this.emit("initialized");
		} catch (error) {
			throw new Error(`Failed to initialize plugin: ${error instanceof Error ? error.message : String(error)}`);
		}
	}

	async cleanup(): Promise<void> {
		await this.stakingContract.cleanup();
		await this.veniceToken.cleanup();
		this.provider = null;
		this.signer = null;
		this.emit("cleanup");
	}


	async getVCUsBalance(vvvBalance: string): Promise<string> {
		return (Number(vvvBalance) * 0.212).toString();
	}

	async getModelInfo(modelId: string): Promise<string> {

		const veniceModel = await axios.get(`https://api.venice.xyz/models/${modelId}`);
		return JSON.stringify(veniceModel.data);
	}

	async getTokenBalance(address: string): Promise<TokenOperationResult> {
		try {
			const result = await this.veniceToken.getBalance(address);

			if (result.success && result.data) {
				return {
					success: true,
					data: {
						raw: result.data,
						formatted: ethers.formatUnits(result.data, 18)
					}
				};
			}

			return result;
		} catch (error) {
			return {
				success: false,
				error: error instanceof Error ? error.message : String(error)
			};
		}
	}

	/**
	 * Check allowance for staking contract
	 */
	async checkAllowance(ownerAddress: string): Promise<TokenOperationResult> {
		try {
			const result = await this.veniceToken.getAllowance(
				ownerAddress,
				this.config.stakingContractAddress || STAKING_CONTRACT_ADDRESS
			);

			if (result.success && result.data) {
				return {
					success: true,
					data: {
						raw: result.data,
						formatted: ethers.formatUnits(result.data, 18)
					}
				};
			}

			return result;
		} catch (error) {
			return {
				success: false,
				error: error instanceof Error ? error.message : String(error)
			};
		}
	}

	/**
	 * Approve staking contract to spend tokens
	 */
	async approveStaking(amount: string): Promise<TokenOperationResult> {
		try {
			const amountBigInt = ethers.parseUnits(amount, 18);
			return await this.veniceToken.approve(
				this.config.stakingContractAddress || STAKING_CONTRACT_ADDRESS,
				amountBigInt
			);
		} catch (error) {
			return {
				success: false,
				error: error instanceof Error ? error.message : String(error)
			};
		}
	}

	// Staking operations

	async stake(recipient: string, amount: string): Promise<StakingOperationResult> {
		try {
			const amountBigInt = ethers.parseUnits(amount, 18);
			return await this.stakingContract.stake(recipient, amountBigInt);
		} catch (error) {
			return {
				success: false,
				error: error instanceof Error ? error.message : String(error)
			};
		}
	}

	async initiateUnstake(amount: string): Promise<StakingOperationResult> {
		try {
			const amountBigInt = ethers.parseUnits(amount, 18);
			return await this.stakingContract.initiateUnstake(amountBigInt);
		} catch (error) {
			return {
				success: false,
				error: error instanceof Error ? error.message : String(error)
			};
		}
	}

	async finalizeUnstake(): Promise<StakingOperationResult> {
		return await this.stakingContract.finalizeUnstake();
	}

	async claim(): Promise<StakingOperationResult> {
		return await this.stakingContract.claim();
	}

	// View functions

	async getPendingRewards(userAddress: string): Promise<StakingViewResult<string>> {
		try {
			const result = await this.stakingContract.getPendingRewards(userAddress);

			if (result.success && result.data) {
				return {
					success: true,
					data: ethers.formatUnits(result.data, 18)
				};
			}

			return {
				success: false,
				error: result.error
			};
		} catch (error) {
			return {
				success: false,
				error: error instanceof Error ? error.message : String(error)
			};
		}
	}

	async getVenicePercentage(): Promise<StakingViewResult<string>> {
		try {
			const result = await this.stakingContract.getVenicePercentage();

			if (result.success && result.data) {
				return {
					success: true,
					data: ethers.formatUnits(result.data, 18)
				};
			}

			return {
				success: false,
				error: result.error
			};
		} catch (error) {
			return {
				success: false,
				error: error instanceof Error ? error.message : String(error)
			};
		}
	}

	async getUserStakeInfo(userAddress: string): Promise<StakingViewResult<any>> {
		try {
			const result = await this.stakingContract.getUserStakeInfo(userAddress);

			if (result.success && result.data) {
				const { stakeAmount, pendingRewards, stakeInfo } = result.data;

				return {
					success: true,
					data: {
						stakeAmount: ethers.formatUnits(stakeAmount, 18),
						pendingRewards: ethers.formatUnits(pendingRewards, 18),
						stakeInfo: {
							rewardDebt: ethers.formatUnits(stakeInfo.rewardDebt, 18),
							cooldownEnd: stakeInfo.cooldownEnd.toString(),
							cooldownAmount: ethers.formatUnits(stakeInfo.cooldownAmount, 18)
						}
					}
				};
			}

			return {
				success: false,
				error: result.error
			};
		} catch (error) {
			return {
				success: false,
				error: error instanceof Error ? error.message : String(error)
			};
		}
	}

	async getStakingStats(): Promise<StakingViewResult<any>> {
		try {
			const result = await this.stakingContract.getStakingStats();

			if (result.success && result.data) {
				const {
					totalStaked,
					emissionRate,
					cooldownDuration,
					venicePercentage,
					accRewardPerShare,
					lastRewardTimestamp
				} = result.data;

				return {
					success: true,
					data: {
						totalStaked: ethers.formatUnits(totalStaked, 18),
						emissionRate: ethers.formatUnits(emissionRate, 18),
						cooldownDuration: cooldownDuration.toString(),
						venicePercentage: ethers.formatUnits(venicePercentage, 18),
						accRewardPerShare: ethers.formatUnits(accRewardPerShare, 36), // ACC_REWARD_SCALE is 1e36
						lastRewardTimestamp: lastRewardTimestamp.toString()
					}
				};
			}

			return {
				success: false,
				error: result.error
			};
		} catch (error) {
			return {
				success: false,
				error: error instanceof Error ? error.message : String(error)
			};
		}
	}

	async getRateLimits(): Promise<RateLimitResponse | null> {
		if (!this.apiClient) {
			throw new Error("Venice API client not initialized. Please provide API key and base URL.");
		}

		try {
			const response = await this.apiClient.get<RateLimitResponse>('/api/v1/api_keys/rate_limits');
			return response.data;
		} catch (error) {
			log.error(`Failed to fetch rate limits: ${error}`);
			return null;
		}
	}

	async getModelRateLimits(modelId: string): Promise<ModelRateLimit | null> {
		const rateLimits = await this.getRateLimits();
		if (!rateLimits) return null;

		return rateLimits.data.rateLimits.find(limit => limit.apiModelId === modelId) || null;
	}

	async getVeniceBalances(): Promise<VeniceBalances | null> {
		const rateLimits = await this.getRateLimits();
		return rateLimits?.data.balances || null;
	}
}