diff --git a/agent/agent_executors.py b/agent/agent_executors.py index fe314cb..d17a90f 100644 --- a/agent/agent_executors.py +++ b/agent/agent_executors.py @@ -1,4 +1,5 @@ import os +import httpx from openai import OpenAI from langgraph.prebuilt import create_react_agent @@ -8,6 +9,31 @@ from agent.tools import create_investor_agent_toolkit, create_analytics_agent_toolkit from onchain.tokens.metadata import TokenMetadataRepo from server import config +from web3 import Web3 +from x402.clients.base import x402Client +from x402.types import x402PaymentRequiredResponse +from langchain_openai import ChatOpenAI +from .x402 import X402Auth + +WEB3_CONFIG = Web3(Web3.HTTPProvider(config.OG_RPC_URL)) +WALLET_ACCOUNT = WEB3_CONFIG.eth.account.from_key( + config.WALLET_PRIV_KEY +) + +TIMEOUT = httpx.Timeout( + timeout=90.0, + connect=15.0, + read=15.0, + write=30.0, + pool=10.0, +) + +LIMITS = httpx.Limits( + max_keepalive_connections=100, + max_connections=500, + keepalive_expiry=60 * 20, # 20 minutes +) + ## # Subnet LLM Configuration @@ -35,11 +61,21 @@ ) GROK_MODEL = "x-ai/grok-2-1212" # $2/M input tokens; $10/M output tokens +x402_http_client = httpx.AsyncClient( + base_url=config.LLM_SERVER_URL, + headers={"Authorization": f"Bearer {config.DUMMY_X402_API_KEY}"}, + timeout=TIMEOUT, + limits=LIMITS, + http2=False, + follow_redirects=False, + auth=X402Auth(account=WALLET_ACCOUNT), # type: ignore +) + # Select model based on configuration if not config.SUBNET_MODE: SUGGESTIONS_MODEL = GOOGLE_GEMINI_20_FLASH_MODEL - ROUTING_MODEL = GOOGLE_GEMINI_FLASH_15_8B_MODEL + ROUTING_MODEL = GOOGLE_GEMINI_20_FLASH_MODEL REASONING_MODEL = GOOGLE_GEMINI_20_FLASH_MODEL BASE_URL = "https://generativelanguage.googleapis.com/v1beta/" API_KEY = os.getenv("GEMINI_API_KEY") @@ -52,30 +88,44 @@ def create_routing_model() -> BaseChatModel: - return ChatGoogleGenerativeAI( - model=ROUTING_MODEL, - temperature=0.0, - google_api_key=API_KEY, - max_tokens=500, - ) + return ChatOpenAI( + model=ROUTING_MODEL, + temperature=0.0, + max_tokens=500, + api_key=config.DUMMY_X402_API_KEY, + http_async_client=x402_http_client, + stream_usage=False, + streaming=False, + base_url=config.LLM_SERVER_URL, + ) + def create_suggestions_model() -> BaseChatModel: - return ChatGoogleGenerativeAI( - model=SUGGESTIONS_MODEL, - temperature=0.3, - google_api_key=API_KEY, - max_tokens=1000, - ) + return ChatOpenAI( + model=SUGGESTIONS_MODEL, + temperature=0.3, + max_tokens=1000, + api_key=config.DUMMY_X402_API_KEY, + http_async_client=x402_http_client, + stream_usage=False, + streaming=False, + base_url=config.LLM_SERVER_URL, + ) + def create_investor_executor() -> any: - openai_model = ChatGoogleGenerativeAI( - model=REASONING_MODEL, - temperature=0.0, - google_api_key=API_KEY, - max_tokens=4096, - ) + openai_model = ChatOpenAI( + model=REASONING_MODEL, + temperature=0.0, + api_key=config.DUMMY_X402_API_KEY, + http_async_client=x402_http_client, + stream_usage=False, + streaming=False, + base_url=config.LLM_SERVER_URL, + ) + agent_executor = create_react_agent( model=openai_model, tools=create_investor_agent_toolkit() ) @@ -84,12 +134,17 @@ def create_investor_executor() -> any: def create_analytics_executor(token_metadata_repo: TokenMetadataRepo) -> any: - openai_model = ChatGoogleGenerativeAI( - model=REASONING_MODEL, - temperature=0.0, - google_api_key=API_KEY, - max_tokens=4096, - ) + openai_model = ChatOpenAI( + model=REASONING_MODEL, + temperature=0.0, + max_tokens=4096, + api_key=config.DUMMY_X402_API_KEY, + http_async_client=x402_http_client, + stream_usage=False, + streaming=False, + base_url=config.LLM_SERVER_URL, + ) + analytics_executor = create_react_agent( model=openai_model, tools=create_analytics_agent_toolkit(token_metadata_repo), diff --git a/agent/x402.py b/agent/x402.py new file mode 100644 index 0000000..4fc69c2 --- /dev/null +++ b/agent/x402.py @@ -0,0 +1,60 @@ +import httpx +import typing +import logging + +from x402.clients.base import x402Client +from x402.types import x402PaymentRequiredResponse, PaymentRequirements + + +class X402Auth(httpx.Auth): + """Auth class for handling x402 payment requirements.""" + + def __init__( + self, + account: typing.Any, + max_value: typing.Optional[int] = None, + payment_requirements_selector: typing.Optional[ + typing.Callable[ + [ + list[PaymentRequirements], + typing.Optional[str], + typing.Optional[str], + typing.Optional[int], + ], + PaymentRequirements, + ] + ] = None, + ): + self.x402_client = x402Client( + account, + max_value=max_value, + payment_requirements_selector=payment_requirements_selector, # type: ignore + ) + + async def async_auth_flow( + self, request: httpx.Request + ) -> typing.AsyncGenerator[httpx.Request, httpx.Response]: + response = yield request + + if response.status_code == 402: + try: + await response.aread() + data = response.json() + + payment_response = x402PaymentRequiredResponse(**data) + + selected_requirements = self.x402_client.select_payment_requirements( + payment_response.accepts + ) + + payment_header = self.x402_client.create_payment_header( + selected_requirements, payment_response.x402_version + ) + + request.headers["X-Payment"] = payment_header + request.headers["Access-Control-Expose-Headers"] = "X-Payment-Response" + yield request + + except Exception as e: + logging.error(f"X402Auth: Error handling payment: {e}") + return diff --git a/requirements.txt b/requirements.txt index 075afe0..79accba 100644 --- a/requirements.txt +++ b/requirements.txt @@ -26,3 +26,6 @@ langchain_google_genai>=2.0.0 aioboto3>=1.38.0 async_lru>=2.0.0 aiolimiter>=1.2.0 +og-test-x402==0.0.5 +eth-account>=0.13.4 +web3>=7.3.0 diff --git a/server/config.py b/server/config.py index dd6bcd5..3f9c24c 100644 --- a/server/config.py +++ b/server/config.py @@ -8,6 +8,11 @@ SUBNET_MODE = os.getenv("subnet_mode", "false").lower() == "true" logging.info(f"Running in subnet mode: {SUBNET_MODE}") +DUMMY_X402_API_KEY = os.getenv("DUMMY_X402_API_KEY", "dummy") +LLM_SERVER_URL: str = os.getenv("LLM_SERVER_URL", "http://127.0.0.1:8080/v1") +OG_RPC_URL: str = os.getenv("OG_RPC_URL", "https://eth-devnet.opengradient.ai") +WALLET_PRIV_KEY: str = os.getenv("WALLET_PRIV_KEY") + # Bypass daily limit for miner wallet MINER_TOKEN = os.getenv("MINER_TOKEN")