JAIM_INFERENCE_ENGINE.md

# JAIM Inference Engine — Build & Integration Guide

> **How to use this file:** Read the overview, then follow each phase in order.
> Every phase has a goal, exact file targets, and copy-paste-ready instructions.
> Do not skip phases — each one depends on the previous.

---

## Project context

JAIM (Jnana AI for Marga) is an Ayurvedic plant care RAG assistant powered by
Vrikshayurveda by Surapala. The existing symbolic rule engine (keyword overlap +
phrase indicators) must be **fully replaced** by a forward + backward chaining
inference engine. A visual Symptom Chain Explorer tab is added to the Gradio UI.

**Existing stack (do not change):**
- `app.py` — Gradio 4.x UI + pipeline orchestration
- LLM: Llama 3.3 70B via Groq API
- Embeddings: BAAI/bge-large-en-v1.5
- Vector DB: Pinecone (cosine similarity, deduplicated by `parent_id`)
- Query expansion: 4 variants via Groq

---

## Folder structure (target state after all phases)

```
project_root/
├── app.py                          ← modified (3 surgical changes only)
├── vrikshayurveda_rules.yaml       ← new (Phase 1)
├── inference/
│   ├── __init__.py                 ← new empty file (Phase 2)
│   ├── symptom_extractor.py        ← new (Phase 2)
│   ├── forward_chain.py            ← new (Phase 3)
│   ├── backward_chain.py           ← new (Phase 4)
│   ├── engine.py                   ← new (Phase 5)
│   └── smoke_test.py               ← new (Phase 6)
```

---

## Phase 1 — Rule base

**Goal:** Create `vrikshayurveda_rules.yaml` in the project root.

**File:** `vrikshayurveda_rules.yaml`

Each rule follows this exact schema:

```yaml
rules:
  - id: string          # e.g. R001
    name: string        # short human-readable label
    source: string      # "vrikshayurveda" or "derived"
    priority: int       # higher fires first during conflict resolution
    conditions: list    # ALL must be true simultaneously (AND logic)
    conclusions: list   # facts asserted when rule fires
    dosha: string       # vata | kapha | pitta | none
    explanation: string # one sentence Ayurvedic justification
```

**Populate with exactly these 20 rules:**

```yaml
rules:

  - id: R001
    name: "Yellowing + stunted growth → Vata disorder"
    source: vrikshayurveda
    priority: 10
    conditions: [yellowing, stunted_growth]
    conclusions: [vata_disorder]
    dosha: vata
    explanation: >
      Yellowing with stunted growth indicates disturbed Vata,
      which governs growth and upward movement in plants.

  - id: R002
    name: "Wilting alone → Pitta imbalance"
    source: vrikshayurveda
    priority: 9
    conditions: [wilting]
    conclusions: [pitta_imbalance]
    dosha: pitta
    explanation: >
      Sudden wilting without moisture deficit signals excess Pitta heat
      disrupting the plant's fluid regulation.

  - id: R003
    name: "Bark lesions → Kapha obstruction"
    source: vrikshayurveda
    priority: 9
    conditions: [bark_lesions]
    conclusions: [kapha_obstruction]
    dosha: kapha
    explanation: >
      Lesions on bark indicate Kapha blocking the plant's
      nutrient channels (srotas).

  - id: R004
    name: "Root rot → root-based cause"
    source: vrikshayurveda
    priority: 9
    conditions: [root_rot]
    conclusions: [root_based_cause]
    dosha: vata
    explanation: >
      Physical root rot confirms the disorder originates
      in the root system.

  - id: R005
    name: "Vata disorder + monsoon → root cause + moisture excess"
    source: vrikshayurveda
    priority: 10
    conditions: [vata_disorder, monsoon]
    conclusions: [root_based_cause, moisture_excess]
    dosha: vata
    explanation: >
      Vata disorders worsened by monsoon rain point to a
      waterlogged root environment.

  - id: R006
    name: "Pitta imbalance + summer → heat stress + solar damage"
    source: vrikshayurveda
    priority: 10
    conditions: [pitta_imbalance, summer]
    conclusions: [heat_stress, solar_damage]
    dosha: pitta
    explanation: >
      Pitta disorders peak in summer — solar radiation amplifies
      the heat imbalance in plant tissue.

  - id: R007
    name: "Kapha obstruction + monsoon → fungal Kapha disorder"
    source: vrikshayurveda
    priority: 8
    conditions: [kapha_obstruction, monsoon]
    conclusions: [fungal_kapha_disorder, moisture_excess]
    dosha: kapha
    explanation: >
      Kapha obstruction in monsoon season indicates fungal
      infiltration of the srotas.

  - id: R008
    name: "Root-based cause → recommend drainage"
    source: derived
    priority: 7
    conditions: [root_based_cause]
    conclusions: [recommend_drainage]
    dosha: none
    explanation: >
      Root-based causes require improving soil drainage
      to eliminate waterlogging.

  - id: R009
    name: "Moisture excess → recommend neem bark"
    source: derived
    priority: 6
    conditions: [moisture_excess]
    conclusions: [recommend_neem_bark]
    dosha: none
    explanation: >
      Neem bark decoction applied to soil counteracts fungal
      growth caused by excess moisture.

  - id: R010
    name: "Heat stress → recommend chandana paste"
    source: derived
    priority: 7
    conditions: [heat_stress]
    conclusions: [recommend_chandana_paste]
    dosha: none
    explanation: >
      Chandana (sandalwood) paste on bark cools Pitta excess
      and prevents further solar damage.

  - id: R011
    name: "Solar damage → recommend shade cloth"
    source: derived
    priority: 5
    conditions: [solar_damage]
    conclusions: [recommend_shade_cloth]
    dosha: none
    explanation: >
      Physical shading reduces direct radiation on
      heat-damaged leaves and bark.

  - id: R012
    name: "Kapha obstruction → recommend bark scraping"
    source: derived
    priority: 6
    conditions: [kapha_obstruction]
    conclusions: [recommend_bark_scraping]
    dosha: none
    explanation: >
      Scraping obstructed bark restores sap flow and
      disperses Kapha blockage.

  - id: R013
    name: "Vata disorder + root rot → fungal Vata disorder"
    source: vrikshayurveda
    priority: 11
    conditions: [vata_disorder, root_rot]
    conclusions: [fungal_vata_disorder]
    dosha: vata
    explanation: >
      Combined Vata imbalance with physical root rot confirms
      fungal origin of the Vata disorder.

  - id: R014
    name: "Fungal Vata disorder → drainage + neem bark"
    source: derived
    priority: 6
    conditions: [fungal_vata_disorder]
    conclusions: [recommend_drainage, recommend_neem_bark]
    dosha: none
    explanation: >
      Fungal Vata disorders require both drainage improvement
      and antifungal neem treatment.

  - id: R015
    name: "Yellowing + wilting → Vata-Pitta combined"
    source: vrikshayurveda
    priority: 8
    conditions: [yellowing, wilting]
    conclusions: [vata_pitta_combined]
    dosha: vata
    explanation: >
      Yellowing with wilting suggests combined Vata-Pitta imbalance
      affecting both growth and heat regulation.

  - id: R016
    name: "Vata-Pitta combined → triphala water + shade"
    source: derived
    priority: 7
    conditions: [vata_pitta_combined]
    conclusions: [recommend_triphala_water, recommend_shade_cloth]
    dosha: none
    explanation: >
      Combined Vata-Pitta disorders respond to Triphala water
      irrigation and physical shade protection.

  - id: R017
    name: "Bark lesions + monsoon → fungal Kapha disorder"
    source: vrikshayurveda
    priority: 9
    conditions: [bark_lesions, monsoon]
    conclusions: [fungal_kapha_disorder]
    dosha: kapha
    explanation: >
      Bark lesions in monsoon season are a direct indicator
      of Kapha-type fungal invasion.

  - id: R018
    name: "Fungal Kapha disorder → bark scraping + neem bark"
    source: derived
    priority: 6
    conditions: [fungal_kapha_disorder]
    conclusions: [recommend_bark_scraping, recommend_neem_bark]
    dosha: none
    explanation: >
      Fungal Kapha disorders require mechanical scraping followed
      by neem antifungal treatment.

  - id: R019
    name: "Wilting + summer + root rot → severe Pitta disorder"
    source: vrikshayurveda
    priority: 8
    conditions: [wilting, summer, root_rot]
    conclusions: [severe_pitta_disorder]
    dosha: pitta
    explanation: >
      Wilting during summer with root rot signals severe Pitta disorder
      — the root system is heat-compromised.

  - id: R020
    name: "Severe Pitta disorder → chandana + drainage + vetiver"
    source: derived
    priority: 9
    conditions: [severe_pitta_disorder]
    conclusions: [recommend_chandana_paste, recommend_drainage, recommend_vetiver_root_soak]
    dosha: none
    explanation: >
      Severe Pitta disorders require multi-pronged treatment: cooling
      paste, drainage improvement, and vetiver root soak.
```

---

## Phase 2 — Symptom extractor

**Goal:** Create `inference/symptom_extractor.py` and an empty `inference/__init__.py`.

**File:** `inference/__init__.py`
Leave this file empty. Its only purpose is to make `inference` a Python package.

---

**File:** `inference/symptom_extractor.py`

```python
from __future__ import annotations

SYMPTOM_KEYWORDS: dict[str, list[str]] = {
    "yellowing": [
        "yellow", "yellowing", "pale", "chlorosis",
        "discolored leaves", "fading", "light green"
    ],
    "stunted_growth": [
        "stunted", "slow growth", "not growing", "small",
        "dwarfed", "no new growth", "growth stopped"
    ],
    "wilting": [
        "wilt", "wilting", "drooping", "limp",
        "sagging", "droopy", "collapsed"
    ],
    "bark_lesions": [
        "lesion", "crack", "canker", "sore",
        "wound on bark", "bark damage", "bark cracking",
        "oozing bark", "sunken spots"
    ],
    "root_rot": [
        "root rot", "rotting roots", "black roots",
        "mushy roots", "decaying roots", "smelly roots"
    ],
    "monsoon": [
        "rain", "rainy", "monsoon", "wet season",
        "waterlogged", "flooded", "overwatered", "soggy soil"
    ],
    "summer": [
        "summer", "hot", "heat", "scorching",
        "dry heat", "high temperature", "blazing sun", "drought"
    ],
}


def extract_facts(
    user_query: str,
    expanded_queries: list[str],
    retrieved_chunks: list[str],
) -> set[str]:
    """
    Extract atomic Ayurvedic facts from all text sources.

    Steps:
    1. Combine user_query + expanded_queries + retrieved_chunks into
       a single lowercase string.
    2. For each (fact, keywords) entry in SYMPTOM_KEYWORDS, check if
       ANY keyword appears in the combined text.
    3. Return the set of matched fact strings.
    """
    combined = " ".join(
        [user_query] + expanded_queries + retrieved_chunks
    ).lower()

    facts: set[str] = set()
    for fact, keywords in SYMPTOM_KEYWORDS.items():
        if any(kw in combined for kw in keywords):
            facts.add(fact)

    return facts
```

---

## Phase 3 — Forward chaining engine

**Goal:** Create `inference/forward_chain.py`.

**File:** `inference/forward_chain.py`

```python
from __future__ import annotations
from dataclasses import dataclass, field
import yaml


@dataclass
class Rule:
    id: str
    name: str
    source: str
    priority: int
    conditions: list[str]
    conclusions: list[str]
    dosha: str
    explanation: str


@dataclass
class ForwardChainResult:
    final_facts: set[str]
    fired_rules: list[dict]
    proof_trace: list[str]
    dosha_scores: dict[str, int]
    new_facts: set[str]


class ForwardChainingEngine:
    """
    Fixed-point forward chaining over Vrikshayurveda rules.

    Algorithm:
    - Each pass scans all unfired rules sorted by conflict resolution order.
    - A rule fires when ALL its conditions are in working memory.
    - Conclusions are added to working memory as new facts.
    - Iteration stops when no new rules fire in a full pass (fixed point).
    """

    def __init__(self, rules_path: str = "vrikshayurveda_rules.yaml") -> None:
        self.rules: list[Rule] = self._load_rules(rules_path)

    def _load_rules(self, path: str) -> list[Rule]:
        with open(path, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f)
        return [
            Rule(
                id=r["id"],
                name=r["name"],
                source=r["source"],
                priority=r["priority"],
                conditions=r["conditions"],
                conclusions=r["conclusions"],
                dosha=r["dosha"],
                explanation=r["explanation"],
            )
            for r in data["rules"]
        ]

    def _resolve_conflicts(self, fireable: list[Rule]) -> list[Rule]:
        """
        Sort fireable rules by:
        1. priority descending (explicit)
        2. len(conditions) descending (specificity)
        3. source == "vrikshayurveda" before "derived"
        """
        return sorted(
            fireable,
            key=lambda r: (
                r.priority,
                len(r.conditions),
                1 if r.source == "vrikshayurveda" else 0,
            ),
            reverse=True,
        )

    def run(self, facts: set[str]) -> ForwardChainResult:
        working_memory: set[str] = set(facts)
        fired_ids: set[str] = set()
        fired_rules: list[dict] = []
        proof_trace: list[str] = []
        dosha_scores: dict[str, int] = {"vata": 0, "pitta": 0, "kapha": 0}

        changed = True
        while changed:
            changed = False
            fireable = [
                r for r in self.rules
                if r.id not in fired_ids
                and all(c in working_memory for c in r.conditions)
            ]
            ordered = self._resolve_conflicts(fireable)

            for rule in ordered:
                new_facts = [c for c in rule.conclusions if c not in working_memory]
                working_memory.update(rule.conclusions)
                fired_ids.add(rule.id)
                fired_rules.append({
                    "rule_id": rule.id,
                    "name": rule.name,
                    "conditions": rule.conditions,
                    "conclusions": rule.conclusions,
                    "explanation": rule.explanation,
                    "dosha": rule.dosha,
                })

                cond_str = " ∧ ".join(rule.conditions)
                conc_str = " + ".join(rule.conclusions)
                proof_trace.append(
                    f"[{rule.id}] {cond_str} → {conc_str}"
                )

                if rule.dosha in dosha_scores:
                    dosha_scores[rule.dosha] += 1

                if new_facts:
                    changed = True

        return ForwardChainResult(
            final_facts=working_memory,
            fired_rules=fired_rules,
            proof_trace=proof_trace,
            dosha_scores=dosha_scores,
            new_facts=working_memory - facts,
        )
```

---

## Phase 4 — Backward chaining engine

**Goal:** Create `inference/backward_chain.py`.

**File:** `inference/backward_chain.py`

```python
from __future__ import annotations
from dataclasses import dataclass, field
import yaml
from inference.forward_chain import Rule


@dataclass
class BackwardChainResult:
    goal: str
    proved: bool
    proof_tree: dict
    proof_trace: list[str]
    confidence: float


@dataclass
class DifferentialResult:
    rankings: list[dict]
    primary: str
    primary_confidence: float
    trace: list[str]


class BackwardChainingEngine:
    """
    Recursive backward chaining for differential dosha diagnosis.

    Algorithm:
    - To prove a goal: check if it is already a known fact.
    - If not, find all rules whose conclusions contain the goal.
    - Recursively attempt to prove each condition of those rules.
    - A rule fires if ALL its conditions are proved.
    - Returns a proof tree showing exactly how the goal was established.
    """

    def __init__(self, rules_path: str = "vrikshayurveda_rules.yaml") -> None:
        raw_rules = self._load_rules(rules_path)
        self.rules = raw_rules
        # goal_index: conclusion → list of rules that can prove it
        self.goal_index: dict[str, list[Rule]] = {}
        for rule in raw_rules:
            for conc in rule.conclusions:
                self.goal_index.setdefault(conc, []).append(rule)

    def _load_rules(self, path: str) -> list[Rule]:
        with open(path, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f)
        return [
            Rule(
                id=r["id"],
                name=r["name"],
                source=r["source"],
                priority=r["priority"],
                conditions=r["conditions"],
                conclusions=r["conclusions"],
                dosha=r["dosha"],
                explanation=r["explanation"],
            )
            for r in data["rules"]
        ]

    def prove(
        self,
        goal: str,
        known_facts: set[str],
        depth: int = 0,
        max_depth: int = 6,
        visited: set[str] | None = None,
    ) -> BackwardChainResult:
        if visited is None:
            visited = set()

        trace: list[str] = []
        indent = "  " * depth

        # Base case: goal already in working memory
        if goal in known_facts:
            trace.append(f"{indent}✓ '{goal}' is a known fact")
            return BackwardChainResult(
                goal=goal,
                proved=True,
                proof_tree={"goal": goal, "grounded": True},
                proof_trace=trace,
                confidence=1.0,
            )

        # Depth limit guard
        if depth >= max_depth or goal in visited:
            trace.append(f"{indent}✗ Cannot prove '{goal}' (depth limit or cycle)")
            return BackwardChainResult(
                goal=goal,
                proved=False,
                proof_tree={"goal": goal, "grounded": False},
                proof_trace=trace,
                confidence=0.0,
            )

        visited = visited | {goal}
        candidate_rules = sorted(
            self.goal_index.get(goal, []),
            key=lambda r: (r.priority, len(r.conditions)),
            reverse=True,
        )

        trace.append(f"{indent}? Trying to prove '{goal}'")

        for rule in candidate_rules:
            trace.append(f"{indent}  Trying rule {rule.id}: {rule.name}")
            sub_proofs: list[dict] = []
            all_proved = True
            conditions_proved = 0

            for cond in rule.conditions:
                sub_result = self.prove(
                    cond, known_facts, depth + 1, max_depth, visited
                )
                trace.extend(sub_result.proof_trace)
                sub_proofs.append(sub_result.proof_tree)
                if sub_result.proved:
                    conditions_proved += 1
                else:
                    all_proved = False
                    break

            if all_proved:
                trace.append(
                    f"{indent}  ✓ Rule {rule.id} fires → '{goal}' proved"
                )
                return BackwardChainResult(
                    goal=goal,
                    proved=True,
                    proof_tree={
                        "goal": goal,
                        "rule_used": rule.id,
                        "rule_name": rule.name,
                        "sub_proofs": sub_proofs,
                    },
                    proof_trace=trace,
                    confidence=1.0,
                )

        trace.append(f"{indent}✗ '{goal}' cannot be proved from known facts")
        return BackwardChainResult(
            goal=goal,
            proved=False,
            proof_tree={"goal": goal, "grounded": False},
            proof_trace=trace,
            confidence=0.0,
        )

    def _partial_confidence(self, goal: str, known_facts: set[str]) -> float:
        """
        Score how strongly the known facts support a goal,
        even if the goal cannot be fully proved.
        Returns fraction of best-matching rule's conditions that are satisfied.
        """
        if goal in known_facts:
            return 1.0
        candidates = self.goal_index.get(goal, [])
        if not candidates:
            return 0.0
        best = max(
            sum(1 for c in rule.conditions if c in known_facts) / len(rule.conditions)
            for rule in candidates
        )
        return round(best, 2)

    def differential_diagnosis(self, known_facts: set[str]) -> DifferentialResult:
        """
        Attempt to prove each primary dosha hypothesis.
        Score each by confidence and rank descending.
        """
        hypotheses = [
            ("vata", "vata_disorder"),
            ("pitta", "pitta_imbalance"),
            ("kapha", "kapha_obstruction"),
        ]
        rankings: list[dict] = []
        all_trace: list[str] = []

        for dosha, goal in hypotheses:
            result = self.prove(goal, known_facts)
            all_trace.extend(result.proof_trace)
            confidence = (
                result.confidence
                if result.proved
                else self._partial_confidence(goal, known_facts)
            )
            rankings.append({
                "dosha": dosha,
                "goal": goal,
                "proved": result.proved,
                "confidence": confidence,
                "proof_tree": result.proof_tree,
            })

        rankings.sort(key=lambda x: x["confidence"], reverse=True)
        primary = rankings[0]

        return DifferentialResult(
            rankings=rankings,
            primary=primary["dosha"],
            primary_confidence=primary["confidence"],
            trace=all_trace,
        )
```

---

## Phase 5 — Unified inference engine

**Goal:** Create `inference/engine.py` — the single entry point that
replaces the old symbolic engine.

**File:** `inference/engine.py`

```python
from __future__ import annotations
from dataclasses import dataclass
from inference.forward_chain import ForwardChainingEngine, ForwardChainResult
from inference.backward_chain import BackwardChainingEngine, DifferentialResult
from inference.symptom_extractor import extract_facts


@dataclass
class DiagnosisResult:
    primary_diagnosis: str
    dosha: str
    confidence: float
    forward_trace: list[str]
    backward_trace: list[str]
    full_proof_tree: dict
    all_facts: set[str]
    fired_rules: list[dict]
    llm_context: str


class VrikshayurvedaInferenceEngine:
    """
    Unified entry point for JAIM's inference system.
    Replaces the old keyword-based symbolic rule engine entirely.

    Call diagnose() with the same inputs the old engine received.
    It returns DiagnosisResult which includes llm_context — a compact
    paragraph ready to be injected into the Llama prompt.
    """

    def __init__(self, rules_path: str = "vrikshayurveda_rules.yaml") -> None:
        self.forward_engine = ForwardChainingEngine(rules_path)
        self.backward_engine = BackwardChainingEngine(rules_path)

    def diagnose(
        self,
        user_query: str,
        expanded_queries: list[str],
        retrieved_chunks: list[str],
    ) -> DiagnosisResult:
        # Step 1: Extract atomic facts from all text sources
        facts = extract_facts(user_query, expanded_queries, retrieved_chunks)

        # Step 2: Forward chain — enrich working memory
        forward_result: ForwardChainResult = self.forward_engine.run(facts)

        # Step 3: Backward chain — differential diagnosis
        diff_result: DifferentialResult = self.backward_engine.differential_diagnosis(
            forward_result.final_facts
        )

        # Step 4: Determine primary diagnosis label
        primary_dosha = diff_result.primary
        confidence = diff_result.primary_confidence
        dosha_map = {
            "vata": "Vata disorder",
            "pitta": "Pitta imbalance",
            "kapha": "Kapha obstruction",
        }
        primary_diagnosis = dosha_map.get(primary_dosha, "Undetermined disorder")

        # Step 5: Extract remedy recommendations from final facts
        remedies = [
            f.replace("recommend_", "").replace("_", " ")
            for f in forward_result.final_facts
            if f.startswith("recommend_")
        ]

        # Step 6: Build dosha ranking string
        ranking_str = ", ".join(
            f"{r['dosha']} ({round(r['confidence'] * 100)}%)"
            for r in diff_result.rankings
        )

        # Step 7: Build llm_context paragraph for Llama prompt injection
        chain_str = " → ".join(forward_result.proof_trace) if forward_result.proof_trace else "No rules fired"
        remedy_str = ", ".join(remedies) if remedies else "none identified"

        llm_context = (
            f"Inference engine diagnosis: {primary_diagnosis} "
            f"({round(confidence * 100)}% confidence). "
            f"Reasoning chain: {chain_str}. "
            f"Differential dosha ranking: {ranking_str}. "
            f"Recommended Ayurvedic treatments: {remedy_str}."
        )

        return DiagnosisResult(
            primary_diagnosis=primary_diagnosis,
            dosha=primary_dosha,
            confidence=confidence,
            forward_trace=forward_result.proof_trace,
            backward_trace=diff_result.trace,
            full_proof_tree={"differential": [r["proof_tree"] for r in diff_result.rankings]},
            all_facts=forward_result.final_facts,
            fired_rules=forward_result.fired_rules,
            llm_context=llm_context,
        )
```

---

## Phase 6 — Smoke test

**Goal:** Create `inference/smoke_test.py` and verify all three scenarios
pass before touching `app.py`.

**File:** `inference/smoke_test.py`

```python
"""
Run with: python -m inference.smoke_test
All three scenarios must pass before proceeding to Phase 7.
"""
from inference.engine import VrikshayurvedaInferenceEngine

engine = VrikshayurvedaInferenceEngine()


def run_scenario(name: str, facts_query: str) -> None:
    print(f"\n{'='*60}")
    print(f"SCENARIO: {name}")
    print('='*60)
    result = engine.diagnose(
        user_query=facts_query,
        expanded_queries=[],
        retrieved_chunks=[],
    )
    print(f"Primary diagnosis : {result.primary_diagnosis}")
    print(f"Dosha             : {result.dosha}")
    print(f"Confidence        : {round(result.confidence * 100)}%")
    print(f"\nForward trace:")
    for line in result.forward_trace:
        print(f"  {line}")
    print(f"\nAll facts in working memory:")
    for f in sorted(result.all_facts):
        print(f"  - {f}")
    print(f"\nLLM context paragraph:")
    print(f"  {result.llm_context}")


run_scenario(
    name="Vata monsoon chain",
    facts_query="The plant has yellowing leaves, stunted growth, and it is monsoon season with waterlogged soil.",
)

run_scenario(
    name="Pitta summer chain",
    facts_query="The plant is wilting badly in scorching summer heat.",
)

run_scenario(
    name="Severe Pitta chain",
    facts_query="The plant is wilting in summer, with mushy rotting roots.",
)
```

**Expected outputs:**

| Scenario | Expected rules to fire | Must appear in final facts |
|---|---|---|
| Vata monsoon | R001 → R005 → R008 → R009 | `vata_disorder`, `root_based_cause`, `moisture_excess`, `recommend_drainage`, `recommend_neem_bark` |
| Pitta summer | R002 → R006 → R010 → R011 | `pitta_imbalance`, `heat_stress`, `solar_damage`, `recommend_chandana_paste`, `recommend_shade_cloth` |
| Severe Pitta | R002 → R004 → R006 → R019 → R020 | `severe_pitta_disorder`, `recommend_chandana_paste`, `recommend_drainage`, `recommend_vetiver_root_soak` |

**Run it:**
```bash
python -m inference.smoke_test
```

Do not proceed to Phase 7 until all three pass.

---

## Phase 7 — Integrate into app.py

**Goal:** Three surgical changes to `app.py`. Touch nothing else.

### Change 1 — Add imports (top of file, after existing imports)

```python
from inference.engine import VrikshayurvedaInferenceEngine, DiagnosisResult

# Instantiate once at module level (not inside a function)
_inference_engine = VrikshayurvedaInferenceEngine()
```

### Change 2 — Replace old engine call

Find the function in `app.py` that calls the old symbolic rule engine.
It will look something like:

```python
# OLD — delete this
symbolic_diagnosis = old_symbolic_engine(user_query, retrieved_chunks)
```

Replace with:

```python
# NEW
diagnosis_result: DiagnosisResult = _inference_engine.diagnose(
    user_query=user_query,
    expanded_queries=expanded_queries,   # your existing expanded_queries variable
    retrieved_chunks=retrieved_chunks,   # your existing chunks variable
)
symbolic_diagnosis = diagnosis_result.primary_diagnosis
llm_grounding = diagnosis_result.llm_context
```

### Change 3 — Inject grounding into Llama prompt

Find where the Llama prompt string is assembled. Add `llm_grounding`
as the first section, before retrieved chunks:

```python
prompt = f"""
[Inference Engine Analysis]
{llm_grounding}

[Retrieved Vrikshayurveda Passages]
{chunks_text}

[User Question]
{user_query}
"""
```

Also return `diagnosis_result` alongside the existing return values
of the pipeline function so the Gradio UI can access it in Phase 8:

```python
# If your function currently returns (response_text, sources):
return response_text, sources, diagnosis_result
```

---

## Phase 8 — Symptom Chain Explorer (Gradio tab)

**Goal:** Add a new tab to the existing `gr.Blocks()` layout. Do NOT
restructure existing tabs — only append a new one.

### 8a — Embed the visualizer HTML

Add this string constant near the top of `app.py`,
after the engine import:

```python
VISUALIZER_HTML = """
<!DOCTYPE html><html><head><style>
body{margin:0;font-family:sans-serif;background:transparent;}
*{box-sizing:border-box;}
#canvas{width:100%;display:block;background:#f8f8f6;border-radius:8px;}
.fact-pill{display:inline-block;font-size:11px;padding:2px 8px;border-radius:10px;margin:2px;border:0.5px solid;}
.fp-symptom{background:#E1F5EE;color:#085041;border-color:#0F6E56;}
.fp-derived{background:#EEEDFE;color:#3C3489;border-color:#534AB7;}
.fp-dosha{background:#FAEEDA;color:#633806;border-color:#854F0B;}
.fp-remedy{background:#FAECE7;color:#712B13;border-color:#993C1D;}
.trace-row{font-size:12px;padding:5px 0;border-bottom:0.5px solid #e0e0e0;line-height:1.6;}
.rule-id{background:#E1F5EE;color:#085041;padding:1px 6px;border-radius:4px;font-weight:600;font-size:11px;margin-right:6px;}
.arr-sym{color:#BA7517;margin:0 5px;}
</style></head><body>
<canvas id="canvas" height="340"></canvas>
<div style="margin-top:12px;display:grid;grid-template-columns:1fr 1fr;gap:12px;">
  <div>
    <div style="font-size:12px;font-weight:600;color:#555;margin-bottom:6px;">Working memory</div>
    <div id="fact-store" style="min-height:36px;"></div>
  </div>
  <div>
    <div style="font-size:12px;font-weight:600;color:#555;margin-bottom:6px;">Proof trace</div>
    <div id="proof-trace" style="min-height:36px;max-height:180px;overflow-y:auto;"></div>
  </div>
</div>
<div style="margin-top:10px;font-size:11px;color:#888;">
  <span style="display:inline-block;width:10px;height:10px;background:#E1F5EE;border:1px solid #0F6E56;border-radius:2px;margin-right:4px;vertical-align:middle;"></span>symptom
  <span style="display:inline-block;width:10px;height:10px;background:#EEEDFE;border:1px solid #534AB7;border-radius:2px;margin-left:8px;margin-right:4px;vertical-align:middle;"></span>derived
  <span style="display:inline-block;width:10px;height:10px;background:#FAEEDA;border:1px solid #854F0B;border-radius:2px;margin-left:8px;margin-right:4px;vertical-align:middle;"></span>dosha
  <span style="display:inline-block;width:10px;height:10px;background:#FAECE7;border:1px solid #993C1D;border-radius:2px;margin-left:8px;margin-right:4px;vertical-align:middle;"></span>remedy
</div>
<script>
const RULES=[
  {id:"R001",conditions:["yellowing","stunted_growth"],conclusions:["vata_disorder"],priority:10},
  {id:"R002",conditions:["wilting"],conclusions:["pitta_imbalance"],priority:9},
  {id:"R003",conditions:["bark_lesions"],conclusions:["kapha_obstruction"],priority:9},
  {id:"R004",conditions:["root_rot"],conclusions:["root_based_cause"],priority:9},
  {id:"R005",conditions:["vata_disorder","monsoon"],conclusions:["root_based_cause","moisture_excess"],priority:10},
  {id:"R006",conditions:["pitta_imbalance","summer"],conclusions:["heat_stress","solar_damage"],priority:10},
  {id:"R007",conditions:["kapha_obstruction","monsoon"],conclusions:["fungal_kapha_disorder","moisture_excess"],priority:8},
  {id:"R008",conditions:["root_based_cause"],conclusions:["recommend_drainage"],priority:7},
  {id:"R009",conditions:["moisture_excess"],conclusions:["recommend_neem_bark"],priority:6},
  {id:"R010",conditions:["heat_stress"],conclusions:["recommend_chandana_paste"],priority:7},
  {id:"R011",conditions:["solar_damage"],conclusions:["recommend_shade_cloth"],priority:5},
  {id:"R012",conditions:["kapha_obstruction"],conclusions:["recommend_bark_scraping"],priority:6},
  {id:"R013",conditions:["vata_disorder","root_rot"],conclusions:["fungal_vata_disorder"],priority:11},
  {id:"R014",conditions:["fungal_vata_disorder"],conclusions:["recommend_drainage","recommend_neem_bark"],priority:6},
  {id:"R015",conditions:["yellowing","wilting"],conclusions:["vata_pitta_combined"],priority:8},
  {id:"R016",conditions:["vata_pitta_combined"],conclusions:["recommend_triphala_water","recommend_shade_cloth"],priority:7},
  {id:"R017",conditions:["bark_lesions","monsoon"],conclusions:["fungal_kapha_disorder"],priority:9},
  {id:"R018",conditions:["fungal_kapha_disorder"],conclusions:["recommend_bark_scraping","recommend_neem_bark"],priority:6},
  {id:"R019",conditions:["wilting","summer","root_rot"],conclusions:["severe_pitta_disorder"],priority:8},
  {id:"R020",conditions:["severe_pitta_disorder"],conclusions:["recommend_chandana_paste","recommend_drainage","recommend_vetiver_root_soak"],priority:9},
];
const FACT_TYPE=f=>{
  const s=["yellowing","stunted_growth","wilting","bark_lesions","root_rot","monsoon","summer"];
  if(s.includes(f))return"symptom";
  if(f.startsWith("recommend_"))return"remedy";
  if(["vata_disorder","pitta_imbalance","kapha_obstruction","fungal_vata_disorder",
      "fungal_kapha_disorder","severe_pitta_disorder","vata_pitta_combined"].includes(f))return"dosha";
  return"derived";
};
const NC={symptom:{fill:"#E1F5EE",stroke:"#0F6E56",text:"#085041"},
          derived:{fill:"#EEEDFE",stroke:"#534AB7",text:"#3C3489"},
          dosha:{fill:"#FAEEDA",stroke:"#854F0B",text:"#633806"},
          remedy:{fill:"#FAECE7",stroke:"#993C1D",text:"#712B13"}};
const canvas=document.getElementById("canvas");
const ctx=canvas.getContext("2d");
let nodePos={},edges=[],workingMem=new Set(),activeFacts=new Set();
let chainSteps=[],stepIdx=0;
function dpr(){return window.devicePixelRatio||1;}
function initCanvas(){
  const W=canvas.parentElement.offsetWidth||640;
  canvas.width=W*dpr();canvas.height=340*dpr();
  canvas.style.width=W+"px";canvas.style.height="340px";
  ctx.scale(dpr(),dpr());
}
function layoutNodes(facts){
  const arr=[...facts];
  const W=(canvas.offsetWidth/dpr())||640;
  const cols=Math.max(3,Math.ceil(Math.sqrt(arr.length*1.8)));
  const rows=Math.ceil(arr.length/cols);
  const cellW=(W-80)/cols,cellH=(300-40)/Math.max(rows,1);
  arr.forEach((f,i)=>{
    nodePos[f]={x:40+((i%cols)+0.5)*cellW,y:20+(Math.floor(i/cols)+0.5)*cellH};
  });
}
function roundRect(ctx,x,y,w,h,r){
  ctx.beginPath();ctx.moveTo(x+r,y);ctx.lineTo(x+w-r,y);
  ctx.quadraticCurveTo(x+w,y,x+w,y+r);ctx.lineTo(x+w,y+h-r);
  ctx.quadraticCurveTo(x+w,y+h,x+w-r,y+h);ctx.lineTo(x+r,y+h);
  ctx.quadraticCurveTo(x,y+h,x,y+h-r);ctx.lineTo(x,y+r);
  ctx.quadraticCurveTo(x,y,x+r,y);ctx.closePath();
}
function drawArrow(x1,y1,x2,y2,col){
  const a=Math.atan2(y2-y1,x2-x1);
  const mx=(x1+x2)/2,my=(y1+y2)/2;
  ctx.beginPath();
  ctx.moveTo(mx-7*Math.cos(a-0.4),my-7*Math.sin(a-0.4));
  ctx.lineTo(mx,my);
  ctx.lineTo(mx-7*Math.cos(a+0.4),my-7*Math.sin(a+0.4));
  ctx.strokeStyle=col;ctx.lineWidth=1.5;ctx.stroke();
}
function draw(){
  const W=(canvas.offsetWidth/dpr())||640;
  ctx.clearRect(0,0,W,340);
  edges.forEach(e=>{
    const f=nodePos[e.from],t=nodePos[e.to];if(!f||!t)return;
    ctx.beginPath();ctx.moveTo(f.x,f.y);ctx.lineTo(t.x,t.y);
    if(e.active){ctx.strokeStyle="#1D9E75";ctx.lineWidth=1.8;ctx.setLineDash([6,3]);}
    else{ctx.strokeStyle="#C8C6BE";ctx.lineWidth=0.8;ctx.setLineDash([]);}
    ctx.stroke();ctx.setLineDash([]);
    if(e.active)drawArrow(f.x,f.y,t.x,t.y,"#1D9E75");
  });
  [...workingMem].forEach(fact=>{
    const p=nodePos[fact];if(!p)return;
    const col=NC[FACT_TYPE(fact)];
    ctx.font="11px sans-serif";
    const label=fact.replace(/_/g," ");
    const tw=ctx.measureText(label).width;
    const w=Math.max(tw+22,72),h=26;
    ctx.fillStyle=col.fill;ctx.strokeStyle=col.stroke;
    ctx.lineWidth=activeFacts.has(fact)?1.8:0.7;
    roundRect(ctx,p.x-w/2,p.y-h/2,w,h,6);ctx.fill();ctx.stroke();
    ctx.fillStyle=col.text;ctx.textAlign="center";ctx.textBaseline="middle";
    ctx.fillText(label,p.x,p.y);
  });
}
function buildSteps(facts){
  let mem=new Set(facts),fired=new Set(),steps=[];
  let changed=true;
  while(changed){
    changed=false;
    const sorted=[...RULES].sort((a,b)=>b.priority-a.priority||b.conditions.length-a.conditions.length);
    for(const rule of sorted){
      if(fired.has(rule.id))continue;
      if(rule.conditions.every(c=>mem.has(c))){
        const nf=rule.conclusions.filter(c=>!mem.has(c));
        steps.push({rule,newFacts:nf});
        nf.forEach(f=>mem.add(f));
        fired.add(rule.id);changed=true;
      }
    }
  }
  return steps;
}
function buildEdges(steps){
  return steps.flatMap(s=>
    s.rule.conditions.flatMap(c=>s.rule.conclusions.map(conc=>({from:c,to:conc,ruleId:s.rule.id,active:false})))
  );
}
function appendTrace(rule,newFacts){
  const pt=document.getElementById("proof-trace");
  const d=document.createElement("div");d.className="trace-row";
  const conds=rule.conditions.map(c=>`<span class="fact-pill fp-${FACT_TYPE(c)}">${c.replace(/_/g," ")}</span>`).join(" ∧ ");
  const concs=newFacts.map(c=>`<span class="fact-pill fp-${FACT_TYPE(c)}">${c.replace(/_/g," ")}</span>`).join(" ");
  d.innerHTML=`<span class="rule-id">${rule.id}</span>${conds}<span class="arr-sym">→</span>${concs||'<span style="color:#aaa;font-size:11px;">already known</span>'}`;
  pt.appendChild(d);pt.scrollTop=pt.scrollHeight;
  updateFactStore();
}
function updateFactStore(){
  document.getElementById("fact-store").innerHTML=[...workingMem].map(f=>
    `<span class="fact-pill fp-${FACT_TYPE(f)}">${f.replace(/_/g," ")}</span>`
  ).join("");
}
window.runChainWithFacts=function(factsArray){
  activeFacts=new Set(factsArray);workingMem=new Set(factsArray);
  chainSteps=buildSteps(activeFacts);stepIdx=0;
  document.getElementById("proof-trace").innerHTML="";
  const allFacts=new Set([...activeFacts,...chainSteps.flatMap(s=>s.newFacts)]);
  allFacts.forEach(f=>workingMem.add(f));
  layoutNodes(allFacts);edges=buildEdges(chainSteps);
  updateFactStore();draw();
  const fire=i=>{
    if(i>=chainSteps.length)return;
    const s=chainSteps[i];
    edges.forEach(e=>{if(e.ruleId===s.rule.id)e.active=true;});
    appendTrace(s.rule,s.newFacts);draw();
    setTimeout(()=>fire(i+1),700);
  };
  fire(0);
};
window.stepChainWithFacts=function(factsArray){
  if(!chainSteps.length){
    activeFacts=new Set(factsArray);workingMem=new Set(factsArray);
    chainSteps=buildSteps(activeFacts);stepIdx=0;
    document.getElementById("proof-trace").innerHTML="";
    const allFacts=new Set([...activeFacts,...chainSteps.flatMap(s=>s.newFacts)]);
    allFacts.forEach(f=>workingMem.add(f));
    layoutNodes(allFacts);edges=buildEdges(chainSteps);
    updateFactStore();draw();
  }
  if(stepIdx>=chainSteps.length)return;
  const s=chainSteps[stepIdx];
  edges.forEach(e=>{if(e.ruleId===s.rule.id)e.active=true;});
  appendTrace(s.rule,s.newFacts);draw();stepIdx++;
};
window.resetViz=function(){
  activeFacts=new Set();workingMem=new Set();
  chainSteps=[];stepIdx=0;edges=[];nodePos={};
  document.getElementById("proof-trace").innerHTML="";
  document.getElementById("fact-store").innerHTML="";
  initCanvas();draw();
};
initCanvas();draw();
window.addEventListener("resize",()=>{initCanvas();layoutNodes(workingMem);draw();});
</script></body></html>
"""
```

### 8b — Gradio functions for the explorer tab

Add these three functions to `app.py`:

```python
def _render_viz_with_facts(facts: list[str], mode: str = "run") -> str:
    """
    Injects a JS bootstrap call into the visualizer HTML
    so the canvas fires immediately on load.
    """
    if not facts:
        return VISUALIZER_HTML.replace(
            "initCanvas();draw();",
            "initCanvas();draw();",
        )
    facts_js = str(facts).replace("'", '"')
    call = (
        f"runChainWithFacts({facts_js});"
        if mode == "run"
        else f"stepChainWithFacts({facts_js});"
    )
    return VISUALIZER_HTML.replace(
        "initCanvas();draw();",
        f"initCanvas();draw();setTimeout(()=>{{{call}}},300);",
    )


def explorer_run(selected_symptoms: list[str]) -> str:
    return _render_viz_with_facts(selected_symptoms, mode="run")


def explorer_step(selected_symptoms: list[str]) -> str:
    return _render_viz_with_facts(selected_symptoms, mode="step")


def explorer_reset() -> str:
    return VISUALIZER_HTML
```

### 8c — Gradio tab definition

Inside your existing `gr.Blocks()` context, append:

```python
with gr.Tab("Symptom Chain Explorer"):
    gr.Markdown("### Select symptoms and watch the inference engine reason step by step.")

    symptom_selector = gr.CheckboxGroup(
        choices=[
            ("Yellowing",      "yellowing"),
            ("Stunted growth", "stunted_growth"),
            ("Wilting",        "wilting"),
            ("Bark lesions",   "bark_lesions"),
            ("Root rot",       "root_rot"),
            ("Monsoon season", "monsoon"),
            ("Summer season",  "summer"),
        ],
        label="Observed symptoms",
    )

    with gr.Row():
        run_btn   = gr.Button("Run full chain",  variant="primary")
        step_btn  = gr.Button("Step through")
        reset_btn = gr.Button("Reset")

    viz_html = gr.HTML(value=VISUALIZER_HTML, label="Chain visualizer")

    run_btn.click(fn=explorer_run,   inputs=[symptom_selector], outputs=[viz_html])
    step_btn.click(fn=explorer_step, inputs=[symptom_selector], outputs=[viz_html])
    reset_btn.click(fn=explorer_reset, inputs=[], outputs=[viz_html])
```

---

## Phase 9 — Final checklist

Before marking complete, verify each item:

- [ ] `vrikshayurveda_rules.yaml` exists in project root with exactly 20 rules
- [ ] `inference/__init__.py` exists (empty)
- [ ] `inference/symptom_extractor.py` — `extract_facts()` returns a `set[str]`
- [ ] `inference/forward_chain.py` — `ForwardChainingEngine.run()` returns `ForwardChainResult`
- [ ] `inference/backward_chain.py` — `BackwardChainingEngine.differential_diagnosis()` returns `DifferentialResult`
- [ ] `inference/engine.py` — `VrikshayurvedaInferenceEngine.diagnose()` returns `DiagnosisResult`
- [ ] `python -m inference.smoke_test` runs without errors and all 3 scenarios pass
- [ ] `app.py` has exactly 3 changes (import + engine call replacement + prompt injection)
- [ ] Gradio app launches without import errors
- [ ] "Symptom Chain Explorer" tab appears in the UI
- [ ] Selecting yellowing + stunted_growth + monsoon and clicking "Run full chain" animates R001 → R005 → R008 → R009 in the canvas

---

## Constraints (non-negotiable)

1. Pure Python — no `pyknow`, `owlready2`, or Prolog bindings
2. `PyYAML` for rule loading only — no new pip dependencies beyond this
3. Type hints throughout all `inference/` files
4. `app.py` changes are **surgical** — only the 3 changes in Phase 7
5. The old symbolic engine code is deleted after replacement
6. The `VISUALIZER_HTML` string is embedded verbatim — do not regenerate it from Python logic
7. Smoke test must pass before Phase 7 begins