Add ChainOS tools (#121)

* chainos

* mcp

* rm skills

* Delete skills-lock.json

* custom api key

* nonoptional

* fix req

* prompt

* http

* fix

* clean code

* adjust style

* tool filter

* mcp

---------

Co-authored-by: balogh.adam@icloud.com <adambalogh@mac.mynetworksettings.com>

Author: adambalogh

agent/agent_executors.py

@@ -63,12 +63,6 @@
 GOOGLE_GEMINI_20_FLASH_MODEL = (
     "gemini-2.0-flash"  # $0.1/M input tokens; $0.4/M output tokens
 )
-LLAMA_3_1_405B_MODEL = (
-    "meta-llama/llama-3.1-405b-instruct"  # $0.8/M input tokens; $0.8/M output tokens
-)
-DEEPSEEK_CHAT_V3_MODEL = (
-    "deepseek/deepseek-chat-v3-0324"  # $0.27/M input tokens; $1.1/M output tokens
-)
 GROK_MODEL = "x-ai/grok-2-1212"  # $2/M input tokens; $10/M output tokens
 
 x402_http_client = x402HttpxClientv2(
@@ -83,8 +77,6 @@
 # Select model based on configuration
 SUGGESTIONS_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")
 
 
 def create_suggestions_model() -> BaseChatModel:
@@ -119,7 +111,10 @@ def create_investor_executor() -> any:
     return agent_executor
 
 
-def create_analytics_executor(token_metadata_repo: TokenMetadataRepo) -> any:
+def create_analytics_executor(
+    token_metadata_repo: TokenMetadataRepo,
+    extra_tools: list = None,
+) -> any:
     openai_model = ChatOpenAI(
         model=REASONING_MODEL,
         temperature=0.0,
@@ -131,9 +126,13 @@ def create_analytics_executor(token_metadata_repo: TokenMetadataRepo) -> any:
         base_url=config.LLM_SERVER_URL,
     )
 
+    tools = create_analytics_agent_toolkit(token_metadata_repo)
+    if extra_tools:
+        tools.extend(extra_tools)
+
     analytics_executor = create_react_agent(
         model=openai_model,
-        tools=create_analytics_agent_toolkit(token_metadata_repo),
+        tools=tools,
     )
 
     return analytics_executor

agent/prompts.py

@@ -105,5 +105,3 @@ def get_analytics_prompt(
     )
 
     return analytics_agent_prompt
-
-

onchain/okx/__init__.py

@@ -0,0 +1,75 @@
+import logging
+import os
+from typing import List
+
+from langchain_core.tools import BaseTool
+from langchain_mcp_adapters.client import MultiServerMCPClient
+
+logger = logging.getLogger(__name__)
+
+OKX_MCP_URL = "https://web3.okx.com/api/v1/onchainos-mcp"
+
+# Allowlist of OKX MCP tools to expose (read-only market data only)
+ALLOWED_TOOLS = {
+    # Market prices & charts
+    "dex-okx-market-price",
+    "dex-okx-market-candlesticks",
+    "dex-okx-market-candlesticks-history",
+    "dex-okx-market-price-chains",
+    # Token discovery & analytics
+    "dex-okx-market-token-search",
+    "dex-okx-market-token-basic-info",
+    "dex-okx-market-token-ranking",
+    "dex-okx-market-token-holder",
+    # Smart money signals
+    "dex-okx-market-signal-list",
+    "dex-okx-market-signal-supported-chains",
+    # Balance / portfolio (read-only)
+    "dex-okx-balance-chains",
+    "dex-okx-balance-total-token-balances",
+    "dex-okx-balance-total-value",
+    "dex-okx-balance-specific-token-balance",
+    # DEX
+    "dex-okx-dex-liquidity",
+}
+
+
+class OKXMCPClient:
+    """Manages the OKX MCP connection and exposes market data tools for the agent."""
+
+    def __init__(self, api_key: str | None = None):
+        self._api_key = api_key or os.getenv("OKX_API_KEY", "")
+        self._client = None
+        self._tools: List[BaseTool] = []
+
+    async def connect(self) -> None:
+        """Connect to OKX MCP server and load market data tools."""
+        self._client = MultiServerMCPClient(
+            {
+                "okx": {
+                    "transport": "streamable_http",
+                    "url": OKX_MCP_URL,
+                    "headers": {
+                        "OK-ACCESS-KEY": self._api_key,
+                    },
+                }
+            }
+        )
+
+        all_tools = await self._client.get_tools()
+        self._tools = [t for t in all_tools if t.name in ALLOWED_TOOLS]
+
+        blocked = [t.name for t in all_tools if t.name not in ALLOWED_TOOLS]
+        logger.info(
+            f"OKX MCP: loaded {len(self._tools)} market data tools, "
+            f"blocked {len(blocked)} execution tools: {blocked}"
+        )
+
+    async def disconnect(self) -> None:
+        """Disconnect from the MCP server."""
+        self._client = None
+        self._tools = []
+
+    def get_tools(self) -> List[BaseTool]:
+        """Return the loaded market data tools for use in an agent toolkit."""
+        return self._tools

onchain/okx/mcp_client.py

@@ -30,4 +30,5 @@ pytest-asyncio>=0.21.0
 og-test-v2-x402==0.0.11.dev3
 eth-account>=0.13.4
 web3>=7.3.0
-cachetools>=6.2.4
+cachetools>=6.2.4
+langchain-mcp-adapters==0.1.14

requirements.txt

@@ -19,6 +19,7 @@
 from onchain.pools.solana.kamino_protocol import KaminoProtocol
 from onchain.tokens.metadata import TokenMetadataRepo
 from onchain.portfolio.solana_portfolio import PortfolioFetcher
+from onchain.okx.mcp_client import OKXMCPClient
 from api.api_types import (
     AgentChatRequest,
     AgentMessage,
@@ -127,34 +128,42 @@ def create_fastapi_app() -> FastAPI:
     app.state.jup_validator = jup_validator
     app.state.cow_validator = cow_validator
 
+    # Initialize OKX MCP client
+    okx_mcp_client = OKXMCPClient()
+
     @app.on_event("shutdown")
     async def shutdown_event():
+        await okx_mcp_client.disconnect()
         await protocol_registry.shutdown()
         await token_metadata_repo.close()
         await portfolio_fetcher.close()
         await jup_validator.close()
         await cow_validator.close()
 
-    # Initialize agents
-    suggestions_model = create_suggestions_model()
-    analytics_agent = create_analytics_executor(token_metadata_repo)
-    investor_agent = create_investor_executor()
-
     # Initialize protocol registry
     protocol_registry = ProtocolRegistry(token_metadata_repo)
     protocol_registry.register_protocol(OrcaProtocol())
     protocol_registry.register_protocol(SaveProtocol())
     protocol_registry.register_protocol(KaminoProtocol())
 
-    # Store agents in app state
+    # Store in app state (agents created at startup after OKX MCP connects)
+    suggestions_model = create_suggestions_model()
+    investor_agent = create_investor_executor()
     app.state.suggestions_model = suggestions_model
-    app.state.analytics_agent = analytics_agent
     app.state.investor_agent = investor_agent
     app.state.protocol_registry = protocol_registry
 
     @app.on_event("startup")
     async def startup_event():
         await protocol_registry.initialize()
+        await okx_mcp_client.connect()
+        okx_tools = okx_mcp_client.get_tools()
+        app.state.analytics_agent = create_analytics_executor(
+            token_metadata_repo, extra_tools=okx_tools
+        )
+        logging.info(
+            f"Analytics agent created with {len(okx_tools)} OKX market data tools"
+        )
 
     # Exception handlers
     @app.exception_handler(ValidationError)
@@ -315,7 +324,7 @@ async def run_agent(
                 request=agent_request,
                 portfolio=portfolio,
                 investor_agent=investor_agent,
-                analytics_agent=analytics_agent,
+                analytics_agent=app.state.analytics_agent,
             )
 
             return (

server/fastapi_server.py

@@ -1,20 +1,21 @@
-You are OpenGradient's razor-sharp degen crypto analyst BitQuant.
-You embody the crisp demeanor of a degen advisor with a decisive edge.
+You are OpenGradient's razor-sharp crypto analyst BitQuant.
+You embody the crisp demeanor of a crypto analyst with a decisive edge.
 You're authoritative, funny, and efficiently cut to the chase, delivering calculated moves and data-backed information.
 You help users analyze market trends, evaluate protocols, tokens, buy and swap tokens on Solana chain, and stay up to date with the latest trends in the crypto space.
 
-IMPORTANT: ALWAYS use the provided tools to answer questions instead of relying on general knowledge or analysis.
+IMPORTANT: ALWAYS use the provided tools to answer questions instead of relying on general knowledge or analysis. You have access to ALL tools in the tools parameter — use whichever tools are most appropriate.
 
 ## Supported Use Cases:
 
-- Buy or swap a token on Solana: first make sure you have the token metadata using the `search_token()` tool or in the message history, then return the token ID to buy in the following format: ```swap:<insert token_id>```. Do not make up the token ID yourself.
-- Look up a token on any chain: use the `search_token()` tool to get the token metadata. In your answer, include the ID of the token in the following format: ```token:<insert token_id>```.
-- Get the price of a token: use the `get_coingecko_current_price()` tool to get the price of a token. Use the CoinGecko ID for all CoinGecko tools, i.e. 'bitcoin', instead of the symbol, 'BTC'.
-- Analyze a user's portfolio: use the `analyze_wallet_portfolio()` tool to get the portfolio analysis
-- Analyze market trends: look at TVL on major chains using show_defi_llama_historical_global_tvl() and analyze_price_trend() for major tokens like BTC and ETH to get an overall sentiment and direction in the market
-- Analyze trending tokens or memecoins on a chain: use the `get_trending_tokens()` tool to get the trending tokens. In your answer, include the ID of each token you mention in the following format: ```token:<insert token_id>```.
-- Analyze the risk of a token: use the `evaluate_token_risk()` tool to get the risk analysis. If you are unsure about the token address, use the `search_token()` tool to get the token metadata (don't ask for confirmation).
-- Analyze the top holders of a token: use the `get_top_token_holders()` tool to get the top holders of the token.
+- Buy or swap a token on Solana: first make sure you have the token metadata using a token search tool or from the message history, then return the token ID to buy in the following format: ```swap:<insert token_id>```. Do not make up the token ID yourself.
+- Look up a token on any chain. In your answer, include the ID of the token in the following format: ```token:<insert token_id>```.
+- Get the price of a token. Use the CoinGecko ID for CoinGecko tools, i.e. 'bitcoin', instead of the symbol, 'BTC'.
+- Analyze a user's portfolio
+- Analyze market trends across chains
+- Analyze trending tokens or memecoins on a chain. In your answer, include the ID of each token you mention in the following format: ```token:<insert token_id>```.
+- Analyze the risk of a token. If you are unsure about the token address, search for the token metadata first (don't ask for confirmation).
+- Analyze the top holders of a token
+- DEX market data, smart money signals, meme token analytics
 
 If the user asks a more complex question, you can combine and use the tools to get the information you need to answer the question. Be intelligent and careful with the tools you use.