emdash-core 0.1.25__py3-none-any.whl → 0.1.37__py3-none-any.whl

emdash_core/agent/runner/context.py

@@ -0,0 +1,470 @@
+"""Context management functions for the agent runner.
+
+This module contains functions for estimating, compacting, and managing
+conversation context during agent runs.
+"""
+
+import os
+from typing import Optional, TYPE_CHECKING
+
+from ...utils.logger import log
+
+if TYPE_CHECKING:
+    from ..toolkit import AgentToolkit
+    from ..events import AgentEventEmitter
+
+
+def estimate_context_tokens(messages: list[dict], system_prompt: Optional[str] = None) -> int:
+    """Estimate the current context window size in tokens.
+
+    Args:
+        messages: Conversation messages
+        system_prompt: Optional system prompt to include in estimation
+
+    Returns:
+        Estimated token count for the context
+    """
+    total_chars = 0
+
+    # Count characters in all messages
+    for msg in messages:
+        content = msg.get("content", "")
+        if isinstance(content, str):
+            total_chars += len(content)
+        elif isinstance(content, list):
+            # Handle multi-part messages (e.g., with images)
+            for part in content:
+                if isinstance(part, dict) and "text" in part:
+                    total_chars += len(part["text"])
+
+        # Add role overhead (~4 tokens per message for role/structure)
+        total_chars += 16
+
+    # Also count system prompt
+    if system_prompt:
+        total_chars += len(system_prompt)
+
+    # Estimate: ~4 characters per token
+    return total_chars // 4
+
+
+def get_context_breakdown(
+    messages: list[dict],
+    system_prompt: Optional[str] = None,
+) -> tuple[dict, list[dict]]:
+    """Get breakdown of context usage by message type.
+
+    Args:
+        messages: Conversation messages
+        system_prompt: Optional system prompt
+
+    Returns:
+        Tuple of (breakdown dict, list of largest messages)
+    """
+    breakdown = {
+        "system_prompt": len(system_prompt) // 4 if system_prompt else 0,
+        "user": 0,
+        "assistant": 0,
+        "tool_results": 0,
+    }
+
+    # Track individual message sizes for finding largest
+    message_sizes = []
+
+    for i, msg in enumerate(messages):
+        role = msg.get("role", "unknown")
+        content = msg.get("content", "")
+
+        # Calculate content size
+        if isinstance(content, str):
+            size = len(content)
+        elif isinstance(content, list):
+            size = sum(len(p.get("text", "")) for p in content if isinstance(p, dict))
+        else:
+            size = 0
+
+        tokens = size // 4
+
+        # Categorize
+        if role == "user":
+            breakdown["user"] += tokens
+        elif role == "assistant":
+            breakdown["assistant"] += tokens
+        elif role == "tool":
+            breakdown["tool_results"] += tokens
+
+        # Track for largest messages
+        if tokens > 100:  # Only track substantial messages
+            # Try to get a label for this message
+            label = f"{role}[{i}]"
+            if role == "tool":
+                tool_call_id = msg.get("tool_call_id", "")
+                # Try to find the tool name from previous assistant message
+                for prev_msg in reversed(messages[:i]):
+                    if prev_msg.get("role") == "assistant" and "tool_calls" in prev_msg:
+                        for tc in prev_msg.get("tool_calls", []):
+                            if tc.get("id") == tool_call_id:
+                                label = tc.get("function", {}).get("name", "tool")
+                                break
+                        break
+
+            message_sizes.append({
+                "index": i,
+                "role": role,
+                "label": label,
+                "tokens": tokens,
+                "preview": content[:100] if isinstance(content, str) else str(content)[:100],
+            })
+
+    # Sort by size and get top 5
+    message_sizes.sort(key=lambda x: x["tokens"], reverse=True)
+    largest = message_sizes[:5]
+
+    return breakdown, largest
+
+
+def maybe_compact_context(
+    messages: list[dict],
+    provider: object,
+    emitter: "AgentEventEmitter",
+    system_prompt: Optional[str] = None,
+    threshold: float = 0.8,
+) -> list[dict]:
+    """Proactively compact context if approaching limit.
+
+    Args:
+        messages: Current conversation messages
+        provider: LLM provider instance
+        emitter: Event emitter for notifications
+        system_prompt: System prompt for token estimation
+        threshold: Trigger compaction at this % of context limit (default 80%)
+
+    Returns:
+        Original or compacted messages
+    """
+    context_tokens = estimate_context_tokens(messages, system_prompt)
+    context_limit = provider.get_context_limit()
+
+    # Check if we need to compact
+    if context_tokens < context_limit * threshold:
+        return messages  # No compaction needed
+
+    log.info(
+        f"Context at {context_tokens:,}/{context_limit:,} tokens "
+        f"({context_tokens/context_limit:.0%}), compacting..."
+    )
+
+    return compact_messages_with_llm(
+        messages, emitter, target_tokens=int(context_limit * 0.5)
+    )
+
+
+def compact_messages_with_llm(
+    messages: list[dict],
+    emitter: "AgentEventEmitter",
+    target_tokens: int,
+) -> list[dict]:
+    """Use fast LLM to summarize middle messages.
+
+    Preserves:
+    - First message (original user request)
+    - Last 4 messages (recent context)
+    - Summarizes everything in between
+
+    Args:
+        messages: Current conversation messages
+        emitter: Event emitter for notifications
+        target_tokens: Target token count after compaction
+
+    Returns:
+        Compacted messages list
+    """
+    from ..subagent import get_model_for_tier
+    from ..providers import get_provider
+
+    if len(messages) <= 5:
+        return messages  # Too few to compact
+
+    # Split messages
+    first_msg = messages[0]
+    recent_msgs = messages[-4:]
+    middle_msgs = messages[1:-4]
+
+    if not middle_msgs:
+        return messages
+
+    # Build summary prompt
+    middle_content = format_messages_for_summary(middle_msgs)
+
+    prompt = f"""Summarize this conversation history concisely.
+
+PRESERVE (include verbatim if present):
+- Code snippets and file paths
+- Error messages
+- Key decisions made
+- Important tool results (file contents, search results)
+
+CONDENSE:
+- Repetitive searches
+- Verbose tool outputs
+- Intermediate reasoning
+
+CONVERSATION HISTORY:
+{middle_content}
+
+OUTPUT FORMAT:
+Provide a concise summary (max 2000 tokens) that captures the essential context needed to continue this task."""
+
+    # Use fast model for summarization
+    fast_model = get_model_for_tier("fast")
+    fast_provider = get_provider(fast_model)
+
+    try:
+        emitter.emit_thinking("Compacting context with fast model...")
+
+        response = fast_provider.chat(
+            messages=[{"role": "user", "content": prompt}],
+            system="You are a context summarizer. Be concise but preserve code and technical details.",
+        )
+
+        summary = response.content or ""
+
+        log.info(
+            f"Compacted {len(middle_msgs)} messages into summary "
+            f"({len(summary)} chars)"
+        )
+
+        # Build compacted messages
+        return [
+            first_msg,
+            {
+                "role": "assistant",
+                "content": f"[Context Summary]\n{summary}\n[End Summary]",
+            },
+            *recent_msgs,
+        ]
+    except Exception as e:
+        log.warning(f"LLM compaction failed: {e}, falling back to truncation")
+        return [first_msg] + recent_msgs
+
+
+def format_messages_for_summary(messages: list[dict]) -> str:
+    """Format messages for summarization prompt.
+
+    Args:
+        messages: Messages to format
+
+    Returns:
+        Formatted string for summarization
+    """
+    parts = []
+    for msg in messages:
+        role = msg.get("role", "unknown")
+        content = msg.get("content", "")
+
+        # Handle tool calls in assistant messages
+        if role == "assistant" and "tool_calls" in msg:
+            tool_calls = msg.get("tool_calls", [])
+            tool_info = [
+                f"Called: {tc.get('function', {}).get('name', 'unknown')}"
+                for tc in tool_calls
+            ]
+            content = f"{content}\n[Tools: {', '.join(tool_info)}]" if content else f"[Tools: {', '.join(tool_info)}]"
+
+        # Truncate very long content
+        if len(content) > 4000:
+            content = content[:4000] + "\n[...truncated...]"
+
+        parts.append(f"[{role.upper()}]\n{content}")
+
+    return "\n\n---\n\n".join(parts)
+
+
+def get_reranked_context(
+    toolkit: "AgentToolkit",
+    current_query: str,
+) -> dict:
+    """Get reranked context items based on the current query.
+
+    Args:
+        toolkit: Agent toolkit instance
+        current_query: Current query for relevance ranking
+
+    Returns:
+        Dict with item_count and items list
+    """
+    try:
+        from ...context.service import ContextService
+        from ...context.reranker import rerank_context_items
+
+        # Get exploration steps for context extraction
+        steps = toolkit.get_exploration_steps()
+        if not steps:
+            return {"item_count": 0, "items": [], "query": current_query, "debug": "no exploration steps"}
+
+        # Use context service to extract context items from exploration
+        service = ContextService(connection=toolkit.connection)
+        terminal_id = service.get_terminal_id()
+
+        # Update context with exploration steps
+        service.update_context(
+            terminal_id=terminal_id,
+            exploration_steps=steps,
+        )
+
+        # Get context items
+        items = service.get_context_items(terminal_id)
+        if not items:
+            return {"item_count": 0, "items": [], "query": current_query, "debug": f"no items from service ({len(steps)} steps)"}
+
+        # Get max tokens from env (default 2000)
+        max_tokens = int(os.getenv("CONTEXT_FRAME_MAX_TOKENS", "2000"))
+
+        # Rerank by query relevance
+        if current_query:
+            items = rerank_context_items(
+                items,
+                current_query,
+                top_k=50,  # Get more candidates, then filter by tokens
+            )
+
+        # Convert to serializable format, limiting by token count
+        result_items = []
+        total_tokens = 0
+        for item in items:
+            item_dict = {
+                "name": item.qualified_name,
+                "type": item.entity_type,
+                "file": item.file_path,
+                "score": round(item.score, 3) if hasattr(item, 'score') else None,
+                "description": item.description[:200] if item.description else None,
+                "touch_count": item.touch_count,
+                "neighbors": item.neighbors[:5] if item.neighbors else [],
+            }
+            # Estimate tokens for this item (~4 chars per token)
+            item_chars = len(str(item_dict))
+            item_tokens = item_chars // 4
+
+            if total_tokens + item_tokens > max_tokens:
+                break
+
+            result_items.append(item_dict)
+            total_tokens += item_tokens
+
+        return {
+            "item_count": len(result_items),
+            "items": result_items,
+            "query": current_query,
+            "total_tokens": total_tokens,
+        }
+
+    except Exception as e:
+        log.warning(f"Failed to get reranked context: {e}")
+        return {"item_count": 0, "items": [], "query": current_query, "debug": str(e)}
+
+
+def emit_context_frame(
+    toolkit: "AgentToolkit",
+    emitter: "AgentEventEmitter",
+    messages: list[dict],
+    system_prompt: Optional[str],
+    current_query: str,
+    total_input_tokens: int,
+    total_output_tokens: int,
+) -> None:
+    """Emit a context frame event with current exploration state.
+
+    Args:
+        toolkit: Agent toolkit instance
+        emitter: Event emitter
+        messages: Current conversation messages
+        system_prompt: System prompt for estimation
+        current_query: Current query for reranking
+        total_input_tokens: Total input tokens used
+        total_output_tokens: Total output tokens used
+    """
+    # Get exploration steps from toolkit session
+    steps = toolkit.get_exploration_steps()
+
+    # Estimate current context window tokens and get breakdown
+    context_tokens = 0
+    context_breakdown = {}
+    largest_messages = []
+    if messages:
+        context_tokens = estimate_context_tokens(messages, system_prompt)
+        context_breakdown, largest_messages = get_context_breakdown(messages, system_prompt)
+
+    # Summarize exploration by tool
+    tool_counts: dict[str, int] = {}
+    entities_found = 0
+    step_details: list[dict] = []
+
+    for step in steps:
+        tool_name = getattr(step, 'tool', 'unknown')
+        tool_counts[tool_name] = tool_counts.get(tool_name, 0) + 1
+
+        # Count entities from the step
+        step_entities = getattr(step, 'entities_found', [])
+        entities_found += len(step_entities)
+
+        # Collect step details
+        params = getattr(step, 'params', {})
+        summary = getattr(step, 'result_summary', '')
+
+        # Extract meaningful info based on tool type
+        detail = {
+            "tool": tool_name,
+            "summary": summary,
+        }
+
+        # Add relevant params based on tool
+        if tool_name == 'read_file' and 'file_path' in params:
+            detail["file"] = params['file_path']
+        elif tool_name == 'read_file' and 'path' in params:
+            detail["file"] = params['path']
+        elif tool_name in ('grep', 'semantic_search') and 'query' in params:
+            detail["query"] = params['query']
+        elif tool_name == 'glob' and 'pattern' in params:
+            detail["pattern"] = params['pattern']
+        elif tool_name == 'list_files' and 'path' in params:
+            detail["path"] = params['path']
+
+        # Add content preview if available
+        content_preview = getattr(step, 'content_preview', None)
+        if content_preview:
+            detail["content_preview"] = content_preview
+
+        # Add token count if available
+        token_count = getattr(step, 'token_count', 0)
+        if token_count > 0:
+            detail["tokens"] = token_count
+
+        # Add entities if any
+        if step_entities:
+            detail["entities"] = step_entities[:5]  # Limit to 5
+
+        step_details.append(detail)
+
+    exploration_steps = [
+        {"tool": tool, "count": count}
+        for tool, count in tool_counts.items()
+    ]
+
+    # Build context frame data
+    adding = {
+        "exploration_steps": exploration_steps,
+        "entities_found": entities_found,
+        "step_count": len(steps),
+        "details": step_details[-20:],  # Last 20 steps
+        "input_tokens": total_input_tokens,
+        "output_tokens": total_output_tokens,
+        "context_tokens": context_tokens,  # Current context window size
+        "context_breakdown": context_breakdown,  # Tokens by message type
+        "largest_messages": largest_messages,  # Top 5 biggest messages
+    }
+
+    # Get reranked context items
+    reading = get_reranked_context(toolkit, current_query)
+
+    # Emit the context frame
+    emitter.emit_context_frame(adding=adding, reading=reading)

emdash_core/agent/runner/factory.py

@@ -0,0 +1,108 @@
+"""Runner factory for hybrid SDK/Provider routing.
+
+This module provides factory functions to create the appropriate runner
+based on the model type:
+- Claude models → SDKAgentRunner (uses Anthropic Agent SDK)
+- Other models → AgentRunner (uses OpenAI-compatible API)
+"""
+
+from pathlib import Path
+from typing import Optional, Union, TYPE_CHECKING
+
+from .sdk_runner import SDKAgentRunner, is_claude_model
+from .agent_runner import AgentRunner
+from ...utils.logger import log
+
+if TYPE_CHECKING:
+    from ..events import AgentEventEmitter
+
+
+def get_runner(
+    model: str,
+    cwd: Optional[str] = None,
+    emitter: Optional["AgentEventEmitter"] = None,
+    system_prompt: Optional[str] = None,
+    plan_mode: bool = False,
+    prefer_sdk: bool = True,
+    **kwargs,
+) -> Union[SDKAgentRunner, AgentRunner]:
+    """Get the appropriate runner for the model.
+
+    Routes to SDKAgentRunner for Claude models (uses Anthropic Agent SDK),
+    or AgentRunner for other models (uses OpenAI-compatible API).
+
+    Args:
+        model: Model string (e.g., "claude-sonnet-4", "fireworks:minimax-m2p1")
+        cwd: Working directory
+        emitter: Event emitter for streaming
+        system_prompt: Custom system prompt
+        plan_mode: If True, restrict to read-only tools
+        prefer_sdk: If True, prefer SDK for Claude models (default True)
+        **kwargs: Additional arguments passed to runner
+
+    Returns:
+        SDKAgentRunner or AgentRunner instance
+
+    Example:
+        # Claude model → uses SDK
+        runner = get_runner("claude-sonnet-4")
+
+        # Fireworks model → uses standard provider
+        runner = get_runner("fireworks:accounts/fireworks/models/minimax-m2p1")
+
+        # Force standard provider even for Claude
+        runner = get_runner("claude-sonnet-4", prefer_sdk=False)
+    """
+    use_sdk = prefer_sdk and is_claude_model(model)
+
+    if use_sdk:
+        log.info(f"Using SDKAgentRunner for Claude model: {model}")
+        return SDKAgentRunner(
+            model=model,
+            cwd=cwd,
+            emitter=emitter,
+            system_prompt=system_prompt,
+            plan_mode=plan_mode,
+        )
+    else:
+        log.info(f"Using AgentRunner for model: {model}")
+        # Import toolkit here to avoid circular imports
+        from ..toolkit import AgentToolkit
+
+        # Generate plan file path when in plan mode
+        plan_file_path = None
+        repo_root = Path(cwd) if cwd else Path.cwd()
+        if plan_mode:
+            plan_file_path = str(repo_root / ".emdash" / "plan.md")
+            # Ensure .emdash directory exists
+            (repo_root / ".emdash").mkdir(exist_ok=True)
+
+        toolkit = AgentToolkit(
+            repo_root=repo_root,
+            plan_mode=plan_mode,
+            plan_file_path=plan_file_path,
+        )
+
+        return AgentRunner(
+            toolkit=toolkit,
+            model=model,
+            emitter=emitter,
+            system_prompt=system_prompt,
+            **kwargs,
+        )
+
+
+def create_hybrid_runner(
+    model: str,
+    **kwargs,
+) -> Union[SDKAgentRunner, AgentRunner]:
+    """Convenience alias for get_runner.
+
+    Args:
+        model: Model string
+        **kwargs: Passed to get_runner
+
+    Returns:
+        Appropriate runner instance
+    """
+    return get_runner(model, **kwargs)