""" 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: # noqa: SIM108 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"