mrmd-ai 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

mrmd_ai/juice.py

@@ -34,6 +34,85 @@ class JuiceLevel(IntEnum):
     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."""
@@ -42,6 +121,7 @@ class ModelConfig:
     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:
@@ -60,28 +140,33 @@ class ModelConfig:
 
 
 # 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
     ),
 }
 
@@ -93,24 +178,28 @@ ULTIMATE_MODELS: list[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.1",
+        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
     ),
 ]
 
@@ -120,14 +209,19 @@ SYNTHESIZER_MODEL = ModelConfig(
     temperature=0.7,
     max_tokens=32000,
     reasoning_effort="high",
+    supports_reasoning=True,
 )
 
 
-def get_lm(juice: JuiceLevel | int = JuiceLevel.QUICK) -> dspy.LM:
-    """Get a dspy.LM configured for the specified juice level.
+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.
@@ -139,20 +233,72 @@ def get_lm(juice: JuiceLevel | int = JuiceLevel.QUICK) -> dspy.LM:
         raise ValueError("ULTIMATE juice level requires multi-model merger. Use JuicedProgram instead.")
 
     config = JUICE_MODELS[juice]
-    return dspy.LM(**config.to_lm_kwargs())
+    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 an optimal final answer.
 
-    You are given the original input and responses from multiple AI models.
-    Analyze all responses, identify the best insights from each, resolve
-    any contradictions, and produce the ultimate synthesized response.
+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/question")
-    model_responses: str = dspy.InputField(desc="Responses from multiple AI models, labeled by model name")
-    synthesized_response: str = dspy.OutputField(desc="The optimal synthesized response combining the best from all models")
+    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:
@@ -166,6 +312,7 @@ class JuicedProgram:
         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.
@@ -173,6 +320,7 @@ class JuicedProgram:
         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
@@ -181,6 +329,7 @@ class JuicedProgram:
         """
         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):
@@ -200,13 +349,18 @@ class JuicedProgram:
         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
+            "model_full": config.model,
+            "reasoning_level": self.reasoning.value if self.reasoning else None,
         })
 
-        lm = get_lm(self.juice)
+        lm = get_lm(self.juice, self.reasoning)
         with dspy.context(lm=lm):
             result = self.program(**kwargs)
 
@@ -230,12 +384,39 @@ class JuicedProgram:
         self._emit("status", {
             "step": "starting_multi_model",
             "models": model_names,
-            "total": len(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 = dspy.LM(**config.to_lm_kwargs())
+            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
@@ -250,13 +431,19 @@ class JuicedProgram:
                 with dspy.context(lm=lm):
                     result = self.program(**kwargs)
 
-                # Emit model complete
+                # 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)
+                    "models_status": dict(models_status),
+                    "response": response_data,  # Include actual response!
                 })
 
                 return {"model": model_name, "result": result, "error": None}
@@ -268,7 +455,8 @@ class JuicedProgram:
                     "model": model_name,
                     "success": False,
                     "error": str(e),
-                    "models_status": dict(models_status)
+                    "models_status": dict(models_status),
+                    "response": None,
                 })
                 return {"model": model_name, "result": None, "error": str(e)}
 
@@ -285,15 +473,13 @@ class JuicedProgram:
             "models_completed": len([r for r in model_results if r["result"] is not None])
         })
 
-        # Merge results - combine outputs from all successful models
-        return self._merge_results(model_results)
+        # Merge results using AI synthesis
+        return self._merge_results(model_results, kwargs)
 
-    def _merge_results(self, model_results: list) -> Any:
-        """Merge results from multiple models into a single response.
+    def _merge_results(self, model_results: list, original_input: dict) -> Any:
+        """Merge results from multiple models using AI synthesis.
 
-        For list fields (like synonyms), combines unique values from all models.
-        For string fields, uses the first successful result.
-        Also includes individual model responses for transparency.
+        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]
@@ -302,23 +488,27 @@ class JuicedProgram:
             errors = [r["error"] for r in model_results if r["error"]]
             raise RuntimeError(f"All models failed: {errors}")
 
-        # Use first successful result as base
-        base_result = successful[0]["result"]
-
-        # Get the _store dict from the result (DSPy stores outputs there)
-        if hasattr(base_result, "_store"):
-            merged = dict(base_result._store)
-        else:
-            merged = {}
+        # 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 for display
+        # 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"):
-                # Extract the main output field (usually 'response', 'completion', etc.)
                 store = r["result"]._store
-                # Get the first string output field
+                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:
@@ -336,26 +526,86 @@ class JuicedProgram:
                     "error": r["error"]
                 })
 
-        # Merge fields from other models
-        for r in successful[1:]:
-            result = r["result"]
-            if hasattr(result, "_store"):
-                store = result._store
-                for key, value in store.items():
-                    if key in merged:
-                        # Merge lists by combining unique values
-                        if isinstance(value, list) and isinstance(merged[key], list):
-                            # Combine and dedupe while preserving order
-                            seen = set(merged[key])
-                            for item in value:
-                                if item not in seen:
-                                    merged[key].append(item)
-                                    seen.add(item)
-                        # For strings, keep the first (base) value
-                    else:
-                        merged[key] = value
-
-        # Return a simple object with the merged data + individual responses
+        # 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
 
@@ -364,6 +614,7 @@ class JuicedProgram:
             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
 
@@ -391,18 +642,24 @@ def juiced(juice: JuiceLevel | int = JuiceLevel.QUICK):
     return decorator
 
 
-def run_with_juice(program: dspy.Module, juice: JuiceLevel | int, **kwargs) -> Any:
+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)
+    juiced_program = JuicedProgram(program, juice, reasoning=reasoning)
     return juiced_program(**kwargs)
 
 

mrmd_ai/modules/__init__.py

@@ -38,6 +38,12 @@ from .document import (
 from .notebook import (
     NotebookNamePredict,
 )
+from .edit import (
+    EditAtCursorPredict,
+    AddressCommentPredict,
+    AddressAllCommentsPredict,
+    AddressNearbyCommentPredict,
+)
 
 __all__ = [
     # Finish programs
@@ -71,4 +77,9 @@ __all__ = [
     "DocumentAnalysisPredict",
     # Notebook programs
     "NotebookNamePredict",
+    # Edit programs (Ctrl-K and comments)
+    "EditAtCursorPredict",
+    "AddressCommentPredict",
+    "AddressAllCommentsPredict",
+    "AddressNearbyCommentPredict",
 ]

mrmd_ai/modules/edit.py

@@ -0,0 +1,102 @@
+"""Edit modules for cursor-based editing and comment processing."""
+
+import dspy
+from typing import List, Optional
+from ..signatures.edit import (
+    Edit,
+    CommentInfo,
+    EditAtCursorSignature,
+    AddressCommentSignature,
+    AddressAllCommentsSignature,
+    AddressNearbyCommentSignature,
+)
+
+
+class EditAtCursorPredict(dspy.Module):
+    """Execute user instructions via precise find/replace edits."""
+
+    def __init__(self):
+        super().__init__()
+        self.predict = dspy.Predict(EditAtCursorSignature)
+
+    def forward(
+        self,
+        text_before: str,
+        text_after: str,
+        selection: str,
+        full_document: str,
+        instruction: str,
+    ):
+        return self.predict(
+            text_before=text_before,
+            text_after=text_after,
+            selection=selection,
+            full_document=full_document,
+            instruction=instruction,
+        )
+
+
+class AddressCommentPredict(dspy.Module):
+    """Address a single comment embedded in the document."""
+
+    def __init__(self):
+        super().__init__()
+        self.predict = dspy.Predict(AddressCommentSignature)
+
+    def forward(
+        self,
+        full_document: str,
+        comment_text: str,
+        comment_context_before: str,
+        comment_context_after: str,
+        comment_raw: str,
+    ):
+        return self.predict(
+            full_document=full_document,
+            comment_text=comment_text,
+            comment_context_before=comment_context_before,
+            comment_context_after=comment_context_after,
+            comment_raw=comment_raw,
+        )
+
+
+class AddressAllCommentsPredict(dspy.Module):
+    """Address all comments in a document."""
+
+    def __init__(self):
+        super().__init__()
+        self.predict = dspy.Predict(AddressAllCommentsSignature)
+
+    def forward(
+        self,
+        full_document: str,
+        comments: List[CommentInfo],
+    ):
+        return self.predict(
+            full_document=full_document,
+            comments=comments,
+        )
+
+
+class AddressNearbyCommentPredict(dspy.Module):
+    """Address the comment nearest to the cursor."""
+
+    def __init__(self):
+        super().__init__()
+        self.predict = dspy.Predict(AddressNearbyCommentSignature)
+
+    def forward(
+        self,
+        full_document: str,
+        cursor_context_before: str,
+        cursor_context_after: str,
+        nearby_comment: CommentInfo,
+        nearby_comment_raw: str,
+    ):
+        return self.predict(
+            full_document=full_document,
+            cursor_context_before=cursor_context_before,
+            cursor_context_after=cursor_context_after,
+            nearby_comment=nearby_comment,
+            nearby_comment_raw=nearby_comment_raw,
+        )

mrmd_ai/server.py

@@ -23,7 +23,7 @@ import json
 # Thread pool for running blocking DSPy calls
 _executor = ThreadPoolExecutor(max_workers=10)
 
-from .juice import JuiceLevel, JuicedProgram, get_lm, JUICE_MODELS
+from .juice import JuiceLevel, ReasoningLevel, JuicedProgram, get_lm, JUICE_MODELS, REASONING_DESCRIPTIONS
 from .modules import (
     # Finish
     FinishSentencePredict,
@@ -56,6 +56,11 @@ from .modules import (
     DocumentAnalysisPredict,
     # Notebook
     NotebookNamePredict,
+    # Edit (Ctrl-K and comments)
+    EditAtCursorPredict,
+    AddressCommentPredict,
+    AddressAllCommentsPredict,
+    AddressNearbyCommentPredict,
 )
 
 
@@ -92,21 +97,26 @@ PROGRAMS = {
     "DocumentAnalysisPredict": DocumentAnalysisPredict,
     # Notebook
     "NotebookNamePredict": NotebookNamePredict,
+    # Edit (Ctrl-K and comments)
+    "EditAtCursorPredict": EditAtCursorPredict,
+    "AddressCommentPredict": AddressCommentPredict,
+    "AddressAllCommentsPredict": AddressAllCommentsPredict,
+    "AddressNearbyCommentPredict": AddressNearbyCommentPredict,
 }
 
-# Cached program instances per juice level
-_program_cache: dict[tuple[str, int], JuicedProgram] = {}
+# Cached program instances per juice level and reasoning level
+_program_cache: dict[tuple[str, int, int | None], JuicedProgram] = {}
 
 
-def get_program(name: str, juice: int = 0) -> JuicedProgram:
-    """Get a JuicedProgram instance for the given program and juice level."""
-    cache_key = (name, juice)
+def get_program(name: str, juice: int = 0, reasoning: int | None = None) -> JuicedProgram:
+    """Get a JuicedProgram instance for the given program, juice level, and reasoning level."""
+    cache_key = (name, juice, reasoning)
     if cache_key not in _program_cache:
         if name not in PROGRAMS:
             raise ValueError(f"Unknown program: {name}")
         program_class = PROGRAMS[name]
         program = program_class()
-        _program_cache[cache_key] = JuicedProgram(program, juice=juice)
+        _program_cache[cache_key] = JuicedProgram(program, juice=juice, reasoning=reasoning)
     return _program_cache[cache_key]
 
 
@@ -153,12 +163,31 @@ async def list_programs():
 
 @app.get("/juice")
 async def get_juice_levels():
-    """Get available juice levels."""
-    from .juice import JUICE_DESCRIPTIONS
+    """Get available juice levels with their capabilities."""
+    from .juice import JUICE_DESCRIPTIONS, JUICE_MODELS, JuiceLevel
+    levels = []
+    for level, desc in JUICE_DESCRIPTIONS.items():
+        level_info = {
+            "level": level.value,
+            "description": desc,
+        }
+        # Add supports_reasoning for non-ULTIMATE levels
+        if level != JuiceLevel.ULTIMATE and level in JUICE_MODELS:
+            level_info["supports_reasoning"] = JUICE_MODELS[level].supports_reasoning
+        else:
+            # ULTIMATE level supports reasoning (all its sub-models do)
+            level_info["supports_reasoning"] = True
+        levels.append(level_info)
+    return {"levels": levels}
+
+
+@app.get("/reasoning")
+async def get_reasoning_levels():
+    """Get available reasoning levels."""
     return {
         "levels": [
             {"level": level.value, "description": desc}
-            for level, desc in JUICE_DESCRIPTIONS.items()
+            for level, desc in REASONING_DESCRIPTIONS.items()
         ]
     }
 
@@ -184,6 +213,7 @@ def extract_result(prediction: Any) -> dict:
             "reformatted_text", "text_to_replace", "replacement",
             "response", "summary", "analysis",  # Document-level fields
             "code",  # ProgramCodePredict output
+            "edits",  # EditAtCursor and AddressComment outputs
         ]
 
         for field in output_fields:
@@ -210,6 +240,16 @@ async def run_program(program_name: str, request: Request):
     except ValueError:
         juice_level = 0
 
+    # Get reasoning level from header (optional)
+    reasoning_header = request.headers.get("X-Reasoning-Level")
+    reasoning_level = None
+    if reasoning_header is not None:
+        try:
+            reasoning_level = int(reasoning_header)
+            reasoning_level = max(0, min(5, reasoning_level))  # Clamp to 0-5
+        except ValueError:
+            reasoning_level = None
+
     # Get request body
     try:
         params = await request.json()
@@ -218,14 +258,17 @@ async def run_program(program_name: str, request: Request):
 
     # Get program
     try:
-        juiced_program = get_program(program_name, juice_level)
+        juiced_program = get_program(program_name, juice_level, reasoning_level)
     except ValueError as e:
         raise HTTPException(status_code=404, detail=str(e))
 
     # Log the call and get model info
-    from .juice import JUICE_DESCRIPTIONS, JUICE_MODELS, ULTIMATE_MODELS, JuiceLevel
+    from .juice import JUICE_DESCRIPTIONS, JUICE_MODELS, ULTIMATE_MODELS, JuiceLevel, ReasoningLevel
     juice_desc = JUICE_DESCRIPTIONS.get(JuiceLevel(juice_level), f"Level {juice_level}")
-    print(f"[AI] {program_name} @ {juice_desc}", flush=True)
+    reasoning_desc = ""
+    if reasoning_level is not None:
+        reasoning_desc = f" | {REASONING_DESCRIPTIONS.get(ReasoningLevel(reasoning_level), f'Reasoning {reasoning_level}')}"
+    print(f"[AI] {program_name} @ {juice_desc}{reasoning_desc}", flush=True)
 
     # Get the model name for this juice level
     if juice_level == JuiceLevel.ULTIMATE:
@@ -245,16 +288,35 @@ async def run_program(program_name: str, request: Request):
         # Add model metadata to response
         response["_model"] = model_name
         response["_juice_level"] = juice_level
-        return response
+        response["_reasoning_level"] = reasoning_level
+        # Serialize any Pydantic models to dicts for JSON compatibility
+        return serialize_for_json(response)
     except Exception as e:
         import traceback
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))
 
 
+def serialize_for_json(obj):
+    """Recursively convert Pydantic models and other objects to JSON-serializable form."""
+    if hasattr(obj, 'model_dump'):
+        # Pydantic v2 model
+        return obj.model_dump()
+    elif hasattr(obj, 'dict'):
+        # Pydantic v1 model
+        return obj.dict()
+    elif isinstance(obj, dict):
+        return {k: serialize_for_json(v) for k, v in obj.items()}
+    elif isinstance(obj, (list, tuple)):
+        return [serialize_for_json(item) for item in obj]
+    else:
+        return obj
+
+
 def sse_event(event: str, data: dict) -> str:
     """Format a Server-Sent Event."""
-    return f"event: {event}\ndata: {json.dumps(data)}\n\n"
+    serialized = serialize_for_json(data)
+    return f"event: {event}\ndata: {json.dumps(serialized)}\n\n"
 
 
 @app.post("/{program_name}/stream")
@@ -275,6 +337,16 @@ async def run_program_stream(program_name: str, request: Request):
     except ValueError:
         juice_level = 0
 
+    # Get reasoning level from header (optional)
+    reasoning_header = request.headers.get("X-Reasoning-Level")
+    reasoning_level = None
+    if reasoning_header is not None:
+        try:
+            reasoning_level = int(reasoning_header)
+            reasoning_level = max(0, min(5, reasoning_level))  # Clamp to 0-5
+        except ValueError:
+            reasoning_level = None
+
     # Get request body
     try:
         params = await request.json()
@@ -286,9 +358,12 @@ async def run_program_stream(program_name: str, request: Request):
         raise HTTPException(status_code=404, detail=f"Unknown program: {program_name}")
 
     # Get model info
-    from .juice import JUICE_DESCRIPTIONS, JUICE_MODELS, ULTIMATE_MODELS, JuiceLevel, JuicedProgram
+    from .juice import JUICE_DESCRIPTIONS, JUICE_MODELS, ULTIMATE_MODELS, JuiceLevel, ReasoningLevel, JuicedProgram
     juice_desc = JUICE_DESCRIPTIONS.get(JuiceLevel(juice_level), f"Level {juice_level}")
-    print(f"[AI Stream] {program_name} @ {juice_desc}", flush=True)
+    reasoning_desc = ""
+    if reasoning_level is not None:
+        reasoning_desc = f" | {REASONING_DESCRIPTIONS.get(ReasoningLevel(reasoning_level), f'Reasoning {reasoning_level}')}"
+    print(f"[AI Stream] {program_name} @ {juice_desc}{reasoning_desc}", flush=True)
 
     # Get model name(s) for display
     if juice_level == JuiceLevel.ULTIMATE:
@@ -318,14 +393,16 @@ async def run_program_stream(program_name: str, request: Request):
                 # Create program with progress callback
                 program_class = PROGRAMS[program_name]
                 program = program_class()
-                juiced = JuicedProgram(program, juice=juice_level, progress_callback=progress_callback)
+                juiced = JuicedProgram(program, juice=juice_level, reasoning=reasoning_level, progress_callback=progress_callback)
 
                 # Emit starting event
                 progress_callback("status", {
                     "step": "starting",
                     "model": model_name,
                     "juice_level": juice_level,
-                    "juice_name": juice_desc
+                    "juice_name": juice_desc,
+                    "reasoning_level": reasoning_level,
+                    "reasoning_name": reasoning_desc.strip(" |") if reasoning_desc else None,
                 })
 
                 # Run the program
@@ -357,6 +434,7 @@ async def run_program_stream(program_name: str, request: Request):
                         response = extract_result(result_holder["result"])
                         response["_model"] = model_name
                         response["_juice_level"] = juice_level
+                        response["_reasoning_level"] = reasoning_level
                         yield sse_event("result", response)
                     break
 

mrmd_ai/signatures/__init__.py

@@ -14,6 +14,14 @@ from .correct import (
     CorrectAndFinishLineSignature,
     CorrectAndFinishSectionSignature,
 )
+from .edit import (
+    Edit,
+    CommentInfo,
+    EditAtCursorSignature,
+    AddressCommentSignature,
+    AddressAllCommentsSignature,
+    AddressNearbyCommentSignature,
+)
 
 __all__ = [
     "FinishSentenceSignature",
@@ -24,4 +32,11 @@ __all__ = [
     "FixTranscriptionSignature",
     "CorrectAndFinishLineSignature",
     "CorrectAndFinishSectionSignature",
+    # Edit signatures
+    "Edit",
+    "CommentInfo",
+    "EditAtCursorSignature",
+    "AddressCommentSignature",
+    "AddressAllCommentsSignature",
+    "AddressNearbyCommentSignature",
 ]

mrmd_ai/signatures/edit.py

@@ -0,0 +1,173 @@
+"""Signature definitions for cursor-based editing and comment processing."""
+
+import dspy
+from pydantic import BaseModel, Field
+from typing import List, Optional
+
+
+class Edit(BaseModel):
+    """A single find/replace edit operation.
+
+    For insertions at cursor, use find="" and the text will be inserted
+    at the cursor position.
+    """
+    find: str = Field(
+        description="Exact text to find in document. Use empty string for insertion at cursor."
+    )
+    replace: str = Field(
+        description="Text to replace the found text with, or text to insert if find is empty."
+    )
+
+
+class CommentInfo(BaseModel):
+    """Information about a comment in the document."""
+    text: str = Field(description="The comment text content")
+    context_before: str = Field(description="Text immediately before the comment")
+    context_after: str = Field(description="Text immediately after the comment")
+
+
+class EditAtCursorSignature(dspy.Signature):
+    """
+    Execute a user instruction by generating precise find/replace edits.
+
+    You are given the cursor context and a natural language instruction.
+    Generate a list of edits that implement the instruction.
+
+    CRITICAL RULES:
+    1. Each edit has `find` (exact text to locate) and `replace` (replacement text)
+    2. For INSERTIONS at cursor: use find="" - the replace text will be inserted at cursor
+    3. For MODIFICATIONS: find must match the EXACT text in the document (character-for-character)
+    4. find strings must be UNIQUE enough to match only the intended location
+    5. Include surrounding context in find to ensure uniqueness (e.g., "def process_data(items)" not just "process_data")
+    6. Edits are applied in order - earlier edits may shift positions of later ones
+
+    Examples:
+    - Instruction: "add a docstring" → find the function definition, replace with definition + docstring
+    - Instruction: "rename x to count" → find=" x " (with spaces), replace=" count "
+    - Instruction: "insert a comment here" → find="", replace="# comment\\n"
+    - Instruction: "delete this line" → find="the line content\\n", replace=""
+
+    When the user has selected text, that text is provided in `selection`.
+    Prefer to operate on the selection when it's relevant to the instruction.
+    """
+
+    text_before: str = dspy.InputField(
+        desc="Text immediately before the cursor (up to 500 characters for context)"
+    )
+    text_after: str = dspy.InputField(
+        desc="Text immediately after the cursor (up to 500 characters for context)"
+    )
+    selection: str = dspy.InputField(
+        desc="Currently selected text, or empty string if no selection"
+    )
+    full_document: str = dspy.InputField(
+        desc="The complete document content for full context"
+    )
+    instruction: str = dspy.InputField(
+        desc="User's natural language instruction for what to do"
+    )
+
+    edits: List[Edit] = dspy.OutputField(
+        desc="List of find/replace edits to apply. Order matters - applied sequentially."
+    )
+
+
+class AddressCommentSignature(dspy.Signature):
+    """
+    Address a single comment/instruction embedded in the document.
+
+    Comments are marked with <!--! comment text !--> syntax.
+    The comment contains instructions or notes that should be addressed.
+    Generate edits that fulfill the comment's request.
+
+    After addressing, you may optionally remove the comment marker itself.
+
+    Guidelines:
+    - Read the comment carefully to understand what's requested
+    - Look at the surrounding context to understand where changes should go
+    - Generate precise edits that address the comment
+    - If the comment asks for something that's already done, return empty edits
+    - Consider removing the comment after addressing it (include that as an edit)
+    """
+
+    full_document: str = dspy.InputField(
+        desc="The complete document content"
+    )
+    comment_text: str = dspy.InputField(
+        desc="The text content of the comment (without the <!--! !--> markers)"
+    )
+    comment_context_before: str = dspy.InputField(
+        desc="Text immediately before the comment marker"
+    )
+    comment_context_after: str = dspy.InputField(
+        desc="Text immediately after the comment marker"
+    )
+    comment_raw: str = dspy.InputField(
+        desc="The full raw comment including markers (e.g., '<!--! add error handling !-->')"
+    )
+
+    edits: List[Edit] = dspy.OutputField(
+        desc="List of find/replace edits to address the comment"
+    )
+
+
+class AddressAllCommentsSignature(dspy.Signature):
+    """
+    Address ALL comments/instructions in a document.
+
+    Scan the document for all <!--! ... !--> comment markers and generate
+    edits that address each one.
+
+    Guidelines:
+    - Process comments in document order (top to bottom)
+    - Each comment should be addressed appropriately
+    - Comments that conflict should be resolved sensibly
+    - After addressing, remove the comment markers
+    - Return all edits as a single list (they'll be applied in order)
+    """
+
+    full_document: str = dspy.InputField(
+        desc="The complete document content with embedded comments"
+    )
+    comments: List[CommentInfo] = dspy.InputField(
+        desc="List of all comments found in the document with their context"
+    )
+
+    edits: List[Edit] = dspy.OutputField(
+        desc="List of all find/replace edits to address all comments"
+    )
+
+
+class AddressNearbyCommentSignature(dspy.Signature):
+    """
+    Address the comment nearest to the cursor position.
+
+    Find the comment that's closest to where the user's cursor is and
+    generate edits to address that specific comment.
+
+    Guidelines:
+    - Focus only on the comment nearest to the cursor
+    - Use the cursor context to identify which comment is relevant
+    - Generate edits that address that comment
+    - Optionally remove the comment marker after addressing
+    """
+
+    full_document: str = dspy.InputField(
+        desc="The complete document content"
+    )
+    cursor_context_before: str = dspy.InputField(
+        desc="Text before the cursor position"
+    )
+    cursor_context_after: str = dspy.InputField(
+        desc="Text after the cursor position"
+    )
+    nearby_comment: CommentInfo = dspy.InputField(
+        desc="The comment closest to the cursor"
+    )
+    nearby_comment_raw: str = dspy.InputField(
+        desc="The full raw comment including markers"
+    )
+
+    edits: List[Edit] = dspy.OutputField(
+        desc="List of find/replace edits to address the nearby comment"
+    )

{mrmd_ai-0.1.0.dist-info → mrmd_ai-0.1.1.dist-info}/METADATA

@@ -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.dist-info → mrmd_ai-0.1.1.dist-info}/RECORD

@@ -1,26 +1,28 @@
 mrmd_ai/__init__.py,sha256=RocuoOEEJCUNinxqsXWjm_bIimCu8aDCDha-G5pU4fU,105
-mrmd_ai/juice.py,sha256=0_h7_M89IDhSt4iohdz0J5StcFXOaLF33d3fTxdBO-o,14442
-mrmd_ai/server.py,sha256=Jgl-bsQYbBMtv_tGu_2FwSN9BxgFpt-fmb9_pjIcqNw,14416
+mrmd_ai/juice.py,sha256=pCXFgAv_GZOYmIrhcsn-6qeqfolcexwbp3mr-Z34NQk,25235
+mrmd_ai/server.py,sha256=jQSX12LK0RiFSMbHp81zaDirbWD14CNT9S5lGPQKRNc,17961
 mrmd_ai/metrics/__init__.py,sha256=6BngKqh0a09phOzUdYeWjhsUXznCIx5jEjrEt7DIDu4,62
-mrmd_ai/modules/__init__.py,sha256=4V4IzurDZs_nhBTGlHyXco24iwEkum58kRNrKfm4IjQ,1793
+mrmd_ai/modules/__init__.py,sha256=OO-5alsFmR1bLtsvxAJ2wNO5nOtHRmWChTexEI6Y9xU,2097
 mrmd_ai/modules/code.py,sha256=8cK6LF0ZTrSp2srt1ltvUnmYjPlt5NZTj1yWctpJ7j0,4099
 mrmd_ai/modules/correct.py,sha256=TWnE1HD_Ip7xZ5yQwJi1n01tNXgBtYNvkTK--kAknak,1478
 mrmd_ai/modules/document.py,sha256=o6iLR2amscn-DHZ95JFQuEwhaj8cLs4bBISf7G9cT9Y,1106
+mrmd_ai/modules/edit.py,sha256=lGa0tNB7d9tRG4rtdQ6uNIi3lzSbC588c3r86ccPXZc,2736
 mrmd_ai/modules/finish.py,sha256=VtyE-45-8iM6iWjNg57wRWL0Ln3mSF0RcZq8CO5CoSk,2638
 mrmd_ai/modules/fix.py,sha256=fb4flKWyyyelheigeb1sI0dixd2scL9HZX-0_M_Uh-o,1506
 mrmd_ai/modules/notebook.py,sha256=w8Dg-NKVL6_kPOKkvb84kGrwgv5zHxvFNWBtXHLHww8,477
 mrmd_ai/modules/text.py,sha256=9MCO__EDalwi-iFf__sd8t7orsUy7WiBBf6Lp4bxxGE,2010
 mrmd_ai/optimizers/__init__.py,sha256=Ay6ZrQu8fLQaG7-dl6hTMruQY5AdGOT_YnlRhhZGgag,60
-mrmd_ai/signatures/__init__.py,sha256=wWT2D8beisIpYMPxN4j4JncQmUqcamD6v5BTKp2zFWc,655
+mrmd_ai/signatures/__init__.py,sha256=UoIplXOXxPidkaXi9u-7l2LziDJX2MN4PmwtXVsC04U,1013
 mrmd_ai/signatures/code.py,sha256=zBM_Nl2NImfOw49fVWCGlXcE_sm8xgWCN1ksDbEa6e8,11245
 mrmd_ai/signatures/correct.py,sha256=tIhYCONgGhuTV0eJCiLSXcGZSAEi06XY35ommtTTsRE,2920
 mrmd_ai/signatures/document.py,sha256=4Y-9SeXJGCq098Vy-PIbb_rexS2dYDlkU-kxnKAPVSU,1828
+mrmd_ai/signatures/edit.py,sha256=OBAYsh88Qg_EIoEJHHa28QRuSh2xbGM3OSGWtPJ8u_A,6524
 mrmd_ai/signatures/finish.py,sha256=x-ZB0U8GQJdNoGGO80FBOxHXjYsCmTFq9fnkXlHDeUY,5294
 mrmd_ai/signatures/fix.py,sha256=LJNvu9_XjPl90Wtt3xn6s-jGXA9GB5rdIL0MeFyRGtE,3042
 mrmd_ai/signatures/notebook.py,sha256=ZBioHA9ZTkLUD_UovdfiRYiDaUKuKOCDhiZP1NDFY8o,1226
 mrmd_ai/signatures/text.py,sha256=GhmFtEZqwivbevPI_NSBzh6AlH6JKLt2rA_LaYGK2lQ,5223
 mrmd_ai/utils/__init__.py,sha256=T4e9jmFWDSj1HOyz5_Qv-JQSC08GwT_9CACcAn37vWg,46
-mrmd_ai-0.1.0.dist-info/METADATA,sha256=4w17Do4YbBpwoT3PFN2i01zswTPJCVzR6rgib0zZXME,1167
-mrmd_ai-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mrmd_ai-0.1.0.dist-info/entry_points.txt,sha256=Bq6nXiXxhNSPEgYWBgrrgJH1DeGjKM6hfCijcWClApw,55
-mrmd_ai-0.1.0.dist-info/RECORD,,
+mrmd_ai-0.1.1.dist-info/METADATA,sha256=_Vd_Hf2JHGDyX9ZtTd3ogma9NnwfDyml1mgTzpEHRQI,1167
+mrmd_ai-0.1.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mrmd_ai-0.1.1.dist-info/entry_points.txt,sha256=Bq6nXiXxhNSPEgYWBgrrgJH1DeGjKM6hfCijcWClApw,55
+mrmd_ai-0.1.1.dist-info/RECORD,,