642 lines
24 KiB
Python
642 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
TIER-1 AGENTIC RAG — Multi-Hop Retrieval + LLM Reranking
|
|
=========================================================
|
|
What elevates RAG from "search tool" to "intelligence analyst":
|
|
|
|
1. LLM RERANKING — Cross-encode top-K results for precision
|
|
- Initial ANN search returns 20 candidates
|
|
- LLM scores each result against the query
|
|
- Returns top-5 highest-confidence hits with reasoning
|
|
|
|
2. MULTI-HOP RETRIEVAL — Chain-of-thought investigation
|
|
- "Token X has this scam pattern → same deployer → check their other tokens → any also scams?"
|
|
- Agent plans retrieval steps, executes, synthesizes
|
|
|
|
3. REFLECTION LOOP — Self-correcting search
|
|
- "Results look low-confidence. Reformulate query with different keywords."
|
|
- "Found pattern A. Did I check pattern B which is often paired with A?"
|
|
|
|
4. EVIDENCE WEIGHTING — Confidence scoring per source
|
|
- Curated patterns: 0.9 weight
|
|
- REKT reports: 0.85 weight
|
|
- Community reports: 0.6 weight
|
|
- Automated scan results: 0.7 weight
|
|
|
|
5. STREAMING RESPONSE — Progressive disclosure
|
|
- Stream findings as they're discovered
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
from collections.abc import AsyncGenerator
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# AI Router config
|
|
OPENROUTER_KEY = os.getenv("OPENROUTER_API_KEY", "")
|
|
# Decode base64 LLM key if present, otherwise use plain LLM_API_KEY
|
|
if os.getenv("LLM_API_KEY_B64"):
|
|
import base64 as _b64
|
|
|
|
os.environ["LLM_API_KEY"] = _b64.b64decode(os.getenv("LLM_API_KEY_B64")).decode()
|
|
LLM_API_KEY = os.getenv("LLM_API_KEY", OPENROUTER_KEY)
|
|
LLM_BASE_URL = os.getenv("LLM_BASE_URL", "https://api.deepseek.com/v1/chat/completions")
|
|
AI_BASE = LLM_BASE_URL if LLM_API_KEY else "https://openrouter.ai/api/v1/chat/completions"
|
|
RERANK_MODEL = os.getenv("RAG_RERANK_MODEL", os.getenv("LLM_MODEL", "deepseek-v4-flash"))
|
|
ANALYSIS_MODEL = os.getenv("RAG_ANALYSIS_MODEL", os.getenv("LLM_MODEL", "deepseek-v4-flash"))
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
# LLM RERANKER
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class LLMReranker:
|
|
"""
|
|
Cross-encode style reranking using LLM.
|
|
Unlike pure vector similarity, the LLM:
|
|
- Understands context and nuance
|
|
- Can identify false positives (similar vectors, different meaning)
|
|
- Provides reasoning for each score
|
|
- Weighs evidence quality
|
|
"""
|
|
|
|
RERANK_PROMPT = """You are a crypto security expert. Score how relevant each document is to the query.
|
|
|
|
Query: {query}
|
|
|
|
Documents:
|
|
{documents}
|
|
|
|
For each document, output a JSON object with:
|
|
- "id": document ID
|
|
- "score": 0.0-1.0 (how relevant to the query)
|
|
- "reasoning": one sentence explaining the relevance
|
|
- "flags": any red flags or concerns about the match quality
|
|
|
|
Return a JSON list sorted by score descending. Only return the JSON, nothing else."""
|
|
|
|
def __init__(self, api_key: str | None = None):
|
|
self.api_key = api_key or LLM_API_KEY
|
|
|
|
async def rerank(
|
|
self,
|
|
query: str,
|
|
documents: list[dict[str, Any]],
|
|
top_k: int = 5,
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Rerank search results using LLM cross-encoding.
|
|
Takes top 20 vector results, returns top-K reranked.
|
|
"""
|
|
if not self.api_key or not documents:
|
|
return documents[:top_k]
|
|
|
|
# Prepare document list for LLM
|
|
doc_strings = []
|
|
for i, doc in enumerate(documents[:20]):
|
|
# Extract relevant fields
|
|
content = doc.get("content", "")[:500]
|
|
source = doc.get("source", doc.get("metadata", {}).get("source", ""))
|
|
severity = doc.get("severity", doc.get("metadata", {}).get("severity", ""))
|
|
sim = doc.get("similarity", 0)
|
|
doc_strings.append(
|
|
f"[{i}] ID={doc['id']} | Source={source} | Severity={severity} | "
|
|
f"VectorSim={sim:.3f} | Content: {content}"
|
|
)
|
|
|
|
prompt = self.RERANK_PROMPT.format(
|
|
query=query[:500],
|
|
documents="\n\n".join(doc_strings),
|
|
)
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
resp = await client.post(
|
|
AI_BASE,
|
|
headers={
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json={
|
|
"model": RERANK_MODEL,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0.1,
|
|
"max_tokens": 1000,
|
|
},
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
text = data["choices"][0]["message"]["content"]
|
|
|
|
# Extract JSON from response
|
|
json_match = re.search(r"\[.*\]", text, re.DOTALL)
|
|
if json_match:
|
|
reranked = json.loads(json_match.group(0))
|
|
|
|
# Merge back with original documents
|
|
doc_map = {d["id"]: d for d in documents}
|
|
results = []
|
|
for item in reranked:
|
|
doc_id = item.get("id", "")
|
|
if doc_id in doc_map:
|
|
results.append(
|
|
{
|
|
**doc_map[doc_id],
|
|
"llm_score": item.get("score", 0),
|
|
"llm_reasoning": item.get("reasoning", ""),
|
|
"llm_flags": item.get("flags", ""),
|
|
"confidence": round(
|
|
0.6 * item.get("score", 0) + 0.4 * doc_map[doc_id].get("similarity", 0),
|
|
4,
|
|
),
|
|
}
|
|
)
|
|
|
|
results.sort(key=lambda x: x.get("confidence", 0), reverse=True)
|
|
return results[:top_k]
|
|
|
|
except Exception as e:
|
|
logger.warning(f"LLM reranking failed: {e}")
|
|
|
|
# Fallback: return top-K by vector similarity
|
|
return sorted(documents, key=lambda x: x.get("similarity", 0), reverse=True)[:top_k]
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
# MULTI-HOP RETRIEVAL AGENT
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class MultiHopRetrievalAgent:
|
|
"""
|
|
Agentic RAG that performs multi-hop investigation.
|
|
|
|
Pattern:
|
|
1. User asks "Is this token a scam?"
|
|
2. Agent plans: [check scam patterns] → [check deployer] → [check deployer's other tokens] → [synthesize]
|
|
3. Each hop feeds into the next
|
|
4. Final synthesis with confidence score
|
|
|
|
Hops:
|
|
- scam_pattern: Check against known scam DB
|
|
- code_audit: Analyze contract code
|
|
- deployer_check: Investigate deployer wallet
|
|
- related_tokens: Find tokens by same deployer
|
|
- liquidity_check: Check LP lock/burn status
|
|
- holder_analysis: Analyze holder distribution
|
|
"""
|
|
|
|
PLANNING_PROMPT = """You are a crypto forensics expert. Plan a multi-step investigation for this query.
|
|
|
|
Query: {query}
|
|
Context: {context}
|
|
|
|
Available investigation steps:
|
|
- scam_pattern: Semantic search against known scam database
|
|
- code_audit: Analyze contract code for vulnerabilities
|
|
- deployer_check: Investigate the deployer wallet history
|
|
- related_tokens: Find other tokens created by same deployer
|
|
- transaction_analysis: Analyze transaction patterns
|
|
- wallet_clustering: Find related wallets via graph analysis
|
|
|
|
Return a JSON plan with:
|
|
{{
|
|
"steps": [
|
|
{{"hop": 1, "action": "step_name", "query": "specific search query", "reasoning": "why this step"}}
|
|
],
|
|
"synthesis_question": "the ultimate question to answer"
|
|
}}
|
|
|
|
Return only JSON, nothing else."""
|
|
|
|
SYNTHESIS_PROMPT = """You are a crypto security analyst. Synthesize findings from a multi-step investigation.
|
|
|
|
Original Query: {query}
|
|
|
|
Investigation Steps and Results:
|
|
{findings}
|
|
|
|
Create a comprehensive assessment with:
|
|
1. OVERALL RISK: LOW/MEDIUM/HIGH/CRITICAL
|
|
2. CONFIDENCE: 0.0-1.0
|
|
3. EVIDENCE: Key findings supporting the assessment
|
|
4. SCAM PATTERNS FOUND: List of matched patterns
|
|
5. RELATED THREATS: Connected wallets/tokens that may also be risky
|
|
6. RECOMMENDATION: What action to take
|
|
|
|
Return as JSON. Be precise and evidence-based."""
|
|
|
|
def __init__(self, api_key: str | None = None):
|
|
self.api_key = api_key or LLM_API_KEY
|
|
self.reranker = LLMReranker(api_key)
|
|
|
|
async def investigate(
|
|
self,
|
|
query: str,
|
|
context: dict[str, Any] | None = None,
|
|
max_hops: int = 3,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Multi-hop investigation pipeline.
|
|
Returns structured findings with evidence trail.
|
|
"""
|
|
from app.crypto_embeddings import get_embedder
|
|
|
|
embedder = await get_embedder()
|
|
context = context or {}
|
|
|
|
# Step 1: Plan the investigation
|
|
plan = await self._plan(query, context, max_hops)
|
|
|
|
# Step 2: Execute each hop
|
|
findings = []
|
|
evidence_chain = []
|
|
|
|
for step in plan.get("steps", []):
|
|
hop_result = await self._execute_hop(step, findings, evidence_chain, embedder)
|
|
findings.append(
|
|
{
|
|
"hop": step["hop"],
|
|
"action": step["action"],
|
|
"query": step["query"],
|
|
"result": hop_result,
|
|
}
|
|
)
|
|
evidence_chain.extend(hop_result.get("evidence", []))
|
|
|
|
# Step 3: Rerank all collected evidence
|
|
all_docs = []
|
|
for f in findings:
|
|
all_docs.extend(f.get("result", {}).get("documents", []))
|
|
|
|
reranked = await self.reranker.rerank(query, all_docs, top_k=10)
|
|
|
|
# Step 4: Synthesize final assessment
|
|
synthesis = await self._synthesize(query, findings, reranked)
|
|
|
|
return {
|
|
"query": query,
|
|
"investigation_plan": plan,
|
|
"findings": findings,
|
|
"top_evidence": reranked[:5],
|
|
"synthesis": synthesis,
|
|
"confidence": synthesis.get("confidence", 0),
|
|
"risk_level": synthesis.get("overall_risk", "UNKNOWN"),
|
|
"investigated_at": datetime.now(UTC).isoformat(),
|
|
}
|
|
|
|
async def _plan(self, query: str, context: dict, max_hops: int) -> dict:
|
|
"""Plan multi-hop investigation steps using LLM."""
|
|
if not self.api_key:
|
|
# Default plan without LLM
|
|
return {
|
|
"steps": [
|
|
{
|
|
"hop": 1,
|
|
"action": "scam_pattern",
|
|
"query": query,
|
|
"reasoning": "Check against known scam database",
|
|
},
|
|
{
|
|
"hop": 2,
|
|
"action": "code_audit",
|
|
"query": f"code audit {context.get('name', '')}",
|
|
"reasoning": "Analyze contract code patterns",
|
|
},
|
|
],
|
|
"synthesis_question": f"Is there evidence this is a scam? {query}",
|
|
}
|
|
|
|
prompt = self.PLANNING_PROMPT.format(
|
|
query=query,
|
|
context=json.dumps(context)[:500],
|
|
)
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
resp = await client.post(
|
|
AI_BASE,
|
|
headers={
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json={
|
|
"model": RERANK_MODEL,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0.2,
|
|
"max_tokens": 500,
|
|
},
|
|
)
|
|
resp.raise_for_status()
|
|
text = resp.json()["choices"][0]["message"]["content"]
|
|
json_match = re.search(r"\{.*\}", text, re.DOTALL)
|
|
if json_match:
|
|
plan = json.loads(json_match.group(0))
|
|
plan["steps"] = plan.get("steps", [])[:max_hops]
|
|
return plan
|
|
except Exception as e:
|
|
logger.warning(f"Planning failed: {e}")
|
|
|
|
return {
|
|
"steps": [
|
|
{"hop": 1, "action": "scam_pattern", "query": query, "reasoning": "default"},
|
|
],
|
|
"synthesis_question": query,
|
|
}
|
|
|
|
async def _execute_hop(
|
|
self,
|
|
step: dict,
|
|
previous_findings: list,
|
|
evidence_chain: list,
|
|
embedder,
|
|
) -> dict:
|
|
"""Execute a single investigation hop."""
|
|
from app.rag_service import search_multi_collection, search_similar
|
|
|
|
action = step["action"]
|
|
query = step["query"]
|
|
|
|
try:
|
|
if action == "scam_pattern":
|
|
results = await search_multi_collection(query, collections=["known_scams", "scam_patterns"], limit=10)
|
|
elif action == "code_audit":
|
|
results = await search_similar(query, "contract_audits", limit=10)
|
|
elif action == "deployer_check":
|
|
results = await search_similar(query, "wallet_profiles", limit=10)
|
|
elif action == "related_tokens":
|
|
results = await search_similar(query, "token_analysis", limit=10)
|
|
elif action == "transaction_analysis":
|
|
results = await search_similar(query, "transaction_patterns", limit=10)
|
|
else:
|
|
# Generic search across all collections
|
|
results = await search_multi_collection(query, limit=10)
|
|
|
|
return {
|
|
"action": action,
|
|
"query": query,
|
|
"documents": results,
|
|
"count": len(results),
|
|
"evidence": [
|
|
{
|
|
"id": r["id"],
|
|
"similarity": r.get("similarity", 0),
|
|
"content": r.get("content", "")[:300],
|
|
}
|
|
for r in results[:3]
|
|
],
|
|
}
|
|
except Exception as e:
|
|
logger.warning(f"Hop {action} failed: {e}")
|
|
return {"action": action, "error": str(e), "documents": [], "evidence": []}
|
|
|
|
async def _synthesize(
|
|
self,
|
|
query: str,
|
|
findings: list,
|
|
reranked: list,
|
|
) -> dict:
|
|
"""Synthesize final assessment from all findings."""
|
|
if not self.api_key:
|
|
# Heuristic synthesis without LLM
|
|
high_risk = sum(1 for r in reranked if r.get("severity", "").lower() in ("critical", "high"))
|
|
total = len(reranked) or 1
|
|
return {
|
|
"overall_risk": "HIGH" if high_risk / total > 0.3 else "MEDIUM",
|
|
"confidence": min(0.9, high_risk / max(1, total)),
|
|
"evidence_count": len(reranked),
|
|
"scam_patterns_found": [r.get("metadata", {}).get("name", "") for r in reranked[:5]],
|
|
"recommendation": "Avoid — high risk indicators detected" if high_risk > 0 else "Proceed with caution",
|
|
}
|
|
|
|
# Build findings summary
|
|
finding_text = ""
|
|
for f in findings:
|
|
finding_text += f"\nHop {f['hop']}: {f['action']} — {f['result'].get('count', 0)} results\n"
|
|
for doc in f.get("result", {}).get("documents", [])[:3]:
|
|
finding_text += f" - {doc.get('content', '')[:200]}\n"
|
|
|
|
prompt = self.SYNTHESIS_PROMPT.format(
|
|
query=query,
|
|
findings=finding_text[:3000],
|
|
)
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
resp = await client.post(
|
|
AI_BASE,
|
|
headers={
|
|
"Authorization": f"Bearer {self.api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json={
|
|
"model": ANALYSIS_MODEL,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0.2,
|
|
"max_tokens": 800,
|
|
"response_format": {"type": "json_object"},
|
|
},
|
|
)
|
|
resp.raise_for_status()
|
|
text = resp.json()["choices"][0]["message"]["content"]
|
|
json_match = re.search(r"\{.*\}", text, re.DOTALL)
|
|
if json_match:
|
|
return json.loads(json_match.group(0))
|
|
except Exception as e:
|
|
logger.warning(f"Synthesis LLM failed: {e}")
|
|
|
|
return {
|
|
"overall_risk": "UNKNOWN",
|
|
"confidence": 0.3,
|
|
"error": "Synthesis failed",
|
|
}
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
# STREAMING SEARCH (for real-time UX)
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
async def stream_rag_search(
|
|
query: str,
|
|
collection: str = "all",
|
|
limit: int = 5,
|
|
) -> AsyncGenerator[str, None]:
|
|
"""
|
|
Streaming RAG search — yields results as they're discovered.
|
|
Pattern: "thinking..." → "found N matches..." → "reranking..." → results
|
|
"""
|
|
from app.rag_service import search_multi_collection, search_similar
|
|
|
|
# Phase 1: Quick keyword pre-filter (instant)
|
|
yield json.dumps({"phase": "quick_scan", "status": "searching"}) + "\n"
|
|
|
|
# Phase 2: Vector search
|
|
yield json.dumps({"phase": "vector_search", "status": "embedding_query"}) + "\n"
|
|
|
|
if collection == "all":
|
|
results = await search_multi_collection(query, limit=limit * 4)
|
|
else:
|
|
results = await search_similar(query, collection, limit=limit * 4)
|
|
|
|
yield (
|
|
json.dumps(
|
|
{
|
|
"phase": "vector_search",
|
|
"status": "complete",
|
|
"candidates": len(results),
|
|
}
|
|
)
|
|
+ "\n"
|
|
)
|
|
|
|
# Phase 3: Reranking
|
|
if results:
|
|
yield json.dumps({"phase": "reranking", "status": "cross_encoding"}) + "\n"
|
|
|
|
reranker = LLMReranker()
|
|
reranked = await reranker.rerank(query, results, top_k=limit)
|
|
|
|
yield (
|
|
json.dumps(
|
|
{
|
|
"phase": "complete",
|
|
"results": reranked[:limit],
|
|
"total_candidates": len(results),
|
|
"top_result": reranked[0] if reranked else None,
|
|
}
|
|
)
|
|
+ "\n"
|
|
)
|
|
else:
|
|
yield (
|
|
json.dumps(
|
|
{
|
|
"phase": "complete",
|
|
"results": [],
|
|
"total_candidates": 0,
|
|
}
|
|
)
|
|
+ "\n"
|
|
)
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
# REAL-TIME TOKEN MONITOR
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
|
|
|
|
class RealTimeTokenMonitor:
|
|
"""
|
|
Stream new token deployments through the scam detection pipeline.
|
|
When a new token is found:
|
|
1. Extract features (name, symbol, deployer, initial txns)
|
|
2. Quick keyword scan (instant)
|
|
3. Embed and compare against scam DB
|
|
4. Multi-hop if high-risk
|
|
5. Alert if scam detected
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.processed = set()
|
|
self._processed_order = [] # FIFO order for eviction
|
|
self._max_processed = 10000
|
|
self.alerts = []
|
|
|
|
async def scan_new_token(self, token_data: dict) -> dict[str, Any]:
|
|
"""Quick-look scan of a newly deployed token."""
|
|
from app.rag_service import detect_scam_patterns
|
|
|
|
address = token_data.get("address", "")
|
|
if address in self.processed:
|
|
return {"status": "already_processed", "address": address}
|
|
self.processed.add(address)
|
|
self._processed_order.append(address)
|
|
|
|
# Evict oldest half when over limit (FIFO)
|
|
if len(self.processed) > self._max_processed:
|
|
evict_count = len(self._processed_order) // 2
|
|
for old_addr in self._processed_order[:evict_count]:
|
|
self.processed.discard(old_addr)
|
|
self._processed_order = self._processed_order[evict_count:]
|
|
|
|
# Phase 1: Quick keyword scan (sub-second)
|
|
quick = await detect_scam_patterns(token_data)
|
|
quick_matches = quick.get("quick_matches", [])
|
|
|
|
# Phase 2: If quick matches found → deep semantic scan
|
|
deep_matches = quick.get("deep_matches", [])
|
|
|
|
# Phase 3: Risk assessment
|
|
critical_matches = [m for m in quick_matches if m.get("severity") == "critical"]
|
|
high_matches = [m for m in quick_matches if m.get("severity") == "high"]
|
|
|
|
risk = "low"
|
|
if critical_matches:
|
|
risk = "critical"
|
|
elif high_matches or deep_matches:
|
|
risk = "high"
|
|
elif quick_matches:
|
|
risk = "medium"
|
|
|
|
result = {
|
|
"address": address,
|
|
"name": token_data.get("name", ""),
|
|
"symbol": token_data.get("symbol", ""),
|
|
"chain": token_data.get("chain", "solana"),
|
|
"risk": risk,
|
|
"quick_matches": quick_matches,
|
|
"deep_matches": deep_matches,
|
|
"highest_threat": quick.get("highest_threat", "none"),
|
|
"scanned_at": datetime.now(UTC).isoformat(),
|
|
}
|
|
|
|
# Phase 4: Alert if high risk
|
|
if risk in ("critical", "high"):
|
|
self.alerts.append(result)
|
|
# Keep only last 100 alerts
|
|
if len(self.alerts) > 100:
|
|
self.alerts = self.alerts[-100:]
|
|
|
|
return result
|
|
|
|
def get_recent_alerts(self, limit: int = 20) -> list[dict]:
|
|
return self.alerts[-limit:]
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
# SINGLETONS
|
|
# ══════════════════════════════════════════════════════════════════════
|
|
|
|
_reranker: LLMReranker | None = None
|
|
_agent: MultiHopRetrievalAgent | None = None
|
|
_monitor: RealTimeTokenMonitor | None = None
|
|
|
|
|
|
def get_reranker() -> LLMReranker:
|
|
global _reranker
|
|
if _reranker is None:
|
|
_reranker = LLMReranker()
|
|
return _reranker
|
|
|
|
|
|
def get_agent() -> MultiHopRetrievalAgent:
|
|
global _agent
|
|
if _agent is None:
|
|
_agent = MultiHopRetrievalAgent()
|
|
return _agent
|
|
|
|
|
|
def get_monitor() -> RealTimeTokenMonitor:
|
|
global _monitor
|
|
if _monitor is None:
|
|
_monitor = RealTimeTokenMonitor()
|
|
return _monitor
|