phd_research_os/agents.py

"""
PhD Research OS — Agent Layer
==============================
The AI brain that powers all pipeline operations.
Uses fine-tuned Qwen2.5-3B-Instruct (or falls back to API) for:
  - Claim extraction from scientific text
  - Epistemic classification  
  - Confidence scoring
  - Contradiction detection
  - Query decomposition
  - Decision generation

Adheres to Research OS v11.0 rules:
  - Provenance Hierarchy: All agent output is Priority 5 (LLM Inference/Hypothesis)
  - Anchor Divergence: Agent output never overrides human-verified observations
  - Fixed-Point Math: All probabilities use scaled integers
  - Causal Lineage: Every claim traces back to source Observation_ID
"""

import json
import os
from typing import Optional
from dataclasses import dataclass

# System prompts for each agent role
PROMPTS = {
    "researcher": """You are the Researcher Agent of a PhD Research OS. Your role is to extract structured scientific claims from research paper text.

For each claim, output a JSON object with these fields:
- claim_id: string (CLM_XXXX format, auto-assigned)
- text: the claim text as stated in the paper
- epistemic_tag: one of "Fact", "Interpretation", "Hypothesis", "Conflict_Hypothesis"
- confidence: float [0,1] computed as evidence_strength × study_quality_weight × journal_tier_weight × completeness_penalty
- evidence_strength: float [0,1] based on directness of evidence
- study_type: one of "primary_experimental", "in_vitro", "simulation", "review_non_systematic", "meta_analysis", "case_study"
- missing_fields: list of field names that could not be determined from the text
- status: "Complete" if no missing fields, else "Incomplete"
- parameters: dict of key experimental parameters mentioned

Output must be valid JSON: {"claims": [...]}
Always classify epistemic tags conservatively. When uncertain, prefer "Interpretation" over "Fact".""",

    "epistemic_classifier": """You are the Epistemic Classifier of a PhD Research OS. Given a scientific statement, classify it into exactly one category:

- Fact: Directly supported by experimental data with quantitative evidence. Reproducible measurements.
- Interpretation: Author's explanation of data. Goes beyond what numbers strictly show.
- Hypothesis: Proposed mechanism or prediction not yet tested.
- Conflict_Hypothesis: Explicitly contradicts another established claim with evidence on both sides.

Output JSON: {"epistemic_tag": "...", "reasoning": "...", "confidence_in_classification": float}""",

    "confidence_scorer": """You are the Confidence Scorer of a PhD Research OS. Score claim confidence using:

confidence = evidence_strength × study_quality_weight × journal_tier_weight × completeness_penalty

study_quality_weight: primary_experimental=1.0, in_vitro=0.8, simulation=0.6, review_non_systematic=0.4, meta_analysis=1.0, case_study=0.3
journal_tier_weight: tier1=1.0, tier2=0.85, tier3=0.7, preprint=0.5
completeness_penalty: 1.0 if complete, 0.7 if missing key parameters

Output JSON: {"confidence": float, "evidence_strength": float, "study_quality_weight": float, "journal_tier_weight": float, "completeness_penalty": float, "reasoning": "..."}""",

    "verifier": """You are the Verifier Agent of a PhD Research OS. Given two scientific claims, determine if they contradict each other.

Output a Conflict Resolution Object as JSON:
- conflict_detected: boolean
- conflict_type: "value_mismatch", "methodology_difference", "scope_difference", or "no_conflict"
- generated_hypothesis: text explaining possible cause
- hypothesis_confidence: always "low" (never auto-set to high)
- resolution_status: "Unresolved"
- key_differences: list of specific differences
- recommended_action: what to investigate""",

    "query_planner": """You are the Query Planner of a PhD Research OS. Decompose broad research questions into 2-4 specific sub-queries for knowledge base search.

Output JSON: {"original_query": "...", "sub_queries": ["...", "..."], "reasoning": "..."}""",

    "decision_generator": """You are the Decision Agent of a PhD Research OS. Given research goals, knowledge gaps, and low-confidence claims, propose a Decision Object.

Output JSON:
- decision_id: string
- recommended_action: "experiment", "literature_search", "collaboration", "replication", "methodology_review"
- action_description: specific description
- expected_information_gain: float [0,1] = uncertainty × impact
- linked_goal_id: which goal this addresses
- linked_claim_ids: which claims this resolves
- priority: "high", "medium", "low"
- estimated_effort: time estimate"""
}


@dataclass
class AgentResponse:
    """Structured response from any agent."""
    success: bool
    data: dict
    raw_output: str
    provenance_level: int = 5  # LLM Inference — lowest priority per Research OS spec
    observation_id: Optional[str] = None
    tokens_in: int = 0
    tokens_out: int = 0
    cost_usd: float = 0.0


class ResearchOSBrain:
    """
    The AI brain of the PhD Research OS.
    
    Supports multiple backends:
    1. Fine-tuned local model (Qwen2.5-3B + LoRA adapter)
    2. API fallback (Anthropic Claude / OpenAI)
    
    All outputs are tagged as Priority 5 (LLM Hypothesis) per Research OS spec.
    Human verification is required to promote to higher provenance levels.
    """
    
    def __init__(self, model_path: str = None, backend: str = "local"):
        """
        Initialize the Research OS Brain.
        
        Args:
            model_path: Path to fine-tuned model or HF model ID
                       Default: "nkshirsa/phd-research-os-brain"
            backend: "local" for local inference, "api" for API fallback
        """
        self.backend = backend
        self.model_path = model_path or "nkshirsa/phd-research-os-brain"
        self.model = None
        self.tokenizer = None
        
        if backend == "local":
            self._load_local_model()
    
    def _load_local_model(self):
        """Load the fine-tuned model for local inference."""
        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer
            import torch
            
            print(f"Loading model from {self.model_path}...")
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_path,
                torch_dtype=torch.float16,
                device_map="auto"
            )
            print("Model loaded successfully.")
        except Exception as e:
            print(f"Warning: Could not load local model: {e}")
            print("Falling back to API mode. Set ANTHROPIC_API_KEY or OPENAI_API_KEY.")
            self.backend = "api"
    
    def _generate_local(self, messages: list, max_tokens: int = 2048) -> str:
        """Generate response using local model."""
        if self.model is None or self.tokenizer is None:
            raise RuntimeError("Local model not loaded")
        
        text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
        
        with __import__('torch').no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                temperature=0.1,  # Low temperature for structured output
                do_sample=True,
                top_p=0.95,
            )
        
        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        return response
    
    def _generate_api(self, messages: list, max_tokens: int = 2048) -> str:
        """Generate response using API (Anthropic or OpenAI fallback)."""
        api_key = os.environ.get("ANTHROPIC_API_KEY")
        if api_key:
            return self._call_anthropic(messages, max_tokens, api_key)
        
        api_key = os.environ.get("OPENAI_API_KEY")
        if api_key:
            return self._call_openai(messages, max_tokens, api_key)
        
        raise RuntimeError("No API key found. Set ANTHROPIC_API_KEY or OPENAI_API_KEY")
    
    def _call_anthropic(self, messages: list, max_tokens: int, api_key: str) -> str:
        """Call Anthropic API."""
        import httpx
        
        system_msg = ""
        api_messages = []
        for msg in messages:
            if msg["role"] == "system":
                system_msg = msg["content"]
            else:
                api_messages.append(msg)
        
        response = httpx.post(
            "https://api.anthropic.com/v1/messages",
            headers={
                "x-api-key": api_key,
                "anthropic-version": "2023-06-01",
                "content-type": "application/json"
            },
            json={
                "model": "claude-sonnet-4-20250514",
                "max_tokens": max_tokens,
                "system": system_msg,
                "messages": api_messages
            },
            timeout=60
        )
        data = response.json()
        return data["content"][0]["text"]
    
    def _call_openai(self, messages: list, max_tokens: int, api_key: str) -> str:
        """Call OpenAI API."""
        import httpx
        
        response = httpx.post(
            "https://api.openai.com/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            },
            json={
                "model": "gpt-4o-mini",
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": 0.1
            },
            timeout=60
        )
        data = response.json()
        return data["choices"][0]["message"]["content"]
    
    def _call(self, agent_role: str, user_message: str, max_tokens: int = 2048) -> AgentResponse:
        """
        Call the brain with a specific agent role.
        
        Returns AgentResponse with provenance_level=5 (LLM Hypothesis).
        Per Research OS spec, this is the LOWEST priority in the provenance hierarchy.
        """
        messages = [
            {"role": "system", "content": PROMPTS[agent_role]},
            {"role": "user", "content": user_message}
        ]
        
        try:
            if self.backend == "local":
                raw = self._generate_local(messages, max_tokens)
            else:
                raw = self._generate_api(messages, max_tokens)
            
            # Parse JSON from response
            # Try to extract JSON from the response (handle markdown code blocks)
            json_str = raw.strip()
            if json_str.startswith("```"):
                json_str = json_str.split("```")[1]
                if json_str.startswith("json"):
                    json_str = json_str[4:]
                json_str = json_str.strip()
            
            data = json.loads(json_str)
            
            return AgentResponse(
                success=True,
                data=data,
                raw_output=raw,
                provenance_level=5,  # LLM inference — must be human-verified
            )
        except json.JSONDecodeError as e:
            return AgentResponse(
                success=False,
                data={"error": f"Invalid JSON: {str(e)}", "raw": raw},
                raw_output=raw,
                provenance_level=5,
            )
        except Exception as e:
            return AgentResponse(
                success=False,
                data={"error": str(e)},
                raw_output="",
                provenance_level=5,
            )
    
    # ============================================================
    # Public API — one method per Research OS task
    # ============================================================
    
    def extract_claims(self, paper_text: str) -> AgentResponse:
        """
        Task 1: Extract structured claims from scientific text.
        Returns claims with epistemic tags, confidence scores, and parameters.
        
        Research OS Provenance: Level 5 (LLM Hypothesis)
        """
        return self._call("researcher", 
            f"Extract all scientific claims from the following paper excerpt:\n\n{paper_text}")
    
    def classify_epistemic(self, statement: str) -> AgentResponse:
        """
        Task 2: Classify the epistemic status of a scientific statement.
        Returns: Fact | Interpretation | Hypothesis | Conflict_Hypothesis
        
        Research OS Provenance: Level 5 (LLM Hypothesis)
        """
        return self._call("epistemic_classifier",
            f"Classify the epistemic status of this scientific statement:\n\n\"{statement}\"")
    
    def score_confidence(self, claim_text: str, journal: str, 
                         study_type: str, journal_tier: int) -> AgentResponse:
        """
        Task 3: Score confidence using the Research OS formula.
        confidence = evidence_strength × study_quality × journal_tier × completeness
        
        Uses FIXED-POINT arithmetic per Research OS Rule 5.
        """
        return self._call("confidence_scorer",
            f"Score the confidence of this claim:\n\n{claim_text}\n\n"
            f"Source: {journal}\nStudy type: {study_type}\nJournal tier: {journal_tier}")
    
    def detect_conflicts(self, claim_a: str, claim_b: str) -> AgentResponse:
        """
        Task 4: Detect contradictions between two claims.
        hypothesis_confidence is ALWAYS "low" — human review required.
        
        Research OS Rule: Agent hypotheses never auto-promote above Level 5.
        """
        return self._call("verifier",
            f"Analyze these two claims for contradictions:\n\n"
            f"Claim A: \"{claim_a}\"\n\nClaim B: \"{claim_b}\"")
    
    def decompose_query(self, question: str) -> AgentResponse:
        """
        Task 5: Decompose a broad research question into sub-queries.
        """
        return self._call("query_planner",
            f"Decompose this research question into specific sub-queries:\n\n\"{question}\"")
    
    def generate_decision(self, goal: str, gaps: list, 
                          low_confidence_claims: list) -> AgentResponse:
        """
        Task 6: Generate a Decision Object with information gain estimate.
        """
        user_msg = f"""Current research goal: {goal}

Knowledge gaps:
{chr(10).join('- ' + g for g in gaps)}

Low-confidence claims requiring resolution:
{chr(10).join('- ' + c for c in low_confidence_claims)}

Propose a Decision Object with the highest expected information gain."""
        
        return self._call("decision_generator", user_msg)


# ============================================================
# Convenience functions
# ============================================================

def create_brain(model_path: str = None, backend: str = "local") -> ResearchOSBrain:
    """Create a Research OS Brain instance."""
    return ResearchOSBrain(model_path=model_path, backend=backend)