senti/core/brain/inference/engine.py

"""
Inference-Time Scaling for Senti AI.

Three levels of reasoning depth:

LEVEL 1 — FAST (< 1 second):
  Simple lookups. What is PAYE? Direct answer.
  One deterministic calculation. No deep thinking needed.

LEVEL 2 — STANDARD (1-3 seconds):
  Standard financial calculation.
  Parse → Compute (Rust) → Format → Respond.
  Most requests land here.

LEVEL 3 — DEEP REASONING (5-15 seconds):
  Complex multi-step analysis.
  "Should I expand to a second branch?"
  "What is my optimal debt payoff strategy?"
  Uses: Chain-of-thought, Monte Carlo, multi-source synthesis.
"""

import re
import time
from typing import Optional

class InferenceEngine:
  
  LEVEL_3_INTENTS = {
    "SCENARIO_MONTE_CARLO",
    "BUSINESS_VALUATION",
    "EXPANSION_ANALYSIS",
    "DEBT_STRATEGY",
    "RETIREMENT_PLAN",
    "PORTFOLIO_ANALYSIS",
    "COMPARE_INVESTMENTS",
  }
  
  LEVEL_3_KEYWORDS = [
    "should i", "is it worth", "compare",
    "what if", "best option", "optimize",
    "recommend", "strategy", "plan",
    "niambie kama", "ni bora",
  ]
  
  LEVEL_1_KEYWORDS = [
    "what is", "define", "explain", "nini ni",
    "maana ya", "meaning", "how does",
  ]
  
  def determine_level(
    self,
    intent: str,
    query: str,
    computed_result: Optional[dict]
  ) -> int:
    """
    Determine reasoning depth needed.
    Returns 1, 2, or 3.
    """
    
    # Level 3: complex multi-step analysis
    if intent in self.LEVEL_3_INTENTS:
      return 3
    
    q_lower = query.lower()
    if any(kw in q_lower for kw in self.LEVEL_3_KEYWORDS):
      return 3
    
    # Level 1: simple explanation/definition
    if any(kw in q_lower for kw in self.LEVEL_1_KEYWORDS):
      if not computed_result or not computed_result.get("computed"):
        return 1
    
    # Level 2: standard calculation (default)
    return 2
  
  def level_1_response(
    self,
    query: str,
    rag_results: list,
    core_memory: str
  ) -> str:
    """
    Level 1: Direct knowledge lookup.
    Returns context for LLM to format.
    """
    
    if rag_results:
      context = "\n".join([
        f"[Source: {r['source']}]\n{r['text']}"
        for r in rag_results[:2]
      ])
      return f"KNOWLEDGE CONTEXT:\n{context}"
    
    return "INSTRUCTION: Answer from general financial knowledge."
  
  def level_2_response(
    self,
    query: str,
    computed: dict,
    rag_results: list,
    core_memory: str
  ) -> str:
    """
    Level 2: Standard calculation context.
    Rust computed the math. LLM formats it.
    """
    
    parts = []
    
    if core_memory:
      parts.append(core_memory)
    
    if computed and computed.get("computed"):
      result_str = self._format_computed_for_llm(computed)
      parts.append(f"COMPUTED RESULTS (from Rust engine):\n{result_str}")
    
    if rag_results:
      parts.append(f"REGULATORY CONTEXT:\n{rag_results[0]['text'][:400]}")
    
    return "\n\n".join(parts)
  
  def level_3_reasoning(
    self,
    query: str,
    intent: str,
    data: dict,
    computed: dict,
    rag_results: list,
    core_memory: str
  ) -> str:
    """
    Level 3: Chain-of-thought multi-step analysis.
    Decomposes complex query into sub-questions.
    Answers each. Synthesizes.
    """
    
    steps = []
    
    # Step 1: What is being asked?
    sub_questions = self._decompose_query(query, intent)
    
    # Step 2: Answer each sub-question
    for i, sub_q in enumerate(sub_questions):
      answer = self._answer_sub_question(sub_q, data, computed)
      steps.append({
        "step": i + 1,
        "question": sub_q,
        "answer": answer
      })
    
    # Step 3: Add Monte Carlo if financial projection
    if intent == "SCENARIO_MONTE_CARLO" or "what if" in query.lower():
      from core.engines.formulas.registry import calc
      if data.get("initial_value") or computed.get("result", {}).get("gross_profit"):
        initial = data.get("initial_value") or \
                  computed.get("result", {}).get("gross_profit", 100000)
        mc = calc.monte_carlo(
          initial_value=float(initial),
          annual_return_percent=data.get("return_percent", 12.0),
          annual_volatility_percent=data.get("volatility_percent", 20.0),
          years=int(data.get("years", 5))
        )
        steps.append({
          "step": len(steps) + 1,
          "question": "Monte Carlo simulation (10,000 scenarios)",
          "answer": (
            f"P10: KES {mc['percentile_10']:,.0f} | "
            f"P50: KES {mc['percentile_50']:,.0f} | "
            f"P90: KES {mc['percentile_90']:,.0f} | "
            f"Probability of profit: {mc['probability_profit_percent']:.1f}%"
          )
        })
    
    # Format reasoning chain for LLM
    reasoning = "DEEP ANALYSIS (chain-of-thought):\n"
    for step in steps:
      reasoning += f"\nStep {step['step']}: {step['question']}\n"
      reasoning += f"  → {step['answer']}\n"
    
    if core_memory:
      reasoning = core_memory + "\n\n" + reasoning
    
    if rag_results:
      reasoning += f"\nREGULATORY CONTEXT:\n{rag_results[0]['text'][:400]}"
    
    reasoning += "\n\nSYNTHESIS INSTRUCTION: Based on all steps above, provide a clear, specific recommendation with KES amounts."
    
    return reasoning
  
  def _decompose_query(self, query: str, intent: str) -> list[str]:
    """Break complex query into answerable sub-questions."""
    
    q = query.lower()
    
    if "expand" in q or "second branch" in q or "new shop" in q:
      return [
        "What is the current financial health?",
        "What are the projected expansion costs?",
        "What does cash flow trend show?",
        "What is the tax impact of expansion?",
        "What is the risk if revenue drops 20%?"
      ]
    
    if "debt" in q and ("pay" in q or "clear" in q or "strategy" in q):
      return [
        "What are all current debts and their rates?",
        "What is the avalanche payoff order (highest rate first)?",
        "What is the snowball payoff order (smallest balance first)?",
        "How much extra cash is available monthly?",
        "Which strategy saves more money overall?"
      ]
    
    if "invest" in q and "best" in q:
      return [
        "What is the investment amount and timeline?",
        "What products are available in Kenya?",
        "What are the current rates for each product?",
        "What is the risk tolerance?",
        "Which product gives best risk-adjusted return?"
      ]
    
    # Default decomposition for any complex query
    return [
      "What specific information is needed?",
      "What does the current data show?",
      "What is the recommended action?"
    ]
  
  def _answer_sub_question(
    self,
    question: str,
    data: dict,
    computed: dict
  ) -> str:
    """Give a brief answer to a sub-question using available data."""
    
    result = computed.get("result", {}) if computed else {}
    
    if "financial health" in question.lower():
      revenue = data.get("revenue") or result.get("revenue", 0)
      expenses = data.get("total_expenses") or result.get("total_expenses", 0)
      if revenue and expenses:
        margin = (revenue - expenses) / revenue * 100 if revenue else 0
        return f"Revenue KES {revenue:,.0f}, expenses KES {expenses:,.0f}, margin {margin:.1f}%"
      return "Insufficient data for financial health assessment"
    
    if "cash flow" in question.lower():
      profit = result.get("gross_profit") or result.get("net_profit")
      if profit:
        return f"Monthly profit: KES {profit:,.0f}. Cash flow appears {'positive' if profit > 0 else 'negative'}."
      return "Cash flow data not available"
    
    if "tax" in question.lower():
      from core.engines.formulas.registry import calc
      revenue = data.get("revenue") or result.get("revenue")
      if revenue:
        tot = calc.tot(float(revenue))
        return f"TOT at 3%: KES {tot['monthly_tax']:,.0f}/month if applicable"
      return "Tax impact requires revenue data"
    
    if "risk" in question.lower() and "20%" in question.lower():
      revenue = data.get("revenue") or result.get("revenue")
      if revenue:
        reduced = float(revenue) * 0.8
        from core.engines.formulas.registry import calc
        expenses = data.get("total_expenses", float(revenue) * 0.75)
        new_profit = reduced - float(expenses)
        return f"At -20% revenue: KES {reduced:,.0f} revenue, profit would be KES {new_profit:,.0f}"
      return "Cannot assess risk without revenue baseline"
    
    return "Analysis requires more context from conversation"
  
  def _format_computed_for_llm(self, computed: dict) -> str:
    """Format Rust computation results cleanly for LLM."""
    
    if not computed or not computed.get("result"):
      return "No computation result"
    
    result = computed["result"]
    lines = []
    
    for key, value in result.items():
      if key.startswith("_") or key == "formula_version":
        continue
      if isinstance(value, float):
        if key.endswith("percent") or "rate" in key.lower():
          lines.append(f"{key}: {value:.2f}%")
        elif value > 100:
          lines.append(f"{key}: KES {value:,.0f}")
        else:
          lines.append(f"{key}: {value:.2f}")
      elif isinstance(value, (int, str, bool)):
        lines.append(f"{key}: {value}")
    
    return "\n".join(lines)

inference_engine = InferenceEngine()