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

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,8 @@ 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 .custom_programs import custom_registry, register_custom_programs
 from .modules import (
     # Finish
     FinishSentencePredict,
@@ -56,6 +57,11 @@ from .modules import (
     DocumentAnalysisPredict,
     # Notebook
     NotebookNamePredict,
+    # Edit (Ctrl-K and comments)
+    EditAtCursorPredict,
+    AddressCommentPredict,
+    AddressAllCommentsPredict,
+    AddressNearbyCommentPredict,
 )
 
 
@@ -92,21 +98,101 @@ 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
+# NOTE: Cache is only used when no custom API keys are provided
+_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)
+# API key header names
+API_KEY_HEADERS = {
+    "anthropic": "X-Api-Key-Anthropic",
+    "openai": "X-Api-Key-Openai",
+    "groq": "X-Api-Key-Groq",
+    "gemini": "X-Api-Key-Gemini",
+    "openrouter": "X-Api-Key-Openrouter",
+}
+
+
+def extract_api_keys(request: Request) -> dict | None:
+    """Extract API keys from request headers.
+
+    Headers:
+        X-Api-Key-Anthropic: Anthropic API key
+        X-Api-Key-Openai: OpenAI API key
+        X-Api-Key-Groq: Groq API key
+        X-Api-Key-Gemini: Google Gemini API key
+        X-Api-Key-Openrouter: OpenRouter API key
+
+    Returns:
+        Dict of provider -> key if any keys are provided, None otherwise.
+    """
+    api_keys = {}
+    for provider, header in API_KEY_HEADERS.items():
+        key = request.headers.get(header)
+        if key:
+            api_keys[provider] = key
+
+    return api_keys if api_keys else None
+
+
+def get_program(
+    name: str,
+    juice: int = 0,
+    reasoning: int | None = None,
+    api_keys: dict | None = None,
+    model_override: str | None = None,
+) -> JuicedProgram:
+    """Get a JuicedProgram instance for the given program configuration.
+
+    Args:
+        name: Program name (can be built-in or custom)
+        juice: Juice level (0-4)
+        reasoning: Optional reasoning level (0-5)
+        api_keys: Optional dict of provider -> API key
+        model_override: Optional model to use instead of default
+
+    Returns:
+        Configured JuicedProgram instance.
+
+    Note:
+        Programs with custom API keys or model overrides are NOT cached,
+        since they need fresh instances with the provided configuration.
+        Custom programs are never cached.
+    """
+    # Check built-in programs first
+    program_class = PROGRAMS.get(name)
+
+    # If not found, check custom registry
+    if program_class is None:
+        program_class = custom_registry.get(name)
+
+    if program_class is None:
+        raise ValueError(f"Unknown program: {name}")
+
+    # Custom programs and those with custom config are never cached
+    is_custom = name not in PROGRAMS
+    if api_keys or model_override or is_custom:
+        program = program_class()
+        return JuicedProgram(
+            program,
+            juice=juice,
+            reasoning=reasoning,
+            api_keys=api_keys,
+            model_override=model_override,
+        )
+
+    # Use cache for standard built-in requests (no custom keys)
+    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 +239,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 +289,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 +316,22 @@ 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
+
+    # Extract API keys from headers (optional)
+    api_keys = extract_api_keys(request)
+
+    # Get model override from header (optional)
+    model_override = request.headers.get("X-Model-Override")
+
     # Get request body
     try:
         params = await request.json()
@@ -218,14 +340,23 @@ 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,
+            api_keys=api_keys,
+            model_override=model_override,
+        )
     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 +376,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 +425,22 @@ 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
+
+    # Extract API keys from headers (optional)
+    api_keys = extract_api_keys(request)
+
+    # Get model override from header (optional)
+    model_override = request.headers.get("X-Model-Override")
+
     # Get request body
     try:
         params = await request.json()
@@ -286,9 +452,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:
@@ -315,17 +484,26 @@ async def run_program_stream(program_name: str, request: Request):
         def run_with_progress():
             """Run the program in a thread, emitting progress events."""
             try:
-                # Create program with progress callback
+                # Create program with progress callback and optional API keys
                 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,
+                    api_keys=api_keys,
+                    model_override=model_override,
+                )
 
                 # 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 +535,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
 
@@ -390,6 +569,79 @@ async def run_program_stream(program_name: str, request: Request):
     )
 
 
+# =============================================================================
+# CUSTOM PROGRAMS API
+# =============================================================================
+
+class CustomProgramConfig(BaseModel):
+    """Configuration for a custom program."""
+    id: str  # Unique program ID (e.g., "Custom_cmd-123")
+    name: str
+    instructions: str
+    inputType: str = "selection"  # selection | cursor | fullDoc
+    outputType: str = "replace"  # replace | insert
+
+
+class RegisterCustomProgramsRequest(BaseModel):
+    """Request to register multiple custom programs."""
+    programs: list[CustomProgramConfig]
+
+
+@app.post("/api/custom-programs/register")
+async def register_custom_programs_endpoint(request: RegisterCustomProgramsRequest):
+    """Register custom programs from frontend configuration.
+
+    This endpoint is called when the app starts or when settings change.
+    It clears existing custom programs and registers the new ones.
+    """
+    # Clear existing custom programs
+    custom_registry.clear()
+
+    # Register new programs
+    registered = []
+    for prog in request.programs:
+        try:
+            custom_registry.register(
+                program_id=prog.id,
+                config={
+                    "name": prog.name,
+                    "instructions": prog.instructions,
+                    "inputType": prog.inputType,
+                    "outputType": prog.outputType,
+                }
+            )
+            registered.append(prog.id)
+            print(f"[Custom] Registered: {prog.name} ({prog.id})")
+        except Exception as e:
+            print(f"[Custom] Failed to register {prog.id}: {e}")
+
+    return {"registered": registered, "count": len(registered)}
+
+
+@app.get("/api/custom-programs")
+async def list_custom_programs():
+    """List all registered custom programs."""
+    programs = []
+    for program_id in custom_registry.list_programs():
+        config = custom_registry.get_config(program_id)
+        if config:
+            programs.append({
+                "id": program_id,
+                "name": config.get("name", program_id),
+                "inputType": config.get("inputType", "selection"),
+                "outputType": config.get("outputType", "replace"),
+            })
+    return {"programs": programs}
+
+
+@app.delete("/api/custom-programs/{program_id}")
+async def unregister_custom_program(program_id: str):
+    """Unregister a specific custom program."""
+    if custom_registry.unregister(program_id):
+        return {"success": True, "id": program_id}
+    raise HTTPException(status_code=404, detail=f"Program not found: {program_id}")
+
+
 def main():
     """Run the AI server."""
     import argparse

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",
 ]