rag_pipeline.py

"""
JAIM - RAG Pipeline (v3)
Core Retrieval-Augmented Generation engine with:
- BAAI/bge-large-en-v1.5 embeddings (1024 dims)
- Query Expansion via Groq (3 rephrasings + original = 4 queries)
- Deduplication by parent entry ID → single best match to LLM
- Forward + Backward chaining inference engine

Uses `requests` for HTTP calls for Python 3.14 compatibility.
"""

import os
import json
import time
import requests
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer
from inference.engine import VrikshayurvedaInferenceEngine

load_dotenv()

# ─── Config ──────────────────────────────────────────────────────────────────
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "jaim3")
EMBEDDING_MODEL = "BAAI/bge-large-en-v1.5"
GROQ_MODEL = "llama-3.3-70b-versatile"
# BGE spec: queries must be prefixed with this instruction
QUERY_PREFIX = "Represent this sentence for searching relevant passages: "
TOP_K = 3


class JAIMPipeline:
    """
    JAIM RAG Pipeline v3:
    1. Expands user query into 4 variants via Groq
    2. Encodes all variants using BGE-large-en-v1.5 (with query prefix)
    3. Retrieves top-K from Pinecone for each variant
    4. Deduplicates by parent_id, keeps highest score per entry
    5. Runs forward + backward chaining inference engine
    6. Sends the best match's context + inference grounding to Groq LLM
    """

    def __init__(self):
        print("🔄 Initializing JAIM Pipeline...")

        print("   📦 Loading BGE-large-en-v1.5 embedding model...")
        self.embed_model = SentenceTransformer(EMBEDDING_MODEL)

        print("   🌲 Connecting to Pinecone...")
        self.pinecone_host = self._get_pinecone_host()
        print(f"      Index host: {self.pinecone_host}")

        print("   🤖 Groq (Llama 3.3 70B) configured")
        print("   🧠 Initializing forward/backward chaining inference engine...")
        self.inference_engine = VrikshayurvedaInferenceEngine()
        print("✅ JAIM Pipeline ready!\n")

    # ─── Pinecone Host Discovery ──────────────────────────────────────────────
    def _get_pinecone_host(self) -> str:
        """Get the Pinecone index host URL via the control plane API."""
        url = f"https://api.pinecone.io/indexes/{INDEX_NAME}"
        headers = {"Api-Key": PINECONE_API_KEY}
        resp = requests.get(url, headers=headers, timeout=15)
        resp.raise_for_status()
        return resp.json()["host"]

    # ─── Query Expansion via Groq ─────────────────────────────────────────────
    def expand_query(self, query: str) -> list[str]:
        """
        Use Groq to generate 3 rephrasings of the user query.
        Returns [original, rephrasing1, rephrasing2, rephrasing3].
        """
        url = "https://api.groq.com/openai/v1/chat/completions"
        headers = {
            "Authorization": f"Bearer {GROQ_API_KEY}",
            "Content-Type": "application/json",
        }
        prompt = f"""Generate exactly 3 different rephrasings of the following plant symptom query.
Each rephrasing should use different vocabulary while preserving the meaning.
Return ONLY a valid JSON array of 3 strings. No extra text, no markdown.

Original query: "{query}"
"""
        payload = {
            "model": GROQ_MODEL,
            "messages": [
                {"role": "system", "content": "You rephrase queries. Return only valid JSON arrays."},
                {"role": "user", "content": prompt},
            ],
            "temperature": 0.7,
            "max_tokens": 256,
        }

        for attempt in range(3):
            resp = requests.post(url, json=payload, headers=headers, timeout=15)
            if resp.status_code == 429:
                time.sleep(2 ** (attempt + 1))
                continue
            resp.raise_for_status()
            break
        else:
            print("   ⚠️ Query expansion failed (rate limited), using original only")
            return [query]

        content = resp.json()["choices"][0]["message"]["content"].strip()
        if content.startswith("```"):
            content = content.split("\n", 1)[1].rsplit("```", 1)[0].strip()

        try:
            rephrasings = json.loads(content)
            return [query] + rephrasings[:3]
        except json.JSONDecodeError:
            print("   ⚠️ Query expansion parse error, using original only")
            return [query]

    # ─── Pinecone Retrieval ───────────────────────────────────────────────────
    def retrieve_single(self, query_embedding: list, top_k: int = TOP_K) -> list[dict]:
        """Query Pinecone with a single embedding vector."""
        url = f"https://{self.pinecone_host}/query"
        headers = {
            "Api-Key": PINECONE_API_KEY,
            "Content-Type": "application/json",
        }
        payload = {
            "vector": query_embedding,
            "topK": top_k,
            "includeMetadata": True,
        }

        resp = requests.post(url, json=payload, headers=headers, timeout=15)
        resp.raise_for_status()

        matches = []
        for match in resp.json().get("matches", []):
            matches.append({
                "id": match["id"],
                "score": round(match["score"], 4),
                "metadata": match.get("metadata", {})
            })
        return matches

    def retrieve_with_variants(self, variants: list[str], top_k: int = TOP_K) -> list[dict]:
        """
        Retrieval pipeline using pre-expanded query variants:
        1. Embed each variant with BGE-large (query prefix applied)
        2. Retrieve top-K from Pinecone for each variant
        3. Deduplicate by parent_id, keep highest score per entry
        """
        print(f"   🔄 Retrieving with {len(variants)} query variants")

        all_matches = []
        for variant in variants:
            prefixed = f"{QUERY_PREFIX}{variant}"
            embedding = self.embed_model.encode(prefixed).tolist()
            matches = self.retrieve_single(embedding, top_k)
            all_matches.extend(matches)

        # Deduplicate by parent_id, keep highest score per entry
        best_by_parent = {}
        for match in all_matches:
            parent = match["metadata"].get("parent_id", match["id"])
            if parent not in best_by_parent or match["score"] > best_by_parent[parent]["score"]:
                best_by_parent[parent] = match

        # Sort by score descending
        deduped = sorted(best_by_parent.values(), key=lambda x: x["score"], reverse=True)
        return deduped

    # ─── Context Building ─────────────────────────────────────────────────────
    def build_context(self, matches: list[dict]) -> str:
        """Build context from the single highest-scoring deduplicated match."""
        if not matches:
            return "No relevant entries found in the Vrikshayurveda database."

        # Use only the single highest-scoring match
        match = matches[0]
        meta = match["metadata"]
        possible_causes = meta.get("possible_causes", "N/A")
        if isinstance(possible_causes, list):
            possible_causes = ", ".join(possible_causes)

        return (
            f"--- Best Match (Relevance: {match['score']}) ---\n"
            f"Disorder Type: {meta.get('disorder', 'N/A')}\n"
            f"Cause (Dosha): {meta.get('cause_given', 'N/A')}\n"
            f"Symptoms: {meta.get('symptoms', 'N/A')}\n"
            f"Cause Elaborated: {meta.get('cause_elaborated', 'N/A')}\n"
            f"Possible Causes: {possible_causes}\n"
            f"Ayurvedic Treatment: {meta.get('treatment_material', 'N/A')}\n"
        )

    # ─── LLM Response Generation ──────────────────────────────────────────────
    def _format_inference_context(self, inference_result) -> str:
        """Convert chaining inference output into compact prompt context."""
        if inference_result is None:
            return "No deterministic diagnosis available."
        if hasattr(inference_result, 'llm_context'):
            return inference_result.llm_context
        return "No inference context available."

    def generate_response(self, query: str, context: str, inference_result=None) -> str:
        """Use Groq REST API (Llama 3.3 70B) to generate a response."""
        system_prompt = """You are JAIM (Jnana AI for Marga — Knowledge AI for the Path), an expert Ayurvedic plant care assistant powered by the ancient Vrikshayurveda (Science of Plant Life).

You help farmers, gardeners, and plant enthusiasts diagnose plant disorders and recommend traditional Ayurvedic treatments based on the Vrikshayurveda text by Surapala.

IMPORTANT RULES:
1. Base your response PRIMARILY on the retrieved Vrikshayurveda knowledge provided.
2. The Inference Engine Analysis provides a formal reasoning chain — use it to validate and strengthen your response.
3. Provide the specific Ayurvedic treatment from the retrieved data based on the matching symptoms.
4. You may add brief modern context to help the user understand, but always prioritize the traditional treatment.
5. Be clear, structured, and helpful.
6. If the query doesn't match any known disorder well, say so honestly."""

        llm_grounding = self._format_inference_context(inference_result)

        user_message = f"""═══════════════════════════════════════
[Inference Engine Analysis]
═══════════════════════════════════════
{llm_grounding}

═══════════════════════════════════════
[Retrieved Vrikshayurveda Passages]
═══════════════════════════════════════
{context}

═══════════════════════════════════════
[User Question]
═══════════════════════════════════════
{query}

Please provide a comprehensive, well-structured response:"""

        url = "https://api.groq.com/openai/v1/chat/completions"
        headers = {
            "Authorization": f"Bearer {GROQ_API_KEY}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": GROQ_MODEL,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message},
            ],
            "temperature": 0.7,
            "max_tokens": 2048,
        }

        # Retry with backoff for rate limits
        for attempt in range(3):
            resp = requests.post(url, json=payload, headers=headers, timeout=30)
            if resp.status_code == 429:
                wait_time = 2 ** (attempt + 1)
                print(f"   ⏳ Rate limited, retrying in {wait_time}s...")
                time.sleep(wait_time)
                continue
            resp.raise_for_status()
            break
        else:
            return "⚠️ Rate limit exceeded. Please try again in a minute."

        result = resp.json()

        try:
            return result["choices"][0]["message"]["content"]
        except (KeyError, IndexError):
            return f"Error parsing response: {json.dumps(result, indent=2)}"

    # ─── Full RAG Pipeline ────────────────────────────────────────────────────
    def query(self, user_query: str, top_k: int = TOP_K) -> dict:
        """
        Full RAG pipeline:
        expand → retrieve → infer (chaining) → build context → generate response.
        """
        # Step 1: Expand query into variants
        expanded_queries = self.expand_query(user_query)
        print(f"   🔄 Query expanded into {len(expanded_queries)} variants")

        # Step 2: Retrieve with pre-expanded variants
        matches = self.retrieve_with_variants(expanded_queries, top_k)

        # Step 3: Extract symptom text from retrieved matches for inference
        retrieved_chunks = [
            m["metadata"].get("symptoms", "")
            for m in matches
            if m.get("metadata", {}).get("symptoms")
        ]

        # Step 4: Run forward + backward chaining inference engine
        diagnosis_result = self.inference_engine.diagnose(
            user_query=user_query,
            expanded_queries=expanded_queries,
            retrieved_chunks=retrieved_chunks,
        )

        # Step 5: Build context and generate LLM response
        context = self.build_context(matches)
        response = self.generate_response(
            user_query, context, inference_result=diagnosis_result
        )

        return {
            "query": user_query,
            "response": response,
            "diagnosis": diagnosis_result,
            "sources": matches[:top_k],
            "num_sources": len(matches)
        }


# ─── Quick Test ───────────────────────────────────────────────────────────────
if __name__ == "__main__":
    pipeline = JAIMPipeline()

    test_query = "My tree trunk is bent and the fruits are hard and not juicy"
    print(f"🔍 Query: {test_query}\n")

    result = pipeline.query(test_query)

    print("=" * 60)
    print("📋 JAIM RESPONSE:")
    print("=" * 60)
    print(result["response"])
    print("\n" + "=" * 60)
    print(f"📊 Sources used: {result['num_sources']}")
    for src in result["sources"]:
        print(f"   • {src['id']} (score: {src['score']})")