511 lines
18 KiB
Python
511 lines
18 KiB
Python
"""
|
|
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"
|