my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	"""MAWD Solana Agent - Web API Server with WebSocket streaming"""
#2
#3	import asyncio
#4	import json
#5	import uuid
#6	import sys
#7	from pathlib import Path
#8	from typing import Optional
#9	from datetime import datetime
#10	from dataclasses import dataclass, asdict
#11	from contextlib import asynccontextmanager
#12
#13	from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
#14	from fastapi.middleware.cors import CORSMiddleware
#15	from fastapi.staticfiles import StaticFiles
#16	from fastapi.responses import FileResponse
#17	from pydantic import BaseModel
#18
#19	# Add parent to path
#20	sys.path.insert(0, str(Path(__file__).parent.parent))
#21
#22	from mini_agent.config import load_config
#23	from mini_agent.agent import SolanaAgent
#24	from mini_agent.tools.base import ToolResult
#25
#26
#27	# ============================================================
#28	# Event Types for WebSocket streaming
#29	# ============================================================
#30
#31	@dataclass
#32	class AgentEvent:
#33	"""Event sent to frontend via WebSocket"""
#34	type: str
#35	timestamp: str
#36	data: dict
#37
#38	def to_json(self):
#39	return json.dumps(asdict(self))
#40
#41
#42	class ChatMessage(BaseModel):
#43	"""Incoming chat message from frontend"""
#44	message: str
#45
#46
#47	class AgentStatus(BaseModel):
#48	"""Agent status response"""
#49	status: str
#50	wallet: Optional[str] = None
#51	balance: Optional[float] = None
#52	tools_count: int = 0
#53	active_sessions: int = 0
#54
#55
#56	# ============================================================
#57	# Connection Manager for WebSocket clients
#58	# ============================================================
#59
#60	class ConnectionManager:
#61	"""Manages WebSocket connections"""
#62
#63	def __init__(self):
#64	self.active_connections: dict[str, WebSocket] = {}
#65
#66	async def connect(self, websocket: WebSocket, client_id: str):
#67	await websocket.accept()
#68	self.active_connections[client_id] = websocket
#69
#70	def disconnect(self, client_id: str):
#71	self.active_connections.pop(client_id, None)
#72
#73	async def send_event(self, client_id: str, event: AgentEvent):
#74	if client_id in self.active_connections:
#75	await self.active_connections[client_id].send_text(event.to_json())
#76
#77	async def broadcast(self, event: AgentEvent):
#78	for websocket in self.active_connections.values():
#79	await websocket.send_text(event.to_json())
#80
#81
#82	manager = ConnectionManager()
#83
#84
#85	# ============================================================
#86	# Streaming Agent Wrapper
#87	# ============================================================
#88
#89	class StreamingAgent:
#90	"""Agent wrapper that streams events to WebSocket clients"""
#91
#92	def __init__(self, agent: SolanaAgent, client_id: str):
#93	self.agent = agent
#94	self.client_id = client_id
#95
#96	async def emit(self, event_type: str, data: dict):
#97	event = AgentEvent(
#98	type=event_type,
#99	timestamp=datetime.utcnow().isoformat(),
#100	data=data
#101	)
#102	await manager.send_event(self.client_id, event)
#103
#104	async def process_message(self, message: str) -> str:
#105	"""Process a message and stream all events"""
#106
#107	# Emit user message
#108	await self.emit("user_message", {"content": message})
#109
#110	# Add to agent
#111	self.agent.add_user_message(message)
#112
#113	step = 0
#114	while step < self.agent.max_steps:
#115	step += 1
#116
#117	# Emit step start
#118	await self.emit("step_start", {"step": step, "max_steps": self.agent.max_steps})
#119
#120	# Get LLM response
#121	try:
#122	response = await self.agent.llm_client.generate(
#123	self.agent.messages,
#124	self.agent.tools
#125	)
#126	except Exception as e:
#127	await self.emit("error", {"message": str(e)})
#128	return f"Error: {e}"
#129
#130	# Emit thinking
#131	if response.thinking:
#132	await self.emit("thinking", {"content": response.thinking[:500]})
#133
#134	# Emit assistant response
#135	if response.content:
#136	await self.emit("assistant_message", {"content": response.content})
#137
#138	# Add to message history
#139	from mini_agent.llm import Message
#140	self.agent.messages.append(Message(
#141	role="assistant",
#142	content=response.content,
#143	thinking=response.thinking,
#144	tool_calls=response.tool_calls,
#145	))
#146
#147	# If no tool calls, we're done
#148	if not response.tool_calls:
#149	await self.emit("step_complete", {"step": step, "final": True})
#150	return response.content or ""
#151
#152	# Execute tool calls
#153	for tool_call in response.tool_calls:
#154	# Emit tool call start
#155	await self.emit("tool_call", {
#156	"name": tool_call.name,
#157	"arguments": tool_call.arguments,
#158	})
#159
#160	# Find and execute tool
#161	tool = None
#162	for t in self.agent.tools:
#163	if t.name == tool_call.name:
#164	tool = t
#165	break
#166
#167	if tool is None:
#168	result = ToolResult(success=False, error=f"Unknown tool: {tool_call.name}")
#169	else:
#170	try:
#171	result = await tool.execute(**tool_call.arguments)
#172	except Exception as e:
#173	result = ToolResult(success=False, error=str(e))
#174
#175	# Emit tool result
#176	await self.emit("tool_result", {
#177	"name": tool_call.name,
#178	"success": result.success,
#179	"content": result.content if result.success else None,
#180	"error": result.error if not result.success else None,
#181	})
#182
#183	# Add tool result to messages
#184	self.agent.messages.append(Message(
#185	role="tool",
#186	content=result.content if result.success else f"Error: {result.error}",
#187	tool_call_id=tool_call.id,
#188	name=tool_call.name,
#189	))
#190
#191	await self.emit("step_complete", {"step": step, "final": False})
#192
#193	return f"Max steps ({self.agent.max_steps}) reached."
#194
#195
#196	# ============================================================
#197	# Global Agent and Birdeye WebSocket Instances
#198	# ============================================================
#199
#200	agent_instance: Optional[SolanaAgent] = None
#201	agent_lock = asyncio.Lock()
#202	birdeye_ws_client = None
#203
#204
#205	async def get_agent() -> SolanaAgent:
#206	"""Get or create the global agent instance"""
#207	global agent_instance
#208
#209	async with agent_lock:
#210	if agent_instance is None:
#211	env_path = Path(__file__).parent.parent.parent / ".env.local"
#212	if not env_path.exists():
#213	env_path = Path(__file__).parent.parent / ".env.local"
#214
#215	config = load_config(str(env_path) if env_path.exists() else None)
#216	agent_instance = SolanaAgent(config)
#217	await agent_instance.initialize()
#218
#219	return agent_instance
#220
#221
#222	# ============================================================
#223	# FastAPI Application
#224	# ============================================================
#225
#226	async def initialize_birdeye_websocket():
#227	"""Initialize Birdeye WebSocket for real-time data"""
#228	global birdeye_ws_client
#229
#230	try:
#231	# Import here to avoid circular imports and use correct path
#232	import sys
#233	from pathlib import Path
#234	sys.path.insert(0, str(Path(__file__).parent.parent))
#235	from clients.birdeye_client import BirdeyeWebSocketClient
#236	from config import load_config as load_root_config
#237
#238	# Load config to get API key
#239	env_path = Path(__file__).parent.parent.parent / ".env.local"
#240	if not env_path.exists():
#241	env_path = Path(__file__).parent.parent / ".env.local"
#242
#243	config = load_root_config(str(env_path) if env_path.exists() else None)
#244
#245	# Create WebSocket client
#246	birdeye_ws_client = BirdeyeWebSocketClient(api_key=config.birdeye_api_key)
#247
#248	# Register event handlers to broadcast to all WebSocket clients
#249	async def handle_price_update(data):
#250	await manager.broadcast(AgentEvent(
#251	type="birdeye_price_update",
#252	timestamp=datetime.utcnow().isoformat(),
#253	data=data
#254	))
#255
#256	async def handle_new_listing(data):
#257	await manager.broadcast(AgentEvent(
#258	type="birdeye_new_listing",
#259	timestamp=datetime.utcnow().isoformat(),
#260	data=data
#261	))
#262
#263	async def handle_large_trade(data):
#264	await manager.broadcast(AgentEvent(
#265	type="birdeye_large_trade",
#266	timestamp=datetime.utcnow().isoformat(),
#267	data=data
#268	))
#269
#270	async def handle_wallet_tx(data):
#271	await manager.broadcast(AgentEvent(
#272	type="birdeye_wallet_tx",
#273	timestamp=datetime.utcnow().isoformat(),
#274	data=data
#275	))
#276
#277	# Register handlers
#278	birdeye_ws_client.on("PRICE_UPDATE", handle_price_update)
#279	birdeye_ws_client.on("NEW_LISTING", handle_new_listing)
#280	birdeye_ws_client.on("LARGE_TRADE", handle_large_trade)
#281	birdeye_ws_client.on("WALLET_TX", handle_wallet_tx)
#282
#283	# Connect
#284	await birdeye_ws_client.connect()
#285
#286	# Subscribe to new listings
#287	await birdeye_ws_client.subscribe_new_listings()
#288
#289	print("✓ Birdeye WebSocket connected and subscribed to new listings")
#290
#291	except Exception as e:
#292	print(f"⚠️ Failed to initialize Birdeye WebSocket: {e}")
#293	birdeye_ws_client = None
#294
#295
#296	@asynccontextmanager
#297	async def lifespan(app: FastAPI):
#298	"""Application lifespan events"""
#299	# Startup
#300	print("🚀 Starting MAWD API Server...")
#301	await initialize_birdeye_websocket()
#302	yield
#303	# Shutdown
#304	if agent_instance:
#305	await agent_instance.close()
#306	if birdeye_ws_client:
#307	await birdeye_ws_client.disconnect()
#308	print("👋 MAWD API Server stopped")
#309
#310
#311	app = FastAPI(
#312	title="MAWD Solana Agent API",
#313	description="AI-powered Solana trading agent with real-time streaming",
#314	version="1.0.0",
#315	lifespan=lifespan,
#316	)
#317
#318	# CORS for frontend
#319	app.add_middleware(
#320	CORSMiddleware,
#321	allow_origins=["*"],
#322	allow_credentials=True,
#323	allow_methods=["*"],
#324	allow_headers=["*"],
#325	)
#326
#327
#328	# ============================================================
#329	# REST Endpoints
#330	# ============================================================
#331
#332	@app.get("/api/status", response_model=AgentStatus)
#333	async def get_status():
#334	"""Get agent status"""
#335	try:
#336	agent = await get_agent()
#337	balance = None
#338	try:
#339	balance = await agent.helius_client.get_sol_balance(agent.bags_client.wallet_pubkey)
#340	except Exception:
#341	pass
#342
#343	return AgentStatus(
#344	status="ready",
#345	wallet=agent.bags_client.wallet_pubkey,
#346	balance=balance,
#347	tools_count=len(agent.tools),
#348	active_sessions=len(manager.active_connections),
#349	)
#350	except Exception as e:
#351	return AgentStatus(status=f"error: {str(e)}", tools_count=0, active_sessions=0)
#352
#353
#354	@app.get("/api/tools")
#355	async def get_tools():
#356	"""Get list of available tools"""
#357	agent = await get_agent()
#358	return {
#359	"tools": [
#360	{
#361	"name": tool.name,
#362	"description": tool.description,
#363	"parameters": tool.parameters,
#364	}
#365	for tool in agent.tools
#366	]
#367	}
#368
#369
#370	@app.post("/api/chat")
#371	async def chat(msg: ChatMessage):
#372	"""Send a message (non-streaming)"""
#373	agent = await get_agent()
#374	response = await agent.process_message(msg.message)
#375	return {"response": response}
#376
#377
#378	@app.get("/api/history")
#379	async def get_history():
#380	"""Get conversation history"""
#381	agent = await get_agent()
#382	return {
#383	"messages": [
#384	{
#385	"role": m.role,
#386	"content": m.content,
#387	}
#388	for m in agent.messages
#389	if m.role != "system"
#390	]
#391	}
#392
#393
#394	@app.post("/api/clear")
#395	async def clear_history():
#396	"""Clear conversation history"""
#397	agent = await get_agent()
#398	agent.messages = [agent.messages[0]] # Keep system prompt
#399	return {"status": "cleared"}
#400
#401
#402	# ============================================================
#403	# Birdeye WebSocket Subscription Endpoints
#404	# ============================================================
#405
#406	@app.post("/api/birdeye/subscribe/price/{token_address}")
#407	async def subscribe_token_price(token_address: str):
#408	"""Subscribe to real-time price updates for a token"""
#409	if not birdeye_ws_client:
#410	raise HTTPException(status_code=503, detail="Birdeye WebSocket not available")
#411
#412	try:
#413	await birdeye_ws_client.subscribe_price(token_address)
#414	return {"status": "subscribed", "token": token_address, "type": "price"}
#415	except Exception as e:
#416	raise HTTPException(status_code=500, detail=str(e))
#417
#418
#419	@app.post("/api/birdeye/subscribe/wallet/{wallet_address}")
#420	async def subscribe_wallet_transactions(wallet_address: str):
#421	"""Subscribe to real-time transactions for a wallet"""
#422	if not birdeye_ws_client:
#423	raise HTTPException(status_code=503, detail="Birdeye WebSocket not available")
#424
#425	try:
#426	await birdeye_ws_client.subscribe_wallet_transactions(wallet_address)
#427	return {"status": "subscribed", "wallet": wallet_address, "type": "wallet_tx"}
#428	except Exception as e:
#429	raise HTTPException(status_code=500, detail=str(e))
#430
#431
#432	@app.post("/api/birdeye/subscribe/large-trades")
#433	async def subscribe_large_trades(threshold_usd: float = 10000):
#434	"""Subscribe to large trade alerts"""
#435	if not birdeye_ws_client:
#436	raise HTTPException(status_code=503, detail="Birdeye WebSocket not available")
#437
#438	try:
#439	await birdeye_ws_client.subscribe_large_trades(threshold_usd)
#440	return {"status": "subscribed", "type": "large_trades", "threshold": threshold_usd}
#441	except Exception as e:
#442	raise HTTPException(status_code=500, detail=str(e))
#443
#444
#445	@app.get("/api/birdeye/status")
#446	async def birdeye_websocket_status():
#447	"""Get Birdeye WebSocket connection status"""
#448	if not birdeye_ws_client:
#449	return {"connected": False, "message": "WebSocket client not initialized"}
#450
#451	return {
#452	"connected": birdeye_ws_client.is_connected,
#453	"subscriptions": len(birdeye_ws_client.handlers),
#454	"reconnect_attempts": birdeye_ws_client.reconnect_attempts
#455	}
#456
#457
#458	# ============================================================
#459	# WebSocket Endpoint for Streaming
#460	# ============================================================
#461
#462	@app.websocket("/ws/{client_id}")
#463	async def websocket_endpoint(websocket: WebSocket, client_id: str):
#464	"""WebSocket for real-time agent streaming"""
#465	await manager.connect(websocket, client_id)
#466
#467	try:
#468	agent = await get_agent()
#469	streaming_agent = StreamingAgent(agent, client_id)
#470
#471	# Send connected event
#472	await manager.send_event(client_id, AgentEvent(
#473	type="connected",
#474	timestamp=datetime.utcnow().isoformat(),
#475	data={
#476	"client_id": client_id,
#477	"wallet": agent.bags_client.wallet_pubkey,
#478	"tools": len(agent.tools),
#479	}
#480	))
#481
#482	while True:
#483	# Wait for message from client
#484	data = await websocket.receive_text()
#485	message_data = json.loads(data)
#486
#487	if message_data.get("type") == "chat":
#488	user_message = message_data.get("message", "")
#489	await streaming_agent.process_message(user_message)
#490
#491	elif message_data.get("type") == "ping":
#492	await manager.send_event(client_id, AgentEvent(
#493	type="pong",
#494	timestamp=datetime.utcnow().isoformat(),
#495	data={}
#496	))
#497
#498	except WebSocketDisconnect:
#499	manager.disconnect(client_id)
#500	except Exception as e:
#501	await manager.send_event(client_id, AgentEvent(
#502	type="error",
#503	timestamp=datetime.utcnow().isoformat(),
#504	data={"message": str(e)}
#505	))
#506	manager.disconnect(client_id)
#507
#508
#509	# ============================================================
#510	# Serve Frontend
#511	# ============================================================
#512
#513	# Mount static files if frontend exists
#514	frontend_path = Path(__file__).parent / "frontend"
#515	if frontend_path.exists():
#516	app.mount("/static", StaticFiles(directory=str(frontend_path)), name="static")
#517
#518	@app.get("/")
#519	async def serve_frontend():
#520	"""Serve the frontend"""
#521	index_path = frontend_path / "index.html"
#522	if index_path.exists():
#523	return FileResponse(str(index_path))
#524	return {"message": "MAWD API Server - Frontend not found. Use /api/* endpoints."}
#525
#526	@app.get("/swap")
#527	async def serve_swap():
#528	"""Serve the Jupiter swap interface"""
#529	swap_path = frontend_path / "swap.html"
#530	if swap_path.exists():
#531	return FileResponse(str(swap_path))
#532	return {"message": "Swap interface not found"}
#533
#534
#535	# ============================================================
#536	# Main
#537	# ============================================================
#538
#539	if __name__ == "__main__":
#540	import uvicorn
#541	uvicorn.run(app, host="0.0.0.0", port=8001)
#542

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public