mrmd-ai 0.1.0__tar.gz → 0.1.1__tar.gz

{mrmd_ai-0.1.0 → mrmd_ai-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mrmd-ai
-Version: 0.1.0
+Version: 0.1.1
 Summary: AI programs for MRMD editor - completions, fixes, and corrections
 Requires-Python: >=3.11
 Requires-Dist: dspy>=2.6

{mrmd_ai-0.1.0 → mrmd_ai-0.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mrmd-ai"
-version = "0.1.0"
+version = "0.1.1"
 description = "AI programs for MRMD editor - completions, fixes, and corrections"
 readme = "README.md"
 requires-python = ">=3.11"

mrmd_ai-0.1.1/src/mrmd_ai/juice.py

@@ -0,0 +1,673 @@
+"""
+Juice Level System for MRMD AI Programs.
+
+Juice levels control the quality/cost tradeoff of AI responses:
+- Level 0: Kimi K2 on Groq (fast, cheap, default)
+- Level 1: Claude Sonnet 4.5 (better quality)
+- Level 2: Gemini 3 Pro with thinking (deep reasoning)
+- Level 3: Claude Opus 4.5 with high thinking (maximum single-model quality)
+- Level 4: Multi-model merger (Grok 4 + Sonnet 4.5 + Gemini 3 + Opus 4.5, synthesized by Gemini 3)
+"""
+
+from enum import IntEnum
+from typing import Any, Callable
+from dataclasses import dataclass, field
+import dspy
+
+
+class JuiceLevel(IntEnum):
+    """Progressive quality levels for AI responses."""
+
+    # Fast & cheap - Kimi K2 on Groq
+    QUICK = 0
+
+    # Better quality - Sonnet 4.5
+    BALANCED = 1
+
+    # Deep reasoning - Gemini 3 with thinking
+    DEEP = 2
+
+    # Maximum single-model - Opus 4.5 with high thinking
+    MAXIMUM = 3
+
+    # Multi-model merger - all models synthesized
+    ULTIMATE = 4
+
+
+class ReasoningLevel(IntEnum):
+    """Independent reasoning/thinking budget control.
+
+    This is separate from JuiceLevel and controls how much "thinking"
+    the model does, independent of which model is selected.
+    """
+
+    # No extended thinking - fastest responses
+    OFF = 0
+
+    # Minimal reasoning
+    MINIMAL = 1
+
+    # Low reasoning effort
+    LOW = 2
+
+    # Medium reasoning effort
+    MEDIUM = 3
+
+    # High reasoning effort
+    HIGH = 4
+
+    # Maximum reasoning budget
+    MAXIMUM = 5
+
+
+# Map reasoning levels to thinking budgets and reasoning_effort values
+# For Anthropic: uses `thinking={"type": "enabled", "budget_tokens": X}`
+# For others: uses `reasoning_effort` ("low", "medium", "high")
+# Note: Anthropic requires max_tokens > thinking.budget_tokens
+REASONING_CONFIGS: dict[ReasoningLevel, dict] = {
+    ReasoningLevel.OFF: {
+        "budget_tokens": None,  # No thinking
+        "reasoning_effort": None,
+        "temperature": None,  # None means use model default
+    },
+    ReasoningLevel.MINIMAL: {
+        "budget_tokens": 1024,  # Minimum thinking budget
+        "reasoning_effort": "low",
+        "temperature": 1.0,  # Required for Anthropic extended thinking
+        "max_tokens": 4096,  # Must be > budget_tokens
+    },
+    ReasoningLevel.LOW: {
+        "budget_tokens": 4096,
+        "reasoning_effort": "low",
+        "temperature": 1.0,
+        "max_tokens": 8192,
+    },
+    ReasoningLevel.MEDIUM: {
+        "budget_tokens": 8192,
+        "reasoning_effort": "medium",
+        "temperature": 1.0,
+        "max_tokens": 16000,
+    },
+    ReasoningLevel.HIGH: {
+        "budget_tokens": 16384,
+        "reasoning_effort": "high",
+        "temperature": 1.0,
+        "max_tokens": 24000,
+    },
+    ReasoningLevel.MAXIMUM: {
+        "budget_tokens": 32768,  # Maximum thinking budget
+        "reasoning_effort": "high",
+        "temperature": 1.0,
+        "max_tokens": 48000,  # Must be > budget_tokens
+    },
+}
+
+
+REASONING_DESCRIPTIONS = {
+    ReasoningLevel.OFF: "Off - No extended thinking",
+    ReasoningLevel.MINIMAL: "Minimal - Light reasoning",
+    ReasoningLevel.LOW: "Low - Some reasoning",
+    ReasoningLevel.MEDIUM: "Medium - Moderate reasoning",
+    ReasoningLevel.HIGH: "High - Deep reasoning",
+    ReasoningLevel.MAXIMUM: "Maximum - Full reasoning budget",
+}
+
+
+@dataclass
+class ModelConfig:
+    """Configuration for a model at a specific juice level."""
+    model: str
+    temperature: float = 0.7
+    max_tokens: int = 4096
+    reasoning_effort: str | None = None
+    thinking: dict | None = None
+    supports_reasoning: bool = True  # Whether the model supports reasoning_effort
+    extra_kwargs: dict = field(default_factory=dict)
+
+    def to_lm_kwargs(self) -> dict:
+        """Convert to dspy.LM kwargs."""
+        kwargs = {
+            "model": self.model,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+            **self.extra_kwargs,
+        }
+        if self.reasoning_effort:
+            kwargs["reasoning_effort"] = self.reasoning_effort
+        if self.thinking:
+            kwargs["thinking"] = self.thinking
+        return kwargs
+
+
+# Model configurations for each juice level
+# supports_reasoning indicates if the model/provider supports reasoning_effort parameter
+JUICE_MODELS: dict[JuiceLevel, ModelConfig] = {
+    JuiceLevel.QUICK: ModelConfig(
+        model="groq/moonshotai/kimi-k2-instruct-0905",
+        temperature=0.7,
+        max_tokens=4096,
+        supports_reasoning=False,  # Groq does NOT support reasoning_effort
+    ),
+    JuiceLevel.BALANCED: ModelConfig(
+        model="anthropic/claude-sonnet-4-5",
+        temperature=0.7,
+        max_tokens=4096,
+        supports_reasoning=True,  # Anthropic supports reasoning_effort
+    ),
+    JuiceLevel.DEEP: ModelConfig(
+        model="gemini/gemini-3-pro-preview",
+        temperature=1.0,
+        max_tokens=16000,
+        reasoning_effort="high",
+        supports_reasoning=True,  # Gemini supports reasoning_effort
+    ),
+    JuiceLevel.MAXIMUM: ModelConfig(
+        model="anthropic/claude-opus-4-5",
+        temperature=1.0,
+        max_tokens=16000,
+        reasoning_effort="high",
+        supports_reasoning=True,  # Anthropic supports reasoning_effort
+    ),
+}
+
+# For ULTIMATE level, we use all 4 models with highest thinking
+# Grok 4, GPT-5.1, Gemini 3, Opus 4.5
+# NOTE: Anthropic requires temperature=1 when using extended thinking
+ULTIMATE_MODELS: list[ModelConfig] = [
+    ModelConfig(
+        model="openrouter/x-ai/grok-4",
+        temperature=0.7,
+        max_tokens=8192,
+        supports_reasoning=True,  # Grok 4 supports reasoning
+    ),
+    ModelConfig(
+        model="openai/gpt-5.2",
+        temperature=1.0,
+        max_tokens=16000,
+        reasoning_effort="high",
+        supports_reasoning=True,  # OpenAI supports reasoning
+    ),
+    ModelConfig(
+        model="gemini/gemini-3-pro-preview",
+        temperature=1.0,
+        max_tokens=16000,
+        reasoning_effort="high",
+        supports_reasoning=True,  # Gemini supports reasoning
+    ),
+    ModelConfig(
+        model="anthropic/claude-opus-4-5",
+        temperature=1.0,  # Must be 1 for extended thinking
+        max_tokens=16000,
+        reasoning_effort="high",
+        supports_reasoning=True,  # Anthropic supports reasoning
+    ),
+]
+
+# Synthesizer model for ULTIMATE level (Gemini 3 synthesizes all responses)
+SYNTHESIZER_MODEL = ModelConfig(
+    model="gemini/gemini-3-pro-preview",
+    temperature=0.7,
+    max_tokens=32000,
+    reasoning_effort="high",
+    supports_reasoning=True,
+)
+
+
+def get_lm(
+    juice: JuiceLevel | int = JuiceLevel.QUICK,
+    reasoning: ReasoningLevel | int | None = None
+) -> dspy.LM:
+    """Get a dspy.LM configured for the specified juice and reasoning levels.
+
+    Args:
+        juice: Juice level (0-3). Level 4 (ULTIMATE) requires special handling.
+        reasoning: Optional reasoning level (0-5). If None, uses juice level's default.
+
+    Returns:
+        Configured dspy.LM instance.
+    """
+    if isinstance(juice, int):
+        juice = JuiceLevel(juice)
+
+    if juice == JuiceLevel.ULTIMATE:
+        raise ValueError("ULTIMATE juice level requires multi-model merger. Use JuicedProgram instead.")
+
+    config = JUICE_MODELS[juice]
+    kwargs = config.to_lm_kwargs()
+
+    # Apply reasoning level overrides if specified AND model supports reasoning
+    if reasoning is not None and config.supports_reasoning:
+        if isinstance(reasoning, int):
+            reasoning = ReasoningLevel(reasoning)
+
+        # Skip if reasoning is OFF
+        if reasoning == ReasoningLevel.OFF:
+            # Remove any existing reasoning params
+            kwargs.pop("reasoning_effort", None)
+            kwargs.pop("thinking", None)
+            return dspy.LM(**kwargs)
+
+        reasoning_config = REASONING_CONFIGS[reasoning]
+        model = config.model.lower()
+
+        # Determine provider and use appropriate parameter format
+        is_anthropic = "anthropic/" in model or "claude" in model
+        is_gemini = "gemini" in model
+        is_openai = "openai/" in model or "gpt" in model
+
+        # Apply temperature (required for Anthropic extended thinking)
+        if reasoning_config.get("temperature") is not None:
+            kwargs["temperature"] = reasoning_config["temperature"]
+
+        # Apply max_tokens
+        if reasoning_config.get("max_tokens") is not None:
+            kwargs["max_tokens"] = reasoning_config["max_tokens"]
+
+        if is_anthropic:
+            # Anthropic uses explicit thinking parameter with budget_tokens
+            budget = reasoning_config.get("budget_tokens", 1024)
+            kwargs["thinking"] = {"type": "enabled", "budget_tokens": budget}
+            # Remove reasoning_effort if present (not used for thinking)
+            kwargs.pop("reasoning_effort", None)
+        else:
+            # Other providers use reasoning_effort
+            if reasoning_config["reasoning_effort"] is not None:
+                kwargs["reasoning_effort"] = reasoning_config["reasoning_effort"]
+
+    return dspy.LM(**kwargs)
+
+
+class SynthesizeResponses(dspy.Signature):
+    """Synthesize multiple AI model responses into one optimal final answer.
+
+    You are given responses from multiple AI models for the same task.
+    Your job is to create the BEST possible response by:
+    1. Identifying the strongest elements from each model's response
+    2. Resolving any contradictions (prefer the most accurate/well-reasoned answer)
+    3. Combining complementary insights that don't conflict
+    4. Maintaining the original format and style expected for the task
+    5. Being concise - don't add unnecessary elaboration
+
+    For grammar/spelling fixes: Pick the most correct version, don't over-correct.
+    For text completion: Choose the most natural, coherent continuation.
+    For code: Select the cleanest, most idiomatic solution.
+    For lists: You may combine unique items if appropriate.
+
+    Output ONLY the synthesized response - no explanations or meta-commentary.
+    """
+
+    original_input: str = dspy.InputField(desc="The original input/task that was given to all models")
+    model_responses: str = dspy.InputField(desc="Responses from multiple AI models, each labeled with model name")
+    synthesized_response: str = dspy.OutputField(desc="The single best response, synthesized from all model outputs. Output ONLY the response content.")
+
+
+class JuicedProgram:
+    """Wrapper that runs any DSPy program with configurable juice levels.
+
+    For levels 0-3, uses a single model with increasing capability.
+    For level 4 (ULTIMATE), runs all models in parallel and synthesizes.
+    """
+
+    def __init__(
+        self,
+        program: dspy.Module,
+        juice: JuiceLevel | int = JuiceLevel.QUICK,
+        reasoning: ReasoningLevel | int | None = None,
+        progress_callback: Callable[[str, dict], None] | None = None
+    ):
+        """Initialize a juiced program.
+
+        Args:
+            program: The DSPy program/module to wrap.
+            juice: Juice level (0-4).
+            reasoning: Optional reasoning level (0-5). If None, uses juice level's default.
+            progress_callback: Optional callback for progress events.
+                              Called with (event_type, data) where event_type is:
+                              - "status": General status update
+                              - "model_start": A model is starting (ultimate mode)
+                              - "model_complete": A model finished (ultimate mode)
+        """
+        self.program = program
+        self.juice = JuiceLevel(juice) if isinstance(juice, int) else juice
+        self.reasoning = ReasoningLevel(reasoning) if isinstance(reasoning, int) else reasoning
+        self.progress_callback = progress_callback
+
+    def _emit(self, event_type: str, data: dict):
+        """Emit a progress event if callback is set."""
+        if self.progress_callback:
+            self.progress_callback(event_type, data)
+
+    def __call__(self, **kwargs) -> Any:
+        """Run the program with the configured juice level."""
+        if self.juice == JuiceLevel.ULTIMATE:
+            return self._run_ultimate(**kwargs)
+        else:
+            return self._run_single(**kwargs)
+
+    def _run_single(self, **kwargs) -> Any:
+        """Run with a single model at the specified juice level."""
+        config = JUICE_MODELS[self.juice]
+        model_name = config.model.split("/")[-1]
+
+        reasoning_desc = ""
+        if self.reasoning is not None:
+            reasoning_desc = f" (reasoning={self.reasoning.name})"
+
+        self._emit("status", {
+            "step": "calling_model",
+            "model": model_name,
+            "model_full": config.model,
+            "reasoning_level": self.reasoning.value if self.reasoning else None,
+        })
+
+        lm = get_lm(self.juice, self.reasoning)
+        with dspy.context(lm=lm):
+            result = self.program(**kwargs)
+
+        self._emit("status", {
+            "step": "model_complete",
+            "model": model_name
+        })
+
+        return result
+
+    def _run_ultimate(self, **kwargs) -> Any:
+        """Run with all models in PARALLEL and merge results."""
+        from concurrent.futures import ThreadPoolExecutor, as_completed
+        import threading
+
+        # Track which models are running
+        model_names = [cfg.model.split("/")[-1] for cfg in ULTIMATE_MODELS]
+        models_status = {name: "pending" for name in model_names}
+        status_lock = threading.Lock()
+
+        self._emit("status", {
+            "step": "starting_multi_model",
+            "models": model_names,
+            "total": len(model_names),
+            "reasoning_level": self.reasoning.value if self.reasoning else None,
+        })
+
+        def run_model(config):
+            """Run a single model - called in parallel."""
+            lm_kwargs = config.to_lm_kwargs()
+
+            # Apply reasoning level overrides if specified AND model supports reasoning
+            if self.reasoning is not None and self.reasoning != ReasoningLevel.OFF and config.supports_reasoning:
+                reasoning_config = REASONING_CONFIGS[self.reasoning]
+                model = config.model.lower()
+
+                # Determine provider
+                is_anthropic = "anthropic/" in model or "claude" in model
+
+                # Apply temperature and max_tokens
+                if reasoning_config.get("temperature") is not None:
+                    lm_kwargs["temperature"] = reasoning_config["temperature"]
+                if reasoning_config.get("max_tokens") is not None:
+                    lm_kwargs["max_tokens"] = reasoning_config["max_tokens"]
+
+                if is_anthropic:
+                    # Anthropic uses thinking parameter with budget_tokens
+                    budget = reasoning_config.get("budget_tokens", 1024)
+                    lm_kwargs["thinking"] = {"type": "enabled", "budget_tokens": budget}
+                    lm_kwargs.pop("reasoning_effort", None)
+                else:
+                    # Other providers use reasoning_effort
+                    if reasoning_config["reasoning_effort"] is not None:
+                        lm_kwargs["reasoning_effort"] = reasoning_config["reasoning_effort"]
+
+            lm = dspy.LM(**lm_kwargs)
+            model_name = config.model.split("/")[-1]
+
+            # Emit model start
+            with status_lock:
+                models_status[model_name] = "running"
+            self._emit("model_start", {
+                "model": model_name,
+                "models_status": dict(models_status)
+            })
+
+            try:
+                with dspy.context(lm=lm):
+                    result = self.program(**kwargs)
+
+                # Extract response text from DSPy Prediction for streaming
+                response_data = {}
+                if hasattr(result, "_store") and result._store:
+                    response_data = dict(result._store)
+
+                # Emit model complete WITH the actual response
+                with status_lock:
+                    models_status[model_name] = "complete"
+                self._emit("model_complete", {
+                    "model": model_name,
+                    "success": True,
+                    "models_status": dict(models_status),
+                    "response": response_data,  # Include actual response!
+                })
+
+                return {"model": model_name, "result": result, "error": None}
+            except Exception as e:
+                # Emit model error
+                with status_lock:
+                    models_status[model_name] = "error"
+                self._emit("model_complete", {
+                    "model": model_name,
+                    "success": False,
+                    "error": str(e),
+                    "models_status": dict(models_status),
+                    "response": None,
+                })
+                return {"model": model_name, "result": None, "error": str(e)}
+
+        # Run all 4 models in parallel
+        model_results = []
+        with ThreadPoolExecutor(max_workers=4) as executor:
+            futures = [executor.submit(run_model, config) for config in ULTIMATE_MODELS]
+            for future in as_completed(futures):
+                model_results.append(future.result())
+
+        # Emit synthesizing status
+        self._emit("status", {
+            "step": "synthesizing",
+            "models_completed": len([r for r in model_results if r["result"] is not None])
+        })
+
+        # Merge results using AI synthesis
+        return self._merge_results(model_results, kwargs)
+
+    def _merge_results(self, model_results: list, original_input: dict) -> Any:
+        """Merge results from multiple models using AI synthesis.
+
+        Uses SYNTHESIZER_MODEL to intelligently combine responses from all models.
+        """
+        # Get successful results
+        successful = [r for r in model_results if r["result"] is not None]
+        if not successful:
+            # All failed - return error
+            errors = [r["error"] for r in model_results if r["error"]]
+            raise RuntimeError(f"All models failed: {errors}")
+
+        # If only one model succeeded, just return its result
+        if len(successful) == 1:
+            result = successful[0]["result"]
+            if hasattr(result, "_store"):
+                result._individual_responses = [{
+                    "model": successful[0]["model"],
+                    "response": str(result._store),
+                    "error": None
+                }]
+            return result
+
+        # Collect individual responses
+        individual_responses = []
+        model_outputs = {}  # model_name -> {field: value}
+
+        for r in model_results:
+            model_name = r["model"]
+            if r["result"] is not None and hasattr(r["result"], "_store"):
+                store = r["result"]._store
+                model_outputs[model_name] = dict(store)
+                # Get main output text for display
+                output_text = None
+                for key, value in store.items():
+                    if isinstance(value, str) and len(value) > 10:
+                        output_text = value
+                        break
+                individual_responses.append({
+                    "model": model_name,
+                    "response": output_text or str(store),
+                    "error": None
+                })
+            elif r["error"]:
+                individual_responses.append({
+                    "model": model_name,
+                    "response": None,
+                    "error": r["error"]
+                })
+
+        # Use first result as template for output fields
+        base_result = successful[0]["result"]
+        base_store = base_result._store if hasattr(base_result, "_store") else {}
+
+        # Format original input for synthesizer
+        input_text = self._format_input(original_input)
+
+        # Create synthesized result
+        merged = {}
+
+        # Configure synthesizer LM
+        synth_lm = dspy.LM(**SYNTHESIZER_MODEL.to_lm_kwargs())
+
+        # Synthesize each output field
+        for field_name, base_value in base_store.items():
+            # Collect this field's values from all models
+            field_values = {}
+            for model_name, outputs in model_outputs.items():
+                if field_name in outputs:
+                    field_values[model_name] = outputs[field_name]
+
+            if not field_values:
+                merged[field_name] = base_value
+                continue
+
+            # Check if it's a list field (like synonyms)
+            if isinstance(base_value, list):
+                # For lists, combine unique values from all models
+                combined = []
+                seen = set()
+                for model_name, values in field_values.items():
+                    if isinstance(values, list):
+                        for item in values:
+                            # Get hashable key for deduplication
+                            # Pydantic models aren't hashable, so convert to JSON
+                            try:
+                                if hasattr(item, 'model_dump_json'):
+                                    # Pydantic v2 model
+                                    item_key = item.model_dump_json()
+                                elif hasattr(item, 'json'):
+                                    # Pydantic v1 model
+                                    item_key = item.json()
+                                else:
+                                    # Regular hashable item
+                                    item_key = item
+                            except TypeError:
+                                # Fallback: convert to string representation
+                                item_key = str(item)
+
+                            if item_key not in seen:
+                                combined.append(item)
+                                seen.add(item_key)
+                merged[field_name] = combined
+            else:
+                # For string/text fields, use AI synthesis
+                responses_text = "\n\n".join([
+                    f"=== {model_name} ===\n{value}"
+                    for model_name, value in field_values.items()
+                ])
+
+                self._emit("status", {
+                    "step": "synthesizing_field",
+                    "field": field_name,
+                    "model": SYNTHESIZER_MODEL.model.split("/")[-1]
+                })
+
+                try:
+                    with dspy.context(lm=synth_lm):
+                        predictor = dspy.Predict(SynthesizeResponses)
+                        synth_result = predictor(
+                            original_input=input_text,
+                            model_responses=responses_text
+                        )
+                        merged[field_name] = synth_result.synthesized_response
+                except Exception as e:
+                    # Fallback to first model's response on synthesis error
+                    print(f"[Synthesis] Error synthesizing {field_name}: {e}")
+                    merged[field_name] = base_value
+
+        # Return a result object with merged data
+        class MergedResult:
+            pass
+
+        result = MergedResult()
+        for key, value in merged.items():
+            setattr(result, key, value)
+        result._store = merged  # For extract_result in server.py
+        result._individual_responses = individual_responses  # For UI display
+        result._synthesized = True  # Mark as AI-synthesized
+
+        return result
+
+    def _format_input(self, kwargs: dict) -> str:
+        """Format input kwargs as a readable string."""
+        parts = []
+        for key, value in kwargs.items():
+            parts.append(f"{key}: {value}")
+        return "\n".join(parts)
+
+
+def juiced(juice: JuiceLevel | int = JuiceLevel.QUICK):
+    """Decorator to run a DSPy program with a specific juice level.
+
+    Usage:
+        @juiced(JuiceLevel.DEEP)
+        def my_program():
+            return dspy.ChainOfThought(MySignature)
+    """
+    def decorator(func: Callable) -> Callable:
+        def wrapper(*args, **kwargs):
+            program = func(*args, **kwargs)
+            return JuicedProgram(program, juice)
+        return wrapper
+    return decorator
+
+
+def run_with_juice(
+    program: dspy.Module,
+    juice: JuiceLevel | int,
+    reasoning: ReasoningLevel | int | None = None,
+    **kwargs
+) -> Any:
+    """Convenience function to run a program with a specific juice level.
+
+    Args:
+        program: The DSPy program to run.
+        juice: Juice level (0-4).
+        reasoning: Optional reasoning level (0-5). If None, uses juice level's default.
+        **kwargs: Arguments to pass to the program.
+
+    Returns:
+        The program result.
+    """
+    juiced_program = JuicedProgram(program, juice, reasoning=reasoning)
+    return juiced_program(**kwargs)
+
+
+# Juice level descriptions for CLI/UI
+JUICE_DESCRIPTIONS = {
+    JuiceLevel.QUICK: "Quick (Kimi K2) - Fast & cheap",
+    JuiceLevel.BALANCED: "Balanced (Sonnet 4.5) - Good quality",
+    JuiceLevel.DEEP: "Deep (Gemini 3 thinking) - Thorough reasoning",
+    JuiceLevel.MAXIMUM: "Maximum (Opus 4.5 thinking) - Best single model",
+    JuiceLevel.ULTIMATE: "Ultimate (Multi-model merger) - All models synthesized",
+}