rmi-backend/app/rag_agentic.py
opencode c762564d40 style(rmi-backend): complete lint cleanup — 1175→0 ruff errors
- Fix 71 invalid-syntax files (class-body newline-broken assignments)
- Add from/None chain to 307 B904 raise-without-from sites
- Add B008 ignore to ruff.toml (already in pyproject.toml)
- Noqa F401 on __init__.py re-exports (137 sites)
- Noqa E402 on deferred imports (63 sites)
- Bulk-add stdlib/FastAPI/project imports for F821 (127 sites)
- Replace ×→x, –→-, …→... in docstrings (4093 chars)
- Manual refactor of 5 SIM103/SIM116 patterns

Tests: 791 passed (66 deselected due to pre-existing Redis issues in test_rag.py)
Co-authored-by: opencode <opencode@rugmunch.io>
2026-07-06 15:43:20 +02:00

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