deep_think.py

#!/usr/bin/env python3.12
"""Deep Think Engine - multi-step reasoning that makes a 7B model think like Opus.

Instead of single-shot Q&A, this engine:
1. DECOMPOSE: Break complex question into sub-questions
2. RESEARCH: Answer each sub-question (with optional web search)
3. CRITIQUE: Use a DIFFERENT model to find flaws in the answer
4. REVISE: Fix the flaws, incorporating the critique
5. VERIFY: Cross-check claims against AOMS and web sources
6. SYNTHESIZE: Combine everything into a final, high-quality answer

This is the cognitive core that turns cheap local models into deep reasoners.
"""

import json
import logging
import re
import subprocess
import tempfile
from pathlib import Path

import requests
import yaml

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [deep_think] %(levelname)s %(message)s",
)
log = logging.getLogger(__name__)

BASE_DIR = Path(__file__).resolve().parent


def load_config():
    with open(BASE_DIR / "config.yaml") as f:
        return yaml.safe_load(f)


def call_model(cfg, prompt, model=None, temperature=0.5, timeout=180):
    """Call a specific Ollama model."""
    url = cfg["ollama"]["url"]
    model = model or cfg["ollama"]["model"]
    try:
        resp = requests.post(url, json={
            "model": model,
            "prompt": prompt,
            "stream": False,
            "options": {"temperature": temperature},
        }, timeout=timeout)
        resp.raise_for_status()
        return resp.json().get("response", "").strip()
    except Exception as e:
        log.warning("Model %s failed: %s", model, e)
        return None


def call_primary(cfg, prompt, **kwargs):
    return call_model(cfg, prompt, model=cfg["ollama"]["model"], **kwargs)


def call_critic(cfg, prompt, **kwargs):
    """Use the OTHER model as critic for adversarial review."""
    return call_model(cfg, prompt, model=cfg["ollama"]["fallback_model"], **kwargs)


# --- Step 1: Decompose ---

def decompose(cfg, question):
    """Break a complex question into 2-4 sub-questions."""
    prompt = f"""Break this question into 2-4 simpler sub-questions that together would fully answer it.
Each sub-question should be independently answerable.

Question: {question}

Output ONLY a JSON array of strings. Example: ["sub-q 1?", "sub-q 2?"]
JSON:"""
    raw = call_primary(cfg, prompt, temperature=0.3)
    if not raw:
        return [question]  # fallback: just use original
    try:
        if "```" in raw:
            raw = raw.split("```")[1].lstrip("json\n")
        start = raw.index("[")
        end = raw.rindex("]") + 1
        subs = json.loads(raw[start:end])
        if isinstance(subs, list) and len(subs) >= 2:
            log.info("Decomposed into %d sub-questions", len(subs))
            return subs[:4]
    except (ValueError, json.JSONDecodeError):
        pass
    return [question]


# --- Step 2: Research each sub-question ---

def research_subquestion(cfg, sub_q, use_web=False):
    """Research a single sub-question. Optionally verify with web search."""
    prompt = f"""Research this question thoroughly. Provide specific facts, numbers, examples.
Be precise — if you're not sure about something, say so explicitly.

Question: {sub_q}

Detailed answer:"""
    answer = call_primary(cfg, prompt, temperature=0.5)

    # Optionally enhance with web search
    if use_web and answer:
        from web_researcher import search_and_verify
        web_facts = search_and_verify(sub_q, max_results=3)
        if web_facts:
            answer += f"\n\n[Web verification: {web_facts}]"

    return answer or ""


# --- Step 3: Critique ---

def critique(cfg, question, answer):
    """Use a DIFFERENT model to critique the answer. Adversarial review."""
    prompt = f"""You are a rigorous fact-checker and critical reviewer.
Analyze this answer for:
1. Factual errors or unsupported claims
2. Logical gaps or contradictions
3. Missing important perspectives
4. Oversimplifications

Question: {question}
Answer: {answer[:2000]}

Be specific. Point out EXACTLY what's wrong and what's missing.
If the answer is actually good, say "NO ISSUES FOUND" and explain why it's solid.

Critique:"""
    result = call_critic(cfg, prompt, temperature=0.3)
    if not result:
        return "Critique unavailable (model failure)"
    return result


# --- Step 4: Revise ---

def revise(cfg, question, original_answer, critique_text):
    """Revise the answer incorporating the critique."""
    if "NO ISSUES FOUND" in critique_text.upper():
        log.info("Critique found no issues, keeping original")
        return original_answer

    prompt = f"""Revise this answer to address the critique below.
Fix factual errors, fill gaps, and strengthen weak points.
Keep what's correct, improve what's not.

Original question: {question}
Original answer: {original_answer[:1500]}

Critique: {critique_text[:1000]}

Revised answer (be thorough and precise):"""
    revised = call_primary(cfg, prompt, temperature=0.4)
    return revised or original_answer


# --- Step 5: Verify against AOMS ---

def verify_against_memory(cfg, answer):
    """Cross-check claims against what AOMS already knows."""
    url = cfg["aoms"]["url"] + "/memory/search"
    # Extract key claims to verify
    claims_prompt = f"""Extract the 3 most important factual claims from this text.
Output a JSON array of short strings.

Text: {answer[:1000]}

JSON:"""
    raw = call_primary(cfg, claims_prompt, temperature=0.1)
    if not raw:
        return []

    try:
        start = raw.index("[")
        end = raw.rindex("]") + 1
        claims = json.loads(raw[start:end])
    except (ValueError, json.JSONDecodeError):
        return []

    verified = []
    for claim in claims[:3]:
        try:
            resp = requests.post(url, json={
                "query": claim,
                "tier": ["semantic"],
                "limit": 3,
            }, timeout=5)
            resp.raise_for_status()
            results = resp.json() if isinstance(resp.json(), list) else resp.json().get("results", [])
            if results:
                verified.append({"claim": claim, "supported": True, "sources": len(results)})
            else:
                verified.append({"claim": claim, "supported": False, "sources": 0})
        except Exception:
            verified.append({"claim": claim, "supported": "unknown", "sources": 0})

    return verified


# --- Step 6: Synthesize ---

def synthesize(cfg, question, sub_answers, critique_text, verification):
    """Combine sub-answers into a final, high-quality response."""
    sub_text = "\n\n".join(f"--- Sub-finding {i+1} ---\n{a[:500]}" for i, a in enumerate(sub_answers))

    verify_note = ""
    if verification:
        supported = sum(1 for v in verification if v.get("supported") is True)
        total = len(verification)
        verify_note = f"\n\nMemory verification: {supported}/{total} claims have supporting evidence in knowledge base."

    prompt = f"""Synthesize these research findings into a single, comprehensive answer.

Original question: {question}

Sub-findings:
{sub_text}

Critique notes: {critique_text[:500]}
{verify_note}

Create a clear, well-structured final answer that:
1. Directly addresses the original question
2. Integrates all sub-findings
3. Acknowledges limitations and uncertainties
4. Highlights the most important insights

Final answer:"""
    return call_primary(cfg, prompt, temperature=0.4)


# --- Step 7: Execute code (optional) ---

def execute_code_hypothesis(cfg, question, hypothesis):
    """If a hypothesis can be tested with code, write and run it safely."""
    prompt = f"""Can this hypothesis be tested with a short Python script (under 30 lines)?
If yes, write the script. If no, reply "NOT_TESTABLE".

Question: {question}
Hypothesis: {hypothesis}

Rules:
- No network requests, no file writes, no system calls
- Pure computation, data analysis, or math only
- Must print a clear result

Python script (or NOT_TESTABLE):"""
    code_response = call_primary(cfg, prompt, temperature=0.2)
    if not code_response or "NOT_TESTABLE" in code_response:
        return None

    # Extract code block
    code = code_response
    if "```python" in code:
        code = code.split("```python")[1].split("```")[0]
    elif "```" in code:
        code = code.split("```")[1].split("```")[0]

    # Safety check — reject dangerous code
    dangerous = ["import os", "import sys", "import subprocess", "open(", "exec(", "eval(",
                 "__import__", "requests", "urllib", "socket", "shutil", "pathlib"]
    if any(d in code for d in dangerous):
        log.warning("Code contains dangerous operations, skipping execution")
        return None

    # Execute in sandbox
    try:
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write(code)
            f.flush()
            result = subprocess.run(
                ["/home/dhawal/curiosity/.venv/bin/python3.12", f.name],
                capture_output=True, text=True, timeout=10,
            )
            if result.returncode == 0:
                log.info("Code execution successful")
                return result.stdout.strip()
            else:
                log.warning("Code execution failed: %s", result.stderr[:200])
                return None
    except subprocess.TimeoutExpired:
        log.warning("Code execution timed out")
        return None
    except Exception as e:
        log.warning("Code execution error: %s", e)
        return None


# --- Main orchestrator ---

def deep_think(question, use_web=False, use_code=False, cfg=None):
    """Run the full deep thinking pipeline on a question.

    Returns a ThinkResult with the final answer and metadata.
    """
    if cfg is None:
        cfg = load_config()

    log.info("=== DEEP THINK: %s ===", question[:80])

    # Step 1: Decompose
    sub_questions = decompose(cfg, question)
    log.info("Sub-questions: %s", [q[:50] for q in sub_questions])

    # Step 2: Research each sub-question
    sub_answers = []
    for i, sq in enumerate(sub_questions):
        log.info("Researching sub-question %d/%d: %s", i + 1, len(sub_questions), sq[:60])
        answer = research_subquestion(cfg, sq, use_web=use_web)
        if answer:
            sub_answers.append(answer)

    if not sub_answers:
        log.error("All sub-question research failed")
        return None

    # Combine for critique
    combined = "\n\n".join(sub_answers)

    # Step 3: Critique (adversarial — uses different model)
    log.info("Running adversarial critique...")
    critique_text = critique(cfg, question, combined)
    log.info("Critique: %s", critique_text[:100])

    # Step 4: Revise based on critique
    log.info("Revising based on critique...")
    revised = revise(cfg, question, combined, critique_text)

    # Step 5: Verify against AOMS memory
    log.info("Verifying against knowledge base...")
    verification = verify_against_memory(cfg, revised)
    supported = sum(1 for v in verification if v.get("supported") is True)
    log.info("Verification: %d/%d claims supported by existing knowledge", supported, len(verification))

    # Step 5b: Optional code execution
    code_result = None
    if use_code:
        log.info("Attempting code hypothesis testing...")
        code_result = execute_code_hypothesis(cfg, question, revised[:500])
        if code_result:
            log.info("Code result: %s", code_result[:100])
            revised += f"\n\n[Computational verification: {code_result}]"

    # Step 6: Synthesize final answer
    log.info("Synthesizing final answer...")
    final = synthesize(cfg, question, sub_answers, critique_text, verification)

    if not final:
        final = revised  # fallback to revised if synthesis fails

    result = {
        "question": question,
        "final_answer": final,
        "sub_questions": sub_questions,
        "sub_answers_count": len(sub_answers),
        "critique": critique_text,
        "verification": verification,
        "code_result": code_result,
        "steps": len(sub_questions) + 3,  # research + critique + revise + synthesize
    }

    log.info("=== DEEP THINK COMPLETE (%d steps, %d chars) ===",
             result["steps"], len(final or ""))
    return result


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Deep Think Engine")
    parser.add_argument("question", nargs="?", help="Question to think deeply about")
    parser.add_argument("--web", action="store_true", help="Enable web search verification")
    parser.add_argument("--code", action="store_true", help="Enable code hypothesis testing")
    args = parser.parse_args()

    if not args.question:
        print("Usage: deep_think.py 'your question here' [--web] [--code]")
        return

    result = deep_think(args.question, use_web=args.web, use_code=args.code)
    if result:
        print("\n" + "=" * 60)
        print("DEEP THINK RESULT")
        print("=" * 60)
        print(f"\nQuestion: {result['question']}")
        print(f"Steps: {result['steps']}")
        print(f"Sub-questions: {len(result['sub_questions'])}")
        print(f"Verification: {len(result['verification'])} claims checked")
        if result.get("code_result"):
            print(f"Code result: {result['code_result'][:200]}")
        print(f"\n--- ANSWER ---\n")
        print(result["final_answer"])
        print(f"\n--- CRITIQUE ---\n")
        print(result["critique"][:500])
    else:
        print("Deep think failed.")


if __name__ == "__main__":
    main()