llm_interface.py

from __future__ import annotations
import os
import re
from typing import Optional, Tuple

from .inference import InferenceResult
import json
try:
    from openai import OpenAI  # new SDK
except Exception:
    OpenAI = None  # type: ignore


def parse_natural_query(text: str) -> Optional[Tuple[str, str, str]]:
    # Very naive patterns for demo purposes
    m = re.search(r"Is\s+(\w+)\s+(\w+)\s+of\s+(\w+)[?]?$", text, re.I)
    if m:
        subj, rel, obj = m.group(1), m.group(2), m.group(3)
        # Capitalize relation first letter to match KG labels by default
        return (subj, rel[0].upper() + rel[1:].lower(), obj)
    # Try explicit triple format: John Grandparent Alice
    parts = text.strip().split()
    if len(parts) == 3:
        return parts[0], parts[1], parts[2]
    return None


def narrate_result(res: InferenceResult) -> str:
    h, r, t = res.triple
    lines = [f"Query: ({h}, {r}, {t})"]
    lines.append(f"Probability: {res.probability:.2f}")
    if res.supporting_rules:
        lines.append("Supported by rules:")
        for s in res.supporting_rules[:5]:
            lines.append(f"  - {s}")
    else:
        lines.append("No supporting rules fired; probability from priors.")
    return "\n".join(lines)


def build_llm_context(
    question: str,
    res: InferenceResult,
    context_weights: list[tuple[int, float]] | None = None,
    rule_witnesses: list[str] | None = None,
) -> str:
    h, r, t = res.triple
    ctx = []
    ctx.append("Task: Answer the KG question with calibrated probability and a brief rationale.")
    ctx.append("")
    ctx.append(f"Question: {question}")
    ctx.append(f"KG Query: ({h}, {r}, {t})")
    ctx.append(f"Estimated Probability: {res.probability:.3f}")
    ctx.append("")
    if context_weights:
        ctx.append("Context Weights (learned reliabilities):")
        for cid, th in sorted(context_weights, key=lambda x: -x[1])[:5]:
            ctx.append(f"  - context {cid}: theta={th:.3f}")
        ctx.append("")
    if res.supporting_rules:
        ctx.append("Supporting Rules Fired:")
        for s in res.supporting_rules[:8]:
            ctx.append(f"  - {s}")
        ctx.append("")
    if rule_witnesses:
        ctx.append("Example Witnesses (bindings supporting the rules):")
        for w in rule_witnesses[:8]:
            ctx.append(f"  - {w}")
        ctx.append("")
    ctx.append("Guidelines:")
    ctx.append("- Be concise (2-4 sentences).")
    ctx.append("- Mention the probability and key rule(s) involved.")
    ctx.append("- If rules are weak or absent, note low confidence.")
    return "\n".join(ctx)


def llm_wrap_explanation(prompt: str, res: InferenceResult) -> str:
    explanation = narrate_result(res)
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key and OpenAI is not None:
        try:
            client = OpenAI(api_key=api_key)
            system = (
                "You are an assistant explaining knowledge-graph probabilistic reasoning."
            )
            user = (
                f"Question: {prompt}\n\n"
                f"Evidence:\n{explanation}\n\n"
                f"Explain the probability and key rules in simple terms."
            )
            resp = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                temperature=0.2,
            )
            return resp.choices[0].message.content or explanation
        except Exception:
            # Fall back to templated explanation
            pass
    return (
        f"You asked: {prompt}\n\n"
        f"Here is an explainable KG answer using probabilistic contexts:\n\n"
        f"{explanation}\n\n"
        f"Interpretation: Based on mined rules and learned context reliabilities, the system combines evidence with a Noisy-OR model to estimate the probability."
    )