ripperdoc 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

ripperdoc/core/query.py

@@ -6,48 +6,60 @@ the query-response loop including tool execution.
 
 import asyncio
 import inspect
+import os
 import time
-from typing import Any, AsyncGenerator, Dict, Iterable, List, Optional, Tuple, Union, cast
+from asyncio import CancelledError
+from typing import (
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
 
-from anthropic import AsyncAnthropic
-from openai import AsyncOpenAI
 from pydantic import ValidationError
 
-from ripperdoc.core.config import ProviderType, provider_protocol
+from ripperdoc.core.config import provider_protocol
+from ripperdoc.core.providers import ProviderClient, get_provider_client
 from ripperdoc.core.permissions import PermissionResult
 from ripperdoc.core.query_utils import (
-    anthropic_usage_tokens,
-    build_anthropic_tool_schemas,
     build_full_system_prompt,
-    build_openai_tool_schemas,
-    content_blocks_from_anthropic_response,
-    content_blocks_from_openai_choice,
     determine_tool_mode,
     extract_tool_use_blocks,
     format_pydantic_errors,
     log_openai_messages,
-    openai_usage_tokens,
     resolve_model_profile,
     text_mode_history,
     tool_result_message,
 )
 from ripperdoc.core.tool import Tool, ToolProgress, ToolResult, ToolUseContext
+from ripperdoc.utils.file_watch import ChangedFileNotice, FileSnapshot, detect_changed_files
 from ripperdoc.utils.log import get_logger
 from ripperdoc.utils.messages import (
     AssistantMessage,
+    MessageContent,
     ProgressMessage,
     UserMessage,
     create_assistant_message,
+    create_user_message,
     create_progress_message,
     normalize_messages_for_api,
     INTERRUPT_MESSAGE,
     INTERRUPT_MESSAGE_FOR_TOOL_USE,
 )
-from ripperdoc.utils.session_usage import record_usage
 
 
 logger = get_logger()
 
+DEFAULT_REQUEST_TIMEOUT_SEC = float(os.getenv("RIPPERDOC_API_TIMEOUT", "120"))
+MAX_LLM_RETRIES = int(os.getenv("RIPPERDOC_MAX_RETRIES", "10"))
+
 
 def _resolve_tool(
     tool_registry: "ToolRegistry", tool_name: str, tool_use_id: str
@@ -62,11 +74,23 @@ def _resolve_tool(
     )
 
 
+ToolPermissionCallable = Callable[
+    [Tool[Any, Any], Any],
+    Union[
+        PermissionResult,
+        Dict[str, Any],
+        Tuple[bool, Optional[str]],
+        bool,
+        Awaitable[Union[PermissionResult, Dict[str, Any], Tuple[bool, Optional[str]], bool]],
+    ],
+]
+
+
 async def _check_tool_permissions(
     tool: Tool[Any, Any],
     parsed_input: Any,
     query_context: "QueryContext",
-    can_use_tool_fn: Optional[Any],
+    can_use_tool_fn: Optional[ToolPermissionCallable],
 ) -> tuple[bool, Optional[str]]:
     """Evaluate whether a tool call is allowed."""
     try:
@@ -102,6 +126,155 @@ async def _check_tool_permissions(
         return False, None
 
 
+def _format_changed_file_notice(notices: List[ChangedFileNotice]) -> str:
+    """Render a system notice about files that changed on disk."""
+    lines: List[str] = [
+        "System notice: Files you previously read have changed on disk.",
+        "Please re-read the affected files before making further edits.",
+        "",
+    ]
+    for notice in notices:
+        lines.append(f"- {notice.file_path}")
+        summary = (notice.summary or "").rstrip()
+        if summary:
+            indented = "\n".join(f"    {line}" for line in summary.splitlines())
+            lines.append(indented)
+    return "\n".join(lines)
+
+
+async def _run_tool_use_generator(
+    tool: Tool[Any, Any],
+    tool_use_id: str,
+    tool_name: str,
+    parsed_input: Any,
+    sibling_ids: set[str],
+    tool_context: ToolUseContext,
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Execute a single tool_use and yield progress/results."""
+    try:
+        async for output in tool.call(parsed_input, tool_context):
+            if isinstance(output, ToolProgress):
+                yield create_progress_message(
+                    tool_use_id=tool_use_id,
+                    sibling_tool_use_ids=sibling_ids,
+                    content=output.content,
+                )
+                logger.debug(f"[query] Progress from tool_use_id={tool_use_id}: {output.content}")
+            elif isinstance(output, ToolResult):
+                result_content = output.result_for_assistant or str(output.data)
+                result_msg = tool_result_message(
+                    tool_use_id, result_content, tool_use_result=output.data
+                )
+                yield result_msg
+                logger.debug(
+                    f"[query] Tool completed tool_use_id={tool_use_id} name={tool_name} "
+                    f"result_len={len(result_content)}"
+                )
+    except Exception as exc:
+        logger.exception(
+            f"Error executing tool '{tool_name}'",
+            extra={"tool": tool_name, "tool_use_id": tool_use_id},
+        )
+        yield tool_result_message(tool_use_id, f"Error executing tool: {str(exc)}", is_error=True)
+
+
+def _group_tool_calls_by_concurrency(prepared_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Group consecutive tool calls by their concurrency safety."""
+    groups: List[Dict[str, Any]] = []
+    for call in prepared_calls:
+        is_safe = bool(call.get("is_concurrency_safe"))
+        if groups and groups[-1]["is_concurrency_safe"] == is_safe:
+            groups[-1]["items"].append(call)
+        else:
+            groups.append({"is_concurrency_safe": is_safe, "items": [call]})
+    return groups
+
+
+async def _execute_tools_sequentially(
+    items: List[Dict[str, Any]], tool_results: List[UserMessage]
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Run tool generators one by one."""
+    for item in items:
+        gen = item.get("generator")
+        if not gen:
+            continue
+        async for message in gen:
+            if isinstance(message, UserMessage):
+                tool_results.append(message)
+            yield message
+
+
+async def _execute_tools_in_parallel(
+    items: List[Dict[str, Any]], tool_results: List[UserMessage]
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Run tool generators concurrently."""
+    generators = [call["generator"] for call in items if call.get("generator")]
+    async for message in _run_concurrent_tool_uses(generators, tool_results):
+        yield message
+
+
+async def _run_tools_concurrently(
+    prepared_calls: List[Dict[str, Any]], tool_results: List[UserMessage]
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Run tools grouped by concurrency safety (parallel for safe groups)."""
+    for group in _group_tool_calls_by_concurrency(prepared_calls):
+        if group["is_concurrency_safe"]:
+            logger.debug(
+                f"[query] Executing {len(group['items'])} concurrency-safe tool(s) in parallel"
+            )
+            async for message in _execute_tools_in_parallel(group["items"], tool_results):
+                yield message
+        else:
+            logger.debug(
+                f"[query] Executing {len(group['items'])} tool(s) sequentially (not concurrency safe)"
+            )
+            async for message in _run_tools_serially(group["items"], tool_results):
+                yield message
+
+
+async def _run_tools_serially(
+    prepared_calls: List[Dict[str, Any]], tool_results: List[UserMessage]
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Run all tools sequentially (helper for clarity)."""
+    async for message in _execute_tools_sequentially(prepared_calls, tool_results):
+        yield message
+
+
+async def _run_concurrent_tool_uses(
+    generators: List[AsyncGenerator[Union[UserMessage, ProgressMessage], None]],
+    tool_results: List[UserMessage],
+) -> AsyncGenerator[Union[UserMessage, ProgressMessage], None]:
+    """Drain multiple tool generators concurrently and stream outputs."""
+    if not generators:
+        return
+
+    queue: asyncio.Queue[Optional[Union[UserMessage, ProgressMessage]]] = asyncio.Queue()
+
+    async def _consume(gen: AsyncGenerator[Union[UserMessage, ProgressMessage], None]) -> None:
+        try:
+            async for message in gen:
+                await queue.put(message)
+        except Exception:
+            logger.exception("[query] Unexpected error while consuming tool generator")
+        finally:
+            await queue.put(None)
+
+    tasks = [asyncio.create_task(_consume(gen)) for gen in generators]
+    active = len(tasks)
+
+    try:
+        while active:
+            message = await queue.get()
+            if message is None:
+                active -= 1
+                continue
+            if isinstance(message, UserMessage):
+                tool_results.append(message)
+            yield message
+    finally:
+        await asyncio.gather(*tasks, return_exceptions=True)
+
+
 class ToolRegistry:
     """Track available tools, including deferred ones, and expose search/activation helpers."""
 
@@ -204,6 +377,8 @@ class QueryContext:
         safe_mode: bool = False,
         model: str = "main",
         verbose: bool = False,
+        pause_ui: Optional[Callable[[], None]] = None,
+        resume_ui: Optional[Callable[[], None]] = None,
     ) -> None:
         self.tool_registry = ToolRegistry(tools)
         self.max_thinking_tokens = max_thinking_tokens
@@ -211,6 +386,9 @@ class QueryContext:
         self.model = model
         self.verbose = verbose
         self.abort_controller = asyncio.Event()
+        self.file_state_cache: Dict[str, FileSnapshot] = {}
+        self.pause_ui = pause_ui
+        self.resume_ui = resume_ui
 
     @property
     def tools(self) -> List[Tool[Any, Any]]:
@@ -238,6 +416,11 @@ async def query_llm(
     max_thinking_tokens: int = 0,
     model: str = "main",
     abort_signal: Optional[asyncio.Event] = None,
+    *,
+    progress_callback: Optional[Callable[[str], Awaitable[None]]] = None,
+    request_timeout: Optional[float] = None,
+    max_retries: int = MAX_LLM_RETRIES,
+    stream: bool = True,
 ) -> AssistantMessage:
     """Query the AI model and return the response.
 
@@ -248,10 +431,16 @@ async def query_llm(
         max_thinking_tokens: Maximum tokens for thinking (0 = disabled)
         model: Model pointer to use
         abort_signal: Event to signal abortion
+        progress_callback: Optional async callback invoked with streamed text chunks
+        request_timeout: Max seconds to wait for a provider response before retrying
+        max_retries: Number of retries on timeout/errors (total attempts = retries + 1)
+        stream: Enable streaming for providers that support it (text-only mode)
 
     Returns:
         AssistantMessage with the model's response
     """
+    request_timeout = request_timeout or DEFAULT_REQUEST_TIMEOUT_SEC
+    request_timeout = request_timeout or DEFAULT_REQUEST_TIMEOUT_SEC
     model_profile = resolve_model_profile(model)
 
     # Normalize messages based on protocol family (Anthropic allows tool blocks; OpenAI-style prefers text-only)
@@ -266,7 +455,7 @@ async def query_llm(
     else:
         messages_for_model = messages
 
-    normalized_messages = normalize_messages_for_api(
+    normalized_messages: List[Dict[str, Any]] = normalize_messages_for_api(
         messages_for_model, protocol=protocol, tool_mode=tool_mode
     )
     logger.info(
@@ -295,95 +484,36 @@ async def query_llm(
     start_time = time.time()
 
     try:
-        # Create the appropriate client based on provider
-        if model_profile.provider == ProviderType.ANTHROPIC:
-            async with AsyncAnthropic(api_key=model_profile.api_key) as client:
-                tool_schemas = await build_anthropic_tool_schemas(tools)
-                response = await client.messages.create(
-                    model=model_profile.model,
-                    max_tokens=model_profile.max_tokens,
-                    system=system_prompt,
-                    messages=normalized_messages,  # type: ignore[arg-type]
-                    tools=tool_schemas if tool_schemas else None,  # type: ignore
-                    temperature=model_profile.temperature,
-                )
-
-                duration_ms = (time.time() - start_time) * 1000
-
-                usage_tokens = anthropic_usage_tokens(getattr(response, "usage", None))
-                record_usage(model_profile.model, duration_ms=duration_ms, **usage_tokens)
-
-                # Calculate cost (simplified, should use actual pricing)
-                cost_usd = 0.0  # TODO: Implement cost calculation
-
-                content_blocks = content_blocks_from_anthropic_response(response, tool_mode)
-                tool_use_blocks = [
-                    block for block in response.content if getattr(block, "type", None) == "tool_use"
-                ]
-                logger.info(
-                    "[query_llm] Received response from Anthropic",
-                    extra={
-                        "model": model_profile.model,
-                        "duration_ms": round(duration_ms, 2),
-                        "usage_tokens": usage_tokens,
-                        "tool_use_blocks": len(tool_use_blocks),
-                    },
-                )
-
-                return create_assistant_message(
-                    content=content_blocks,
-                    cost_usd=cost_usd,
-                    duration_ms=duration_ms,
-                )
-
-        elif model_profile.provider == ProviderType.OPENAI_COMPATIBLE:
-            # OpenAI-compatible APIs (OpenAI, DeepSeek, Mistral, etc.)
-            async with AsyncOpenAI(api_key=model_profile.api_key, base_url=model_profile.api_base) as client:
-                openai_tools = await build_openai_tool_schemas(tools)
-
-                # Prepare messages for OpenAI format
-                openai_messages = [
-                    {"role": "system", "content": system_prompt}
-                ] + normalized_messages
-
-                # Make the API call
-                openai_response: Any = await client.chat.completions.create(
-                    model=model_profile.model,
-                    messages=openai_messages,
-                    tools=openai_tools if openai_tools else None,  # type: ignore[arg-type]
-                    temperature=model_profile.temperature,
-                    max_tokens=model_profile.max_tokens,
-                )
-
-                duration_ms = (time.time() - start_time) * 1000
-                usage_tokens = openai_usage_tokens(getattr(openai_response, "usage", None))
-                record_usage(model_profile.model, duration_ms=duration_ms, **usage_tokens)
-                cost_usd = 0.0  # TODO: Implement cost calculation
-
-                # Convert OpenAI response to our format
-                content_blocks = []
-                choice = openai_response.choices[0]
-
-                logger.info(
-                    "[query_llm] Received response from OpenAI-compatible provider",
-                    extra={
-                        "model": model_profile.model,
-                        "duration_ms": round(duration_ms, 2),
-                        "usage_tokens": usage_tokens,
-                        "finish_reason": getattr(choice, "finish_reason", None),
-                    },
-                )
-
-                content_blocks = content_blocks_from_openai_choice(choice, tool_mode)
-
-                return create_assistant_message(
-                    content=content_blocks, cost_usd=cost_usd, duration_ms=duration_ms
-                )
+        client: Optional[ProviderClient] = get_provider_client(model_profile.provider)
+        if client is None:
+            duration_ms = (time.time() - start_time) * 1000
+            error_msg = create_assistant_message(
+                content=(
+                    "Gemini protocol is not supported yet in Ripperdoc. "
+                    "Please configure an Anthropic or OpenAI-compatible model."
+                ),
+                duration_ms=duration_ms,
+            )
+            error_msg.is_api_error_message = True
+            return error_msg
+
+        provider_response = await client.call(
+            model_profile=model_profile,
+            system_prompt=system_prompt,
+            normalized_messages=normalized_messages,
+            tools=tools,
+            tool_mode=tool_mode,
+            stream=stream,
+            progress_callback=progress_callback,
+            request_timeout=request_timeout,
+            max_retries=max_retries,
+        )
 
-        elif model_profile.provider == ProviderType.GEMINI:
-            raise NotImplementedError("Gemini protocol is not yet supported.")
-        else:
-            raise NotImplementedError(f"Provider {model_profile.provider} not yet implemented")
+        return create_assistant_message(
+            content=provider_response.content_blocks,
+            cost_usd=provider_response.cost_usd,
+            duration_ms=provider_response.duration_ms,
+        )
 
     except Exception as e:
         # Return error message
@@ -392,9 +522,9 @@ async def query_llm(
             extra={
                 "model": getattr(model_profile, "model", None),
                 "model_pointer": model,
-                "provider": getattr(model_profile.provider, "value", None)
-                if model_profile
-                else None,
+                "provider": (
+                    getattr(model_profile.provider, "value", None) if model_profile else None
+                ),
             },
         )
         duration_ms = (time.time() - start_time) * 1000
@@ -410,7 +540,7 @@ async def query(
     system_prompt: str,
     context: Dict[str, str],
     query_context: QueryContext,
-    can_use_tool_fn: Optional[Any] = None,
+    can_use_tool_fn: Optional[ToolPermissionCallable] = None,
 ) -> AsyncGenerator[Union[UserMessage, AssistantMessage, ProgressMessage], None]:
     """Execute a query with tool support.
 
@@ -442,6 +572,9 @@ async def query(
     # Work on a copy so external mutations (e.g., UI appending messages while consuming)
     # do not interfere with recursion or normalization.
     messages = list(messages)
+    change_notices = detect_changed_files(query_context.file_state_cache)
+    if change_notices:
+        messages.append(create_user_message(_format_changed_file_notice(change_notices)))
     model_profile = resolve_model_profile(query_context.model)
     tool_mode = determine_tool_mode(model_profile)
     tools_for_model: List[Tool[Any, Any]] = [] if tool_mode == "text" else query_context.all_tools()
@@ -458,15 +591,74 @@ async def query(
         },
     )
 
-    assistant_message = await query_llm(
-        messages,
-        full_system_prompt,
-        tools_for_model,
-        query_context.max_thinking_tokens,
-        query_context.model,
-        query_context.abort_controller,
+    progress_queue: asyncio.Queue[Optional[ProgressMessage]] = asyncio.Queue()
+
+    async def _stream_progress(chunk: str) -> None:
+        if not chunk:
+            return
+        try:
+            await progress_queue.put(
+                create_progress_message(
+                    tool_use_id="stream",
+                    sibling_tool_use_ids=set(),
+                    content=chunk,
+                )
+            )
+        except Exception:
+            logger.exception("[query] Failed to enqueue stream progress chunk")
+
+    assistant_task = asyncio.create_task(
+        query_llm(
+            messages,
+            full_system_prompt,
+            tools_for_model,
+            query_context.max_thinking_tokens,
+            query_context.model,
+            query_context.abort_controller,
+            progress_callback=_stream_progress,
+            request_timeout=DEFAULT_REQUEST_TIMEOUT_SEC,
+            max_retries=MAX_LLM_RETRIES,
+            stream=True,
+        )
     )
 
+    assistant_message: Optional[AssistantMessage] = None
+
+    while True:
+        if query_context.abort_controller.is_set():
+            assistant_task.cancel()
+            try:
+                await assistant_task
+            except CancelledError:
+                pass
+            yield create_assistant_message(INTERRUPT_MESSAGE)
+            return
+        if assistant_task.done():
+            assistant_message = await assistant_task
+            break
+        try:
+            progress = progress_queue.get_nowait()
+        except asyncio.QueueEmpty:
+            waiter = asyncio.create_task(progress_queue.get())
+            done, pending = await asyncio.wait(
+                {assistant_task, waiter}, return_when=asyncio.FIRST_COMPLETED
+            )
+            if assistant_task in done:
+                for task in pending:
+                    task.cancel()
+                assistant_message = await assistant_task
+                break
+            progress = waiter.result()
+        if progress:
+            yield progress
+
+    while not progress_queue.empty():
+        residual = progress_queue.get_nowait()
+        if residual:
+            yield residual
+
+    assert assistant_message is not None
+
     # Check for abort
     if query_context.abort_controller.is_set():
         yield create_assistant_message(INTERRUPT_MESSAGE)
@@ -474,7 +666,7 @@ async def query(
 
     yield assistant_message
 
-    tool_use_blocks = extract_tool_use_blocks(assistant_message)
+    tool_use_blocks: List[MessageContent] = extract_tool_use_blocks(assistant_message)
     text_blocks = (
         len(assistant_message.message.content)
         if isinstance(assistant_message.message.content, list)
@@ -495,6 +687,7 @@ async def query(
     sibling_ids = set(
         getattr(t, "tool_use_id", None) or getattr(t, "id", None) or "" for t in tool_use_blocks
     )
+    prepared_calls: List[Dict[str, Any]] = []
 
     for tool_use in tool_use_blocks:
         tool_name = tool_use.name
@@ -511,14 +704,6 @@ async def query(
             continue
         assert tool is not None
 
-        tool_context = ToolUseContext(
-            safe_mode=query_context.safe_mode,
-            verbose=query_context.verbose,
-            permission_checker=can_use_tool_fn,
-            tool_registry=query_context.tool_registry,
-            abort_signal=query_context.abort_controller,
-        )
-
         try:
             parsed_input = tool.input_schema(**tool_input)
             logger.debug(
@@ -526,6 +711,17 @@ async def query(
                 f"{str(parsed_input)[:500]}"
             )
 
+            tool_context = ToolUseContext(
+                safe_mode=query_context.safe_mode,
+                verbose=query_context.verbose,
+                permission_checker=can_use_tool_fn,
+                tool_registry=query_context.tool_registry,
+                file_state_cache=query_context.file_state_cache,
+                abort_signal=query_context.abort_controller,
+                pause_ui=query_context.pause_ui,
+                resume_ui=query_context.resume_ui,
+            )
+
             validation = await tool.validate_input(parsed_input, tool_context)
             if not validation.result:
                 logger.debug(
@@ -555,26 +751,19 @@ async def query(
                     permission_denied = True
                     break
 
-            async for output in tool.call(parsed_input, tool_context):
-                if isinstance(output, ToolProgress):
-                    progress = create_progress_message(
-                        tool_use_id=tool_use_id,
-                        sibling_tool_use_ids=sibling_ids,
-                        content=output.content,
-                    )
-                    yield progress
-                    logger.debug(f"[query] Progress from tool_use_id={tool_use_id}: {output.content}")
-                elif isinstance(output, ToolResult):
-                    result_content = output.result_for_assistant or str(output.data)
-                    result_msg = tool_result_message(
-                        tool_use_id, result_content, tool_use_result=output.data
-                    )
-                    tool_results.append(result_msg)
-                    yield result_msg
-                    logger.debug(
-                        f"[query] Tool completed tool_use_id={tool_use_id} name={tool_name} "
-                        f"result_len={len(result_content)}"
-                    )
+            prepared_calls.append(
+                {
+                    "is_concurrency_safe": tool.is_concurrency_safe(),
+                    "generator": _run_tool_use_generator(
+                        tool,
+                        tool_use_id,
+                        tool_name,
+                        parsed_input,
+                        sibling_ids,
+                        tool_context,
+                    ),
+                }
+            )
 
         except ValidationError as ve:
             detail_text = format_pydantic_errors(ve)
@@ -600,6 +789,13 @@ async def query(
         if permission_denied:
             break
 
+    if permission_denied:
+        return
+
+    if prepared_calls:
+        async for message in _run_tools_concurrently(prepared_calls, tool_results):
+            yield message
+
     # Check for abort after tools
     if query_context.abort_controller.is_set():
         yield create_assistant_message(INTERRUPT_MESSAGE_FOR_TOOL_USE)