RugMunchMedia/rmi-backend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

merge: chore/cleanup-remove-bloat-and-secrets into main

Browse source
This commit is contained in:
Crypto Rug Munch 2026-07-02 01:24:22 +07:00
commit bde2f3a97d
1173 changed files with 437609 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

511
app/query_transform.py Normal file
View file

@ -0,0 +1,511 @@
"""
Query Transformation Pipeline for RMI RAG System
===================================================
Improves search recall by transforming queries before retrieval.
Three strategies:
1. HyDE -- Hypothetical Document Embedding (generate hypothetical answer, embed that)
2. Query Expansion -- Generate variant phrasings, search all, merge with RRF
3. Step-Back Prompting -- Abstract the query to a broader form, retrieve for both
All strategies have LLM-backed and rule-based fallback paths.
"""
import hashlib
import json
import logging
import os
import re
from dataclasses import dataclass, field
import httpx
logger = logging.getLogger(__name__)
# -- Redis config (same as rag_service) --
REDIS_HOST = os.getenv("REDIS_HOST", "rmi-redis")
REDIS_PORT = int(os.getenv("REDIS_PORT", "6379"))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", "")
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"
AI_MODEL = os.getenv("QUERY_TRANSFORM_MODEL", os.getenv("LLM_MODEL", "deepseek-v4-flash"))
# ======================================================================
# DATA MODEL
# ======================================================================
@dataclass
class TransformedQuery:
original_query: str
strategy: str # "hyde" | "expand" | "step_back" | "none"
transformed_queries: list[str] # list of strings to embed & search
metadata: dict = field(default_factory=dict)
# ======================================================================
# REDIS HELPER
# ======================================================================
async def _get_redis():
import redis.asyncio as ai_redis
return ai_redis.Redis(
host=REDIS_HOST,
port=REDIS_PORT,
password=REDIS_PASSWORD or None,
db=0,
decode_responses=True,
)
# ======================================================================
# LLM HELPER
# ======================================================================
async def _call_llm(prompt: str, max_tokens: int = 512) -> str | None:
"""Call LLM chat API (DeepSeek primary, OpenRouter fallback). Returns None if unavailable."""
if not LLM_API_KEY:
return None
try:
async with httpx.AsyncClient(timeout=20) as client:
resp = await client.post(
AI_BASE,
headers={
"Authorization": f"Bearer {LLM_API_KEY}",
"Content-Type": "application/json",
},
json={
"model": AI_MODEL,
"messages": [{"role": "user", "content": prompt}],
"temperature": 0.3,
"max_tokens": max_tokens,
},
)
resp.raise_for_status()
return resp.json()["choices"][0]["message"]["content"].strip()
except Exception as e:
logger.warning(f"LLM call failed: {e}")
return None
# ======================================================================
# CRYPTO SYNONYM MAPPINGS (rule-based fallback for query expansion)
# ======================================================================
CRYPTO_SYNONYMS: dict[str, list[str]] = {
"rug pull": ["honeypot", "token scam", "liquidity drain", "exit scam", "pull liquidity"],
"wash trade": ["volume manipulation", "fake trading", "artificial volume", "self-trading"],
"honeypot": ["rug pull", "sell-disabled token", "trapped buyers", "cant-sell token"],
"pump and dump": ["pump dump", "price manipulation", "artificial inflation", "dump scheme"],
"flash loan": [
"instant loan",
"uncollateralized loan",
"flash loan attack",
"atomic arbitrage",
],
"sandwich attack": ["MEV attack", "front-running", "sandwich bot", "frontrun"],
"money laundering": ["mixer", "tornado cash", "coin mixing", "layered transactions"],
"sybil attack": ["fake identities", "multi-account", "sock puppet", "astroturfing"],
"dusting attack": ["dust attack", "micro-deposit", "privacy violation", "address poisoning"],
"reentrancy": ["recursive call", "reentrancy attack", "call-back vulnerability"],
"front running": ["MEV", "sandwich", "transaction ordering", "priority gas auction"],
"impersonation": ["phishing", "fake project", "clone scam", "counterfeit token"],
"exit scam": ["rug pull", "abandonment", "vanishing act", "team disappearance"],
"liquidity lock": ["locked liquidity", "LP lock", "liquidity freeze", "trustlock"],
"token drain": ["wallet drainer", "approval scam", "sweep tokens", "drain contract"],
"smart contract vulnerability": [
"code exploit",
"contract bug",
"security flaw",
"audit finding",
],
"whale manipulation": ["whale dump", "large holder sell", "whale influence", "big wallet move"],
"defi exploit": ["protocol hack", "vault drain", "yield farm exploit", "pool exploit"],
"bridge exploit": ["cross-chain attack", "bridge hack", "relayer vulnerability"],
"governance attack": [
"dao takeover",
"voting manipulation",
"governance exploit",
"flash loan governance",
],
}
# Token symbol pattern: $SYMBOL
_TOKEN_RE = re.compile(r"\$([A-Za-z]{2,10})")
# Category mappings for step-back
_TOKEN_CATEGORIES = {
# Major chains
"SOL": "Solana ecosystem token",
"ETH": "Ethereum ecosystem token",
"BTC": "Bitcoin ecosystem token",
"BNB": "Binance Smart Chain token",
"MATIC": "Polygon ecosystem token",
"AVAX": "Avalanche ecosystem token",
"ARB": "Arbitrum ecosystem token",
"OP": "Optimism ecosystem token",
"BASE": "Base ecosystem token",
}
# ======================================================================
# 1. HyDE -- HYPOTHETICAL DOCUMENT EMBEDDING
# ======================================================================
async def hyde_transform(query: str) -> str:
"""
Generate a hypothetical answer document for the query.
The hypothetical doc is what we'd *want* to retrieve.
Cache results in Redis (key: hyde:{hash}, TTL 24h).
"""
# Check Redis cache
query_hash = hashlib.sha256(query.encode()).hexdigest()[:32]
cache_key = f"hyde:{query_hash}"
try:
r = await _get_redis()
cached = await r.get(cache_key)
if cached:
logger.debug(f"HyDE cache hit: {query[:60]}")
return cached
except Exception:
pass
# Try LLM-based HyDE
llm_result = await _call_llm(
f"You are a cryptocurrency scam detection expert. Write a detailed, factual paragraph "
f"that would be the ideal answer to this question. Include specific details, indicators, "
f"and evidence that a real investigation report would contain.\n\n"
f"Question: {query}\n\n"
f"Ideal answer paragraph:",
max_tokens=300,
)
if llm_result:
result = llm_result
else:
# Rule-based fallback: prepend crypto-specific context
result = _hyde_rule_based(query)
# Cache in Redis (24h TTL)
try:
r = await _get_redis()
await r.setex(cache_key, 86400, result)
except Exception:
pass
return result
def _hyde_rule_based(query: str) -> str:
"""Rule-based HyDE: prepend crypto investigation context."""
query_lower = query.lower()
prefixes = []
if any(w in query_lower for w in ["rug", "pull", "drain", "liquidity"]):
prefixes.append("Cryptocurrency rug pull investigation:")
elif any(w in query_lower for w in ["honeypot", "can't sell", "sell disabled"]):
prefixes.append("Honeypot token analysis reveals:")
elif any(w in query_lower for w in ["wash", "volume", "fake trading"]):
prefixes.append("Wash trading detection report:")
elif any(w in query_lower for w in ["phishing", "drainer", "approval", "wallet drain"]):
prefixes.append("Wallet drainer and approval scam investigation:")
elif any(w in query_lower for w in ["mint", "unlimited", "supply", "dilution"]):
prefixes.append("Unlimited mint vulnerability detection:")
else:
prefixes.append("Cryptocurrency scam detection analysis:")
return f"{prefixes[0]} {query}"
# ======================================================================
# 2. QUERY EXPANSION
# ======================================================================
async def expand_query(query: str) -> list[str]:
"""
Generate 3-5 variant phrasings of the same query.
All variants will be searched and results merged with RRF.
"""
# Try LLM-based expansion
llm_result = await _call_llm(
f"Generate 3 to 5 alternative phrasings of this crypto fraud detection query. "
f"Each variant should use different terminology but seek the same information. "
f"Return ONLY a JSON array of strings, nothing else.\n\n"
f"Query: {query}\n\n"
f"Alternative phrasings (JSON array):",
max_tokens=256,
)
if llm_result:
try:
json_match = re.search(r"\[.*\]", llm_result, re.DOTALL)
if json_match:
parsed = json.loads(json_match.group(0))
if isinstance(parsed, list) and len(parsed) >= 2:
return [query] + [str(v) for v in parsed[:5] if v]
except Exception as e:
logger.warning(f"Failed to parse expansion LLM output: {e}")
# Rule-based fallback
return _expand_rule_based(query)
def _expand_rule_based(query: str) -> list[str]:
"""Rule-based query expansion using crypto synonym mappings."""
variants = [query]
query_lower = query.lower()
for key, synonyms in CRYPTO_SYNONYMS.items():
if key in query_lower:
added = 0
for syn in synonyms:
if syn.lower() not in query_lower and added < 3:
variant = re.sub(re.escape(key), syn, query, flags=re.IGNORECASE)
if variant != query and variant not in variants:
variants.append(variant)
added += 1
# Also add a broader version if query mentions token symbol
tokens = _TOKEN_RE.findall(query)
if tokens and len(variants) < 3:
symbol = tokens[0].upper()
category = _TOKEN_CATEGORIES.get(symbol, "cryptocurrency token")
broader = re.sub(r"\$" + re.escape(symbol), category, query, flags=re.IGNORECASE)
broader = re.sub(r"\?$", " indicators?", broader)
if broader != query and broader not in variants:
variants.append(broader)
# Deduplicate case-insensitively
seen = set()
unique = []
for v in variants:
key = v.lower().strip()
if key not in seen:
seen.add(key)
unique.append(v)
return unique[:6]
# ======================================================================
# 3. STEP-BACK PROMPTING
# ======================================================================
async def step_back_query(query: str) -> str:
"""
Generate a broader abstraction of the query.
e.g. "Is $SOL a rug pull?" -> "What are indicators of a rug pull on Solana?"
Retrieve for BOTH the original and the step-back query.
"""
llm_result = await _call_llm(
f"Given this specific cryptocurrency fraud detection query, generate a broader, "
f"more general version of the question that would retrieve relevant background "
f"knowledge. The broader question should focus on the CATEGORY or TYPE of issue "
f"rather than the specific instance.\n\n"
f"Specific query: {query}\n\n"
f"Broader general query:",
max_tokens=150,
)
if llm_result:
return llm_result.strip()
# Rule-based fallback
return _step_back_rule_based(query)
def _step_back_rule_based(query: str) -> str:
"""Rule-based step-back: replace specific tokens with broader categories."""
result = query
chain_used = None
# Replace $SYMBOL with "a token" and note the chain for context suffix
tokens = _TOKEN_RE.findall(query)
for symbol in tokens:
symbol_upper = symbol.upper()
_TOKEN_CATEGORIES.get(symbol_upper, "a cryptocurrency token")
result = result.replace(f"${symbol}", "a token")
result = result.replace(f"${symbol_upper}", "a token")
chain_used = symbol_upper
# Replace specific addresses with "an address"
addr_pattern = re.compile(r"0x[a-fA-F0-9]{40}|[1-9A-HJ-NP-Za-km-z]{32,44}")
if addr_pattern.search(result):
result = addr_pattern.sub("an address", result)
# Replace specific percentages with general terms
result = re.sub(r"\b\d+\.?\d*%\b", "significant percentage", result)
# Broader phrasing replacements — make it generic
broader_map = {
"is this": "what are indicators of",
"is it": "what are characteristics of",
"check if": "how to detect",
"detect if": "how to identify",
"verify if": "how to verify",
"is a token": "what are indicators of a token being",
"is an address": "what are indicators of an address being",
}
result_lower = result.lower()
for specific, broad in broader_map.items():
if specific in result_lower:
result = result_lower.replace(specific, broad, 1)
if result:
result = result[0].upper() + result[1:]
break
# Make it more question-like if it isn't already
if not result.endswith("?") and "?" not in result:
result = f"What are indicators of {result.lower()}?"
# If we replaced a token symbol, add chain context
if chain_used:
chain_names = {
"SOL": "Solana",
"ETH": "Ethereum",
"BTC": "Bitcoin",
"BNB": "Binance Smart Chain",
"MATIC": "Polygon",
"AVAX": "Avalanche",
"ARB": "Arbitrum",
"OP": "Optimism",
"BASE": "Base",
}
chain_name = chain_names.get(chain_used, chain_used)
if chain_name.lower() not in result.lower():
result = result.rstrip("?") + f" on {chain_name}?"
return result
# ======================================================================
# 4. ROUTER -- auto-select the best strategy
# ======================================================================
async def transform_query(query: str, strategy: str = "auto") -> TransformedQuery:
"""
Transform a query using the specified strategy.
strategy:
"auto" -- pick the best strategy based on query characteristics
"hyde" -- Hypothetical Document Embedding
"expand" -- Query Expansion (multiple variant searches)
"step_back" -- Step-Back Prompting (broader abstraction)
"none" -- no transformation, use query as-is
Returns TransformedQuery with the list of strings to embed and search.
"""
if strategy == "none":
return TransformedQuery(
original_query=query,
strategy="none",
transformed_queries=[query],
metadata={"reason": "no transformation requested"},
)
if strategy == "auto":
strategy = _auto_select_strategy(query)
if strategy == "hyde":
hypo_doc = await hyde_transform(query)
return TransformedQuery(
original_query=query,
strategy="hyde",
transformed_queries=[hypo_doc],
metadata={"hypothetical_doc": hypo_doc[:200]},
)
elif strategy == "expand":
variants = await expand_query(query)
return TransformedQuery(
original_query=query,
strategy="expand",
transformed_queries=variants,
metadata={"variant_count": len(variants)},
)
elif strategy == "step_back":
stepped = await step_back_query(query)
return TransformedQuery(
original_query=query,
strategy="step_back",
transformed_queries=[query, stepped],
metadata={"step_back_query": stepped},
)
# Fallback: no transformation
return TransformedQuery(
original_query=query,
strategy="none",
transformed_queries=[query],
metadata={"reason": f"unknown strategy: {strategy}"},
)
def _auto_select_strategy(query: str) -> str:
"""
Auto-select the best query transformation strategy based on query characteristics.
Heuristics:
- Short factual queries (< 6 words, no question marks) -> "none"
- Ambiguous queries (broad terms, generic questions) -> "hyde"
- Specific entity queries (addresses, token symbols) -> "expand"
- Broad/exploratory queries (how/what/why + topic) -> "step_back"
"""
q = query.strip()
word_count = len(q.split())
q_lower = q.lower()
# Check for specific entities
has_address = bool(re.search(r"0x[a-fA-F0-9]{20,}|[1-9A-HJ-NP-Za-km-z]{32,44}", q))
has_token_symbol = bool(_TOKEN_RE.search(q))
# Check for question patterns
starts_with_question_word = q_lower.startswith(("what", "how", "why", "is", "are", "can", "does", "do"))
has_question_mark = "?" in q
# Check for ambiguity
broad_indicators = [
"scam",
"fraud",
"suspicious",
"safe",
"legit",
"trustworthy",
"risk",
"dangerous",
"honest",
"real",
]
is_ambiguous = any(w in q_lower for w in broad_indicators) and not has_address and not has_token_symbol
# Decision logic
if has_address or has_token_symbol:
return "expand"
if is_ambiguous:
return "hyde"
if word_count < 6 and not has_question_mark and not starts_with_question_word:
return "none"
if starts_with_question_word or has_question_mark:
return "step_back"
return "none"