emdash-core 0.1.7__py3-none-any.whl

emdash_core/agent/runner.py

@@ -0,0 +1,601 @@
+"""Agent runner for LLM-powered exploration."""
+
+import json
+from datetime import datetime, date
+from typing import Any, Optional
+
+from ..utils.logger import log
+from ..core.config import get_config
+from ..core.exceptions import ContextLengthError
+from .toolkit import AgentToolkit
+from .events import AgentEventEmitter, NullEmitter
+from .providers import get_provider
+from .providers.factory import DEFAULT_MODEL
+from .context_manager import (
+    truncate_tool_output,
+    reduce_context_for_retry,
+    is_context_overflow_error,
+)
+from .prompts import BASE_SYSTEM_PROMPT, build_system_prompt
+from .tools.tasks import TaskState
+
+
+class SafeJSONEncoder(json.JSONEncoder):
+    """JSON encoder that handles Neo4j types and other non-serializable objects."""
+
+    def default(self, obj: Any) -> Any:
+        # Handle datetime objects
+        if isinstance(obj, (datetime, date)):
+            return obj.isoformat()
+
+        # Handle Neo4j DateTime
+        if hasattr(obj, 'isoformat'):
+            return obj.isoformat()
+
+        # Handle Neo4j Date, Time, etc.
+        if hasattr(obj, 'to_native'):
+            return str(obj.to_native())
+
+        # Handle sets
+        if isinstance(obj, set):
+            return list(obj)
+
+        # Handle bytes
+        if isinstance(obj, bytes):
+            return obj.decode('utf-8', errors='replace')
+
+        # Fallback to string representation
+        try:
+            return str(obj)
+        except Exception:
+            return f"<non-serializable: {type(obj).__name__}>"
+
+
+class AgentRunner:
+    """Runs an LLM agent with tool access for code exploration.
+
+    Example:
+        runner = AgentRunner()
+        response = runner.run("How does authentication work in this codebase?")
+        print(response)
+    """
+
+    def __init__(
+        self,
+        toolkit: Optional[AgentToolkit] = None,
+        model: str = DEFAULT_MODEL,
+        system_prompt: Optional[str] = None,
+        emitter: Optional[AgentEventEmitter] = None,
+        max_iterations: int = 50,
+        verbose: bool = False,
+        show_tool_results: bool = False,
+    ):
+        """Initialize the agent runner.
+
+        Args:
+            toolkit: AgentToolkit instance. If None, creates default.
+            model: LLM model to use.
+            system_prompt: Custom system prompt. If None, uses default.
+            emitter: Event emitter for streaming output.
+            max_iterations: Maximum tool call iterations.
+            verbose: Whether to print verbose output.
+            show_tool_results: Whether to show detailed tool results.
+        """
+        self.toolkit = toolkit or AgentToolkit()
+        self.provider = get_provider(model)
+        self.model = model
+        self.system_prompt = system_prompt or build_system_prompt(self.toolkit)
+        self.emitter = emitter or NullEmitter()
+        # Inject emitter into tools that need it (e.g., TaskTool for sub-agent streaming)
+        self.toolkit.set_emitter(self.emitter)
+        self.max_iterations = max_iterations
+        self.verbose = verbose
+        self.show_tool_results = show_tool_results
+        # Conversation history for multi-turn support
+        self._messages: list[dict] = []
+        # Token usage tracking
+        self._total_input_tokens: int = 0
+        self._total_output_tokens: int = 0
+        # Store query for reranking
+        self._current_query: str = ""
+        # Todo state tracking for injection
+        self._last_todo_snapshot: str = ""
+
+    def _get_todo_snapshot(self) -> str:
+        """Get current todo state as string for comparison."""
+        state = TaskState.get_instance()
+        return json.dumps(state.get_all_tasks(), sort_keys=True)
+
+    def _format_todo_reminder(self) -> str:
+        """Format current todos as XML reminder for injection into context."""
+        state = TaskState.get_instance()
+        tasks = state.get_all_tasks()
+        if not tasks:
+            return ""
+
+        counts = {"pending": 0, "in_progress": 0, "completed": 0}
+        lines = []
+        for t in tasks:
+            status = t.get("status", "pending")
+            counts[status] = counts.get(status, 0) + 1
+            status_icon = {"pending": "⬚", "in_progress": "🔄", "completed": "✅"}.get(status, "?")
+            lines.append(f'  {t["id"]}. {status_icon} {t["title"]}')
+
+        header = f'Tasks: {counts["completed"]} completed, {counts["in_progress"]} in progress, {counts["pending"]} pending'
+        task_list = "\n".join(lines)
+        return f"<todo-state>\n{header}\n{task_list}\n</todo-state>"
+
+    def run(
+        self,
+        query: str,
+        context: Optional[str] = None,
+        images: Optional[list] = None,
+    ) -> str:
+        """Run the agent to answer a query.
+
+        Args:
+            query: User's question or request
+            context: Optional additional context
+            images: Optional list of images to include
+
+        Returns:
+            Agent's final response
+        """
+        # Store query for reranking context frame
+        self._current_query = query
+        self.emitter.emit_start(goal=query)
+
+        # Build messages
+        messages = []
+
+        if context:
+            messages.append({
+                "role": "user",
+                "content": f"Context:\n{context}\n\nQuestion: {query}",
+            })
+        else:
+            messages.append({
+                "role": "user",
+                "content": query,
+            })
+
+        # TODO: Handle images if provided
+
+        # Get tool schemas
+        tools = self.toolkit.get_all_schemas()
+
+        try:
+            response, final_messages = self._run_loop(messages, tools)
+            # Save conversation history for multi-turn support
+            self._messages = final_messages
+            self.emitter.emit_end(success=True)
+            return response
+
+        except Exception as e:
+            log.exception("Agent run failed")
+            self.emitter.emit_error(str(e))
+            return f"Error: {str(e)}"
+
+    def _run_loop(
+        self,
+        messages: list[dict],
+        tools: list[dict],
+    ) -> tuple[str, list[dict]]:
+        """Run the agent loop until completion.
+
+        Args:
+            messages: Initial messages
+            tools: Tool schemas
+
+        Returns:
+            Tuple of (final response text, conversation messages)
+        """
+        max_retries = 3
+
+        for iteration in range(self.max_iterations):
+            # Try API call with retry on context overflow
+            retry_count = 0
+            response = None
+
+            while retry_count < max_retries:
+                try:
+                    response = self.provider.chat(
+                        messages=messages,
+                        system=self.system_prompt,
+                        tools=tools,
+                    )
+                    break  # Success
+
+                except Exception as exc:
+                    if is_context_overflow_error(exc):
+                        retry_count += 1
+                        log.warning(
+                            "Context overflow on attempt {}/{}, reducing context...",
+                            retry_count,
+                            max_retries,
+                        )
+
+                        if retry_count >= max_retries:
+                            raise ContextLengthError(
+                                f"Failed to reduce context after {max_retries} attempts: {exc}",
+                            )
+
+                        # Reduce context by removing old messages
+                        messages = reduce_context_for_retry(
+                            messages,
+                            keep_recent=max(2, 6 - retry_count * 2),  # Fewer messages each retry
+                        )
+                    else:
+                        raise  # Re-raise non-context errors
+
+            if response is None:
+                raise RuntimeError("Failed to get response from provider")
+
+            # Accumulate token usage
+            self._total_input_tokens += response.input_tokens
+            self._total_output_tokens += response.output_tokens
+
+            # Check for tool calls
+            if response.tool_calls:
+                # Emit any content that accompanies tool calls (e.g., thinking text)
+                if response.content:
+                    self.emitter.emit_message_start()
+                    self.emitter.emit_message_delta(response.content)
+                    self.emitter.emit_message_end()
+
+                # Execute tools and add results
+                for tool_call in response.tool_calls:
+                    # Parse arguments if they're a JSON string
+                    args = tool_call.arguments
+                    if isinstance(args, str):
+                        args = json.loads(args)
+
+                    self.emitter.emit_tool_start(
+                        tool_call.name,
+                        args,
+                    )
+
+                    result = self.toolkit.execute(
+                        tool_call.name,
+                        **args,
+                    )
+
+                    self.emitter.emit_tool_result(
+                        tool_call.name,
+                        result.success,
+                        self._summarize_result(result),
+                    )
+
+                    # Check if tool is asking a clarification question
+                    if (result.success and
+                        result.data and
+                        result.data.get("status") == "awaiting_response" and
+                        "question" in result.data):
+                        self.emitter.emit_clarification(
+                            question=result.data["question"],
+                            context="",
+                            options=result.data.get("options", []),
+                        )
+
+                    # Add assistant message with tool call
+                    messages.append({
+                        "role": "assistant",
+                        "content": response.content or "",
+                        "tool_calls": [{
+                            "id": tool_call.id,
+                            "type": "function",
+                            "function": {
+                                "name": tool_call.name,
+                                "arguments": json.dumps(args),
+                            },
+                        }],
+                    })
+
+                    # Serialize and truncate tool result to prevent context overflow
+                    result_json = json.dumps(result.to_dict(), cls=SafeJSONEncoder)
+                    result_json = truncate_tool_output(result_json)
+
+                    # Check if todos changed and inject reminder
+                    if tool_call.name in ("write_todo", "update_todo_list"):
+                        new_snapshot = self._get_todo_snapshot()
+                        if new_snapshot != self._last_todo_snapshot:
+                            self._last_todo_snapshot = new_snapshot
+                            reminder = self._format_todo_reminder()
+                            if reminder:
+                                result_json += f"\n\n{reminder}"
+
+                    # Add tool result
+                    messages.append({
+                        "role": "tool",
+                        "tool_call_id": tool_call.id,
+                        "content": result_json,
+                    })
+
+
+            else:
+                # No tool calls - check if response was truncated
+                if response.stop_reason in ("max_tokens", "length"):
+                    # Response was truncated, request continuation
+                    log.debug("Response truncated ({}), requesting continuation", response.stop_reason)
+                    if response.content:
+                        messages.append({
+                            "role": "assistant",
+                            "content": response.content,
+                        })
+                    messages.append({
+                        "role": "user",
+                        "content": "Your response was cut off. Please continue.",
+                    })
+                    continue
+
+                # Agent is done - emit final response
+                if response.content:
+                    self.emitter.emit_message_start()
+                    self.emitter.emit_message_delta(response.content)
+                    self.emitter.emit_message_end()
+                    # Add final assistant message to history
+                    messages.append({
+                        "role": "assistant",
+                        "content": response.content,
+                    })
+
+                # Emit final context frame summary
+                self._emit_context_frame(messages)
+
+                return response.content or "", messages
+
+        # Hit max iterations - emit final message and context frame
+        final_message = "Reached maximum iterations without completing."
+        self.emitter.emit_message_start()
+        self.emitter.emit_message_delta(final_message)
+        self.emitter.emit_message_end()
+        self._emit_context_frame(messages)
+        return final_message, messages
+
+    def _summarize_result(self, result: Any) -> str:
+        """Create a brief summary of a tool result."""
+        if not result.success:
+            return f"Error: {result.error}"
+
+        if not result.data:
+            return "Empty result"
+
+        data = result.data
+
+        if "results" in data:
+            return f"{len(data['results'])} results"
+        elif "root_node" in data:
+            node = data["root_node"]
+            name = node.get("qualified_name") or node.get("file_path", "unknown")
+            return f"Expanded: {name}"
+        elif "callers" in data:
+            return f"{len(data['callers'])} callers"
+        elif "callees" in data:
+            return f"{len(data['callees'])} callees"
+
+        return "Completed"
+
+    def _emit_context_frame(self, messages: list[dict] = None) -> None:
+        """Emit a context frame event with current exploration state.
+
+        Args:
+            messages: Current conversation messages to estimate context size
+        """
+        # Get exploration steps from toolkit session
+        steps = self.toolkit.get_exploration_steps()
+
+        # Estimate current context window tokens
+        context_tokens = 0
+        if messages:
+            context_tokens = self._estimate_context_tokens(messages)
+
+        # 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": self._total_input_tokens,
+            "output_tokens": self._total_output_tokens,
+            "context_tokens": context_tokens,  # Current context window size
+        }
+
+        # Get reranked context items
+        reading = self._get_reranked_context()
+
+        # Emit the context frame
+        self.emitter.emit_context_frame(adding=adding, reading=reading)
+
+    def _estimate_context_tokens(self, messages: list[dict]) -> int:
+        """Estimate the current context window size in tokens.
+
+        Args:
+            messages: Conversation messages
+
+        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 self.system_prompt:
+            total_chars += len(self.system_prompt)
+
+        # Estimate: ~4 characters per token
+        return total_chars // 4
+
+    def _get_reranked_context(self) -> dict:
+        """Get reranked context items based on the current query.
+
+        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 = self.toolkit.get_exploration_steps()
+            if not steps:
+                return {"item_count": 0, "items": []}
+
+            # Use context service to extract context items from exploration
+            service = ContextService(connection=self.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": []}
+
+            # Rerank by query relevance
+            if self._current_query:
+                items = rerank_context_items(
+                    items,
+                    self._current_query,
+                    top_k=20,
+                )
+
+            # Convert to serializable format
+            result_items = []
+            for item in items[:20]:  # Limit to 20 items
+                result_items.append({
+                    "name": item.qualified_name,
+                    "type": item.entity_type,
+                    "file": item.file_path,
+                    "score": round(item.score, 3) if hasattr(item, 'score') else None,
+                })
+
+            return {
+                "item_count": len(result_items),
+                "items": result_items,
+            }
+
+        except Exception as e:
+            log.debug(f"Failed to get reranked context: {e}")
+            return {"item_count": 0, "items": []}
+
+    def chat(self, message: str, images: Optional[list] = None) -> str:
+        """Continue a conversation with a new message.
+
+        This method maintains conversation history for multi-turn interactions.
+        Call run() first to start a conversation, then chat() for follow-ups.
+
+        Args:
+            message: User's follow-up message
+            images: Optional list of images to include
+
+        Returns:
+            Agent's response
+        """
+        if not self._messages:
+            # No history, just run fresh
+            return self.run(message, images=images)
+
+        # Store query for reranking context frame
+        self._current_query = message
+        self.emitter.emit_start(goal=message)
+
+        # Add new user message to history
+        self._messages.append({
+            "role": "user",
+            "content": message,
+        })
+
+        # Get tool schemas
+        tools = self.toolkit.get_all_schemas()
+
+        try:
+            response, final_messages = self._run_loop(self._messages, tools)
+            # Update conversation history
+            self._messages = final_messages
+            self.emitter.emit_end(success=True)
+            return response
+
+        except Exception as e:
+            log.exception("Agent chat failed")
+            self.emitter.emit_error(str(e))
+            return f"Error: {str(e)}"
+
+    def reset(self) -> None:
+        """Reset the agent state."""
+        self.toolkit.reset_session()
+        self._total_input_tokens = 0
+        self._total_output_tokens = 0
+        self._current_query = ""