add RAG

rag/ingest.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+import uuid, pathlib, logging
+from typing import List, Dict, Any
+from pypdf import PdfReader
+import trafilatura
+from .utils import Doc, normalize_text
+
+# Silence noisy pypdf warnings from malformed PDFs
+logging.getLogger("pypdf").setLevel(logging.ERROR)
+
+def read_txt(path: str) -> str:
+    with open(path, "r", encoding="utf-8", errors="ignore") as f:
+        return f.read()
+
+def read_pdf(path: str) -> str:
+    text = []
+    reader = PdfReader(path)
+    for page in reader.pages:
+        text.append(page.extract_text() or "")
+    return "\n".join(text)
+
+def read_any(path: str) -> str:
+    ext = pathlib.Path(path).suffix.lower()
+    if ext in [".txt", ".md"]:
+        return read_txt(path)
+    elif ext in [".pdf"]:
+        return read_pdf(path)
+    else:
+        return read_txt(path)
+
+def fetch_url(url: str) -> str:
+    downloaded = trafilatura.fetch_url(url)
+    if not downloaded:
+        return ""
+    return trafilatura.extract(downloaded) or ""
+
+def split_to_chunks(text: str, chunk_size: int = 800, overlap: int = 100) -> List[str]:
+    words = text.split()
+    if not words:
+        return []
+    chunks = []
+    i = 0
+    step = max(1, chunk_size - overlap)
+    while i < len(words):
+        chunk = " ".join(words[i:i+chunk_size])
+        chunks.append(chunk)
+        i += step
+    return chunks or [text]
+
+def guess_coin(label: str) -> str:
+    low = label.lower()
+    if "bitcoin" in low or "btc" in low: return "bitcoin"
+    if "ethereum" in low or "eth" in low: return "ethereum"
+    return ""
+
+def build_docs_from_paths(paths: List[str], source_label: str = "local") -> List[Doc]:
+    docs: List[Doc] = []
+    for p in paths or []:
+        raw = read_any(p)
+        if not raw: 
+            continue
+        coin = guess_coin(p)
+        for i, chunk in enumerate(split_to_chunks(raw)):
+            docs.append(Doc(
+                id=f"{uuid.uuid4()}",
+                text=normalize_text(chunk),
+                metadata={"source": source_label, "path": p, "chunk": i, "coin": coin}
+            ))
+    return docs
+
+def build_docs_from_urls(urls: List[str], source_label: str = "web") -> List[Doc]:
+    docs: List[Doc] = []
+    for u in urls or []:
+        raw = fetch_url(u)
+        if not raw: 
+            continue
+        coin = guess_coin(u)
+        for i, chunk in enumerate(split_to_chunks(raw)):
+            docs.append(Doc(
+                id=f"{uuid.uuid4()}",
+                text=normalize_text(chunk),
+                metadata={"source": source_label, "url": u, "chunk": i, "coin": coin}
+            ))
+    return docs

rag/pipeline.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+from typing import List, Dict, Any
+
+from openai import OpenAI
+from .utils import HybridIndex, Reranker, Doc, select_fewshots
+from .ingest import build_docs_from_paths, build_docs_from_urls
+from prompts import SYSTEM_PROMPT, FEWSHOTS
+
+class CryptoRAGPipeline:
+    def __init__(self, dense_model: str = "sentence-transformers/all-MiniLM-L6-v2", reranker_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
+        self.index = HybridIndex(dense_model_name=dense_model)
+        self.reranker = Reranker(reranker_model)
+        self.client: OpenAI | None = None
+
+    def set_openai(self, api_key: str):
+        self.client = OpenAI(api_key=api_key)
+
+    def add_local_files(self, paths: List[str]):
+        docs = build_docs_from_paths(paths, source_label="local")
+        self.index.add(docs)
+
+    def add_urls(self, urls: List[str]):
+        docs = build_docs_from_urls(urls, source_label="web")
+        self.index.add(docs)
+
+    def build(self):
+        self.index.build()
+
+    def route(self, query: str) -> str:
+        q = query.lower()
+        if any(k in q for k in ["price", "market cap", "marketcap", "ath", "all-time high", "24h", "fear greed", "greed index"]):
+            return "tools"
+        return "rag"
+
+    def build_prompt(self, query: str, contexts: List[Doc]) -> str:
+        fs = select_fewshots(query, FEWSHOTS, self.index.embedder, n=2)
+        few = "\n\n".join([f"Q: {x['q']}\nA: {x['a']}" for x in fs])
+        ctx = "\n\n".join([f"[{i+1}] {c.text[:1200]}" for i, c in enumerate(contexts)])
+        prompt = f"""{SYSTEM_PROMPT}
+
+Few-shot examples:
+{few}
+
+Context (use to answer if relevant; cite [#]):
+{ctx}
+
+User question: {query}
+
+Answer:"""
+        return prompt
+
+    def answer_stream(self, query: str, contexts: List[Doc], model: str = "gpt-4o-mini"):
+        assert self.client is not None, "LLM client not set"
+        prompt = self.build_prompt(query, contexts)
+        with self.client.chat.completions.create(
+            model=model,
+            messages=[{"role":"system","content":SYSTEM_PROMPT},
+                      {"role":"user","content":prompt}],
+            stream=True,
+            temperature=0.3,
+            max_tokens=400
+        ) as stream:
+            for event in stream:
+                if hasattr(event, "choices") and event.choices:
+                    delta = event.choices[0].delta
+                    if delta and delta.content:
+                        yield delta.content
+
+    def ask(self, query: str, k: int = 8, alpha: float = 0.5, top_k_rerank: int = 5, filters: Dict[str, Any] | None = None, stream: bool = True):
+        route = self.route(query)
+        if route == "tools":
+            return {"route": "tools", "contexts": []}
+
+        # Try auto-build if needed
+        if not self.index.ready():
+            self.index.build()
+            if not self.index.ready():
+                return {"route": "not_ready", "contexts": [], "reason": "index_empty" if len(self.index.docs)==0 else "build_failed"}
+
+        hits = self.index.search(query, k=k, alpha=alpha, filters=filters)
+        if not hits:
+            return {"route": "not_ready", "contexts": [], "reason": "no_results"}
+
+        reranked = self.reranker.rerank(query, hits, top_k=top_k_rerank)
+        top_contexts = [d for d,_ in reranked]
+        return {"route": "rag", "contexts": top_contexts}

rag/tools.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+import requests
+
+# Minimal map from common names/symbols → CoinGecko IDs
+COIN_MAP = {
+    "btc": "bitcoin", "bitcoin": "bitcoin",
+    "eth": "ethereum", "ethereum": "ethereum",
+    "sol": "solana", "solana": "solana",
+    "xrp": "ripple", "ripple": "ripple",
+}
+
+def resolve_coin_id(text_or_symbol: str, default: str = "bitcoin") -> str:
+    t = (text_or_symbol or "").lower().strip()
+    # try exact-in-text matches first (longest keys first)
+    for key in sorted(COIN_MAP.keys(), key=len, reverse=True):
+        if key in t.split() or key in t:
+            return COIN_MAP[key]
+    return default
+
+def get_price(coin_id: str = "bitcoin", vs: str = "usd"):
+    url = f"https://api.coingecko.com/api/v3/simple/price?ids={coin_id}&vs_currencies={vs}"
+    r = requests.get(url, timeout=10)
+    r.raise_for_status()
+    data = r.json()
+    return data.get(coin_id, {}).get(vs)
+
+def get_price_any(coin_or_query: str, vs: str = "usd"):
+    coin_id = resolve_coin_id(coin_or_query)
+    return coin_id, get_price(coin_id, vs)
+
+def get_price_multi(coin_ids: list[str], vs: str = "usd") -> dict:
+    # Efficient batch call (one request) e.g. ["bitcoin","ethereum","solana","ripple"]
+    unique = ",".join(sorted(set(coin_ids)))
+    url = f"https://api.coingecko.com/api/v3/simple/price?ids={unique}&vs_currencies={vs}"
+    r = requests.get(url, timeout=10)
+    r.raise_for_status()
+    return r.json()
+
+def get_fear_greed():
+    url = "https://api.alternative.me/fng/"
+    r = requests.get(url, timeout=10)
+    r.raise_for_status()
+    data = r.json()
+    if data.get("data"):
+        return data["data"][0]
+    return None

rag/utils.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+import re, json, hashlib
+from dataclasses import dataclass
+from typing import List, Dict, Any, Tuple
+
+import numpy as np
+from rank_bm25 import BM25Okapi
+from sentence_transformers import SentenceTransformer, CrossEncoder
+
+def cache_key(obj: Any) -> str:
+    return hashlib.sha256(json.dumps(obj, sort_keys=True).encode()).hexdigest()
+
+def normalize_text(s: str) -> str:
+    return re.sub(r"\s+", " ", s).strip()
+
+@dataclass
+class Doc:
+    id: str
+    text: str
+    metadata: Dict[str, Any]
+
+class HybridIndex:
+    def __init__(self, dense_model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
+        self.dense_model_name = dense_model_name
+        self.embedder = SentenceTransformer(dense_model_name)
+        self.docs: List[Doc] = []
+        self.bm25 = None
+        self.embeddings = None
+
+    def add(self, docs: List[Doc]):
+        self.docs.extend(docs)
+
+    def build(self):
+        # Build only if we have docs
+        if not self.docs:
+            self.bm25, self.embeddings = None, None
+            return
+        corpus = [d.text for d in self.docs]
+        tokenized = [c.split() for c in corpus]
+        self.bm25 = BM25Okapi(tokenized)
+        self.embeddings = self.embedder.encode(
+            corpus, convert_to_numpy=True, normalize_embeddings=True
+        )
+
+    def ready(self) -> bool:
+        return (self.bm25 is not None) and (self.embeddings is not None) and (len(self.docs) > 0)
+
+    def search(self, query: str, k: int = 8, alpha: float = 0.5, filters: Dict[str, Any] | None = None):
+        # If index isn't ready, return empty (UI/pipeline should guide the user)
+        if not self.ready():
+            return []
+
+        # Dense embedding for query
+        query_vec = self.embedder.encode([query], convert_to_numpy=True, normalize_embeddings=True)[0]
+
+        # BM25 + dense scores
+        q_tokens = query.split()
+        try:
+            bm25_scores = self.bm25.get_scores(q_tokens)
+        except Exception:
+            # Fallback if BM25 hiccups (e.g., empty tokens)
+            bm25_scores = np.zeros(len(self.docs), dtype=float)
+        dense_scores = (self.embeddings @ query_vec)
+
+        # NumPy 2.x-safe normalization
+        def _norm(x: np.ndarray) -> np.ndarray:
+            x = np.asarray(x, dtype=float)
+            rng = np.ptp(x)
+            return (x - x.min()) / (rng + 1e-8)
+
+        bm25_norm = _norm(bm25_scores)
+        dense_norm = _norm(dense_scores)
+        scores = alpha * bm25_norm + (1 - alpha) * dense_norm
+
+        # Optional metadata filters
+        idxs = np.arange(len(self.docs))
+        if filters:
+            def ok(d: Doc) -> bool:
+                for kf, vf in filters.items():
+                    if kf not in d.metadata:
+                        return False
+                    dv = str(d.metadata[kf]).lower()
+                    if isinstance(vf, (list, tuple, set)):
+                        if not any(str(x).lower() in dv for x in vf):
+                            return False
+                    else:
+                        if str(vf).lower() not in dv:
+                            return False
+                return True
+
+            keep = [i for i in idxs if ok(self.docs[int(i)])]
+            if not keep:
+                return []
+            idxs = np.array(keep, dtype=int)
+            scores = scores[idxs]
+
+        # Top-k results
+        order = np.argsort(-scores)[:k]
+        return [(self.docs[int(idxs[i])], float(scores[i])) for i in order]
+
+class Reranker:
+    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
+        self.model = CrossEncoder(model_name)
+
+    def rerank(self, query: str, docs: List[Tuple[Doc, float]], top_k: int = 5) -> List[Tuple[Doc, float]]:
+        if not docs:
+            return []
+        pairs = [(query, d.text) for d, _ in docs]
+        scores = self.model.predict(pairs)
+        rescored = list(zip([d for d,_ in docs], [float(s) for s in scores]))
+        rescored.sort(key=lambda x: -x[1])
+        return rescored[:top_k]
+
+def select_fewshots(query: str, fewshots: List[Dict[str, str]], embedder: SentenceTransformer, n: int = 2):
+    if not fewshots:
+        return []
+    qv = embedder.encode([query], convert_to_numpy=True, normalize_embeddings=True)[0]
+    ex_vecs = embedder.encode([fs["q"] for fs in fewshots], convert_to_numpy=True, normalize_embeddings=True)
+    sims = ex_vecs @ qv
+    order = np.argsort(-sims)[:n]
+    return [fewshots[i] for i in order]