aria-code 4.1.3__py3-none-any.whl

apps/cli/providers/llm/sse_stream.py

@@ -0,0 +1,216 @@
+"""stream_chat — Aria cloud SSE provider extracted from aria_cli.py.
+
+Streams AI responses from the Aria backend via Server-Sent Events (SSE).
+Supports cancellation, thinking tokens, tool calls, retries, and usage stats.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Callable, Optional
+
+
+async def stream_chat(
+    base_url: str,
+    message: str,
+    history: list,
+    model: str = "qwen2.5:7b",
+    thinking_mode: str = "auto",
+    user_context: Optional[dict] = None,
+    auth_token: Optional[str] = None,
+    on_token: Optional[Callable[[str], None]] = None,
+    on_thinking: Optional[Callable[[str], None]] = None,
+    on_tool_call: Optional[Callable[[str, dict], None]] = None,
+    on_tool_result: Optional[Callable[[str, str], None]] = None,
+    on_status: Optional[Callable[[str, str], None]] = None,
+    cancel_event: Optional[asyncio.Event] = None,
+    project_context: str = "",
+) -> dict:
+    """Stream AI chat via SSE with cancel support and user context.
+
+    Parameters
+    ----------
+    project_context:
+        ARIA.md / CLAUDE.md content to inject into user_context.
+        Callers pass ``_PROJECT_CONTEXT`` from aria_cli, keeping this
+        module free of global state.
+    """
+    import aiohttp
+
+    url = f"{base_url}/api/v2/ai/chat/stream"
+
+    payload: dict = {
+        "message": message,
+        "conversation_history": history[-20:],
+        "model": model,
+        "thinking_mode": thinking_mode,
+        "stream": True,
+    }
+    if user_context:
+        if project_context:
+            user_context = {**user_context, "project_context": project_context}
+        payload["user_context"] = user_context
+
+    headers: dict = {}
+    if auth_token:
+        headers["Authorization"] = f"Bearer {auth_token}"
+
+    full_response = ""
+    thinking_content = ""
+    tools_used: list = []
+    sources: list = []
+    tool_calls_pending: list = []
+    usage: dict = {"prompt_tokens": 0, "completion_tokens": 0, "thinking_tokens": 0}
+
+    _max_connect_retries = 2
+    _last_connect_error: Optional[str] = None
+
+    for _attempt in range(_max_connect_retries + 1):
+        if cancel_event and cancel_event.is_set():
+            return {
+                "success": True, "response": "", "cancelled": True,
+                "tools_used": [], "sources": [], "usage": usage,
+            }
+        # Reset per-attempt accumulators
+        full_response = ""
+        thinking_content = ""
+        tools_used = []
+        sources = []
+        tool_calls_pending = []
+        usage = {"prompt_tokens": 0, "completion_tokens": 0, "thinking_tokens": 0}
+
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    url, json=payload, headers=headers,
+                    timeout=aiohttp.ClientTimeout(total=120),
+                ) as resp:
+                    if resp.status != 200:
+                        error_text = await resp.text()
+                        return {"success": False, "error": f"HTTP {resp.status}: {error_text[:200]}"}
+
+                    buffer = ""
+                    event_type = "delta"
+
+                    async for chunk in resp.content:
+                        if cancel_event and cancel_event.is_set():
+                            try:
+                                await session.post(
+                                    f"{base_url}/api/v2/ai/chat/cancel",
+                                    headers=headers,
+                                    timeout=aiohttp.ClientTimeout(total=3),
+                                )
+                            except Exception:
+                                pass
+                            return {
+                                "success": True, "response": full_response,
+                                "cancelled": True, "tools_used": tools_used,
+                                "sources": sources, "usage": usage,
+                            }
+
+                        text = chunk.decode("utf-8", errors="ignore")
+                        buffer += text
+
+                        while "\n" in buffer:
+                            line, buffer = buffer.split("\n", 1)
+                            line = line.strip()
+
+                            if not line or line.startswith(":"):
+                                continue
+                            if line.startswith("event:"):
+                                event_type = line[6:].strip()
+                                continue
+                            if not line.startswith("data:"):
+                                continue
+
+                            data_str = line[5:].strip()
+                            if data_str == "[DONE]":
+                                break
+                            try:
+                                data = json.loads(data_str)
+                            except json.JSONDecodeError:
+                                continue
+
+                            # Backend error: {"success": false, "error": "..."}
+                            if data.get("success") is False:
+                                err_msg = data.get("error", "Backend error")
+                                return {"success": False, "error": f"Backend: {err_msg}"}
+
+                            evt = data.get("type", event_type)
+
+                            if evt == "delta":
+                                token = data.get("text", data.get("content", ""))
+                                if token:
+                                    full_response += token
+                                    usage["completion_tokens"] += 1
+                                    if on_token:
+                                        on_token(token)
+
+                            elif evt == "thinking_content":
+                                tc = data.get("content", "")
+                                if tc:
+                                    thinking_content += tc
+                                    usage["thinking_tokens"] += 1
+                                    if on_thinking:
+                                        on_thinking(tc)
+
+                            elif evt == "tool_call":
+                                tool = data.get("tool", data.get("name", ""))
+                                params = data.get("params", {})
+                                tools_used.append(tool)
+                                tool_calls_pending.append({"tool": tool, "params": params})
+                                if on_tool_call:
+                                    on_tool_call(tool, params)
+
+                            elif evt == "tool_result":
+                                if on_tool_result:
+                                    on_tool_result(data.get("tool", ""), data.get("summary", ""))
+
+                            elif evt == "status":
+                                if on_status:
+                                    on_status(data.get("state", ""), data.get("message", ""))
+
+                            elif evt == "final":
+                                full_response = data.get("answer", full_response)
+                                sources = data.get("sources", [])
+                                if data.get("usage"):
+                                    u = data["usage"]
+                                    usage["prompt_tokens"] = u.get("prompt_tokens", usage["prompt_tokens"])
+                                    usage["completion_tokens"] = u.get("completion_tokens", usage["completion_tokens"])
+
+                            elif evt == "error":
+                                return {"success": False, "error": data.get("message", "Unknown error")}
+
+            return {
+                "success": True,
+                "response": full_response,
+                "thinking": thinking_content,
+                "tools_used": tools_used,
+                "sources": sources,
+                "tool_calls_pending": tool_calls_pending,
+                "usage": usage,
+            }
+
+        except asyncio.TimeoutError:
+            return {"success": False, "error": "Request timed out (120s)"}
+        except asyncio.CancelledError:
+            return {
+                "success": True, "response": full_response, "cancelled": True,
+                "tools_used": tools_used, "sources": sources, "usage": usage,
+            }
+        except aiohttp.ClientConnectorError as exc:
+            _last_connect_error = str(exc)
+            if _attempt < _max_connect_retries:
+                wait = 1.5 * (_attempt + 1)
+                await asyncio.sleep(wait)
+                if on_status:
+                    on_status("retry", f"Connection failed, retrying ({_attempt + 2}/{_max_connect_retries + 1})...")
+                continue
+            break
+        except Exception as exc:
+            return {"success": False, "error": str(exc)}
+
+    return {
+        "success": False,
+        "error": f"Connection failed after {_max_connect_retries + 1} attempts: {_last_connect_error}",
+    }

apps/cli/providers/runtime_bridge.py

@@ -0,0 +1,185 @@
+"""Phase-3 bridge: run one chat turn through the shared ``runtime.run_agent``.
+
+The documented runtime next step is to route the CLI tool loop through
+``run_agent`` and keep aria_cli as orchestration glue. This module is that engine
+— the two adapters run_agent needs plus a thin driver:
+
+  • build_tool_executor() — wraps the CLI's ``LOCAL_TOOLS`` ({name: (handler, schema)})
+  • make_provider_fn()     — selects the provider (chat_routing) + streams it
+  • run_chat_via_runtime() — drives run_agent, renders via callbacks, returns text
+
+It is opt-in: ``send_message`` only uses it when ``config['use_runtime_loop']`` is
+on, and falls back to the proven inline loop on any error. That keeps the live
+path untouched until the runtime path is verified in a real REPL.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, List, Optional
+
+from .chat_routing import first_round_route, is_placeholder_response, should_fallback
+
+
+def build_tool_executor(local_tools, config: Optional[dict] = None):
+    """Wrap the CLI's LOCAL_TOOLS registry for run_agent."""
+    from runtime.tool_executor import ToolExecutor
+    return ToolExecutor(local_tools, config=config or {})
+
+
+async def run_with_fallback(
+    route: str,
+    *,
+    run_cloud: Callable,
+    run_ollama: Callable,
+    on_token: Optional[Callable[[str], None]] = None,
+) -> dict:
+    """Primary generation per ``route``, with cloud→Ollama fallback parity.
+
+    Mirrors ``send_message``'s inline fallback, but keyed on *route* (via
+    ``should_fallback``) so a genuinely-good forced-backend answer is kept
+    instead of being discarded and re-run:
+
+      • ``skip`` / ``ollama`` → run local Ollama directly (no cloud round)
+      • ``cloud``            → run cloud; if it fails or returns a placeholder
+                               (empty / canned / backend stub), fall back to
+                               local Ollama
+
+    ``run_cloud`` / ``run_ollama`` are async ``(on_token) -> result dict``
+    closures; injecting them keeps this orchestration unit-testable without
+    real providers or network.
+    """
+    if route in ("skip", "ollama"):
+        return await run_ollama(on_token)
+
+    # route == "cloud": count streamed tokens so a long-but-unstreamed canned
+    # backend reply is recognised as a placeholder, not a real generation.
+    _tokens = [0]
+
+    def _counting_on_token(tok: str) -> None:
+        _tokens[0] += 1
+        if on_token is not None:
+            on_token(tok)
+
+    result = await run_cloud(_counting_on_token)
+    if result.get("cancelled"):
+        return result  # user-cancelled — never silently re-run on a different provider
+
+    placeholder = is_placeholder_response(result.get("response", ""), _tokens[0])
+    if should_fallback("cloud", result, is_placeholder=placeholder):
+        return await run_ollama(on_token)
+    return result
+
+
+def make_provider_fn(
+    *,
+    model: str,
+    config: dict,
+    api_url: Optional[str],
+    ollama_url: str,
+    tool_schemas: List[dict],
+    thinking_mode: str = "auto",
+    user_context: Optional[dict] = None,
+    auth_token: Optional[str] = None,
+    project_context: Any = None,
+    system_override: Optional[str] = None,
+) -> Callable:
+    """Build an async ``provider_fn`` for run_agent.
+
+    Selects the provider per chat_routing (cloud → AriaSSE backend; ollama/skip →
+    local Ollama) and streams it through the shared ``stream_provider_result``.
+    A pending system-role override is threaded the same way ``send_message`` does
+    it: cloud via ``user_context['system_role_override']``, Ollama via the
+    provider's ``system_override`` argument.
+    """
+    from apps.cli.providers.base import AriaSSEProvider, OllamaProvider
+    from packages.aria_sdk.streaming import stream_provider_result
+
+    _cloud_uctx = dict(user_context or {})
+    if system_override:
+        _cloud_uctx["system_role_override"] = system_override
+
+    async def _provider_fn(prompt, history, *, on_token=None, on_thinking=None,
+                           on_tool_call=None, on_tool_result=None, on_status=None,
+                           cancel_event=None):
+        route = first_round_route(model, config, api_url)
+
+        async def _stream(provider, _on_token):
+            return await stream_provider_result(
+                provider, prompt, history, tools=tool_schemas,
+                cancel_event=cancel_event, on_token=_on_token, on_thinking=on_thinking,
+                on_tool_call=on_tool_call, on_tool_result=on_tool_result, on_status=on_status,
+            )
+
+        async def run_cloud(_on_token):
+            return await _stream(
+                AriaSSEProvider(
+                    api_url, model, thinking_mode=thinking_mode,
+                    user_context=_cloud_uctx, auth_token=auth_token,
+                    project_context=project_context,
+                ),
+                _on_token,
+            )
+
+        async def run_ollama(_on_token):
+            return await _stream(
+                OllamaProvider(ollama_url, model, system_override=system_override),
+                _on_token,
+            )
+
+        return await run_with_fallback(
+            route, run_cloud=run_cloud, run_ollama=run_ollama, on_token=on_token,
+        )
+
+    return _provider_fn
+
+
+async def run_chat_via_runtime(
+    *,
+    prompt: str,
+    history: list,
+    local_tools,
+    tool_schemas: List[dict],
+    model: str,
+    config: dict,
+    api_url: Optional[str],
+    ollama_url: str,
+    cancel_event=None,
+    on_token: Optional[Callable[[str], None]] = None,
+    on_thinking: Optional[Callable[[str], None]] = None,
+    on_tool_call: Optional[Callable[[str, dict], None]] = None,
+    on_tool_result: Optional[Callable[[str, dict], None]] = None,
+    on_status: Optional[Callable[[str, str], None]] = None,
+    thinking_mode: str = "auto",
+    user_context: Optional[dict] = None,
+    auth_token: Optional[str] = None,
+    project_context: Any = None,
+    system_override: Optional[str] = None,
+    max_rounds: int = 30,
+) -> str:
+    """Run one chat turn through the shared runtime Gateway; return the text.
+
+    This is the CLI *adapter* for ``runtime.gateway.run_turn``: it builds the
+    CLI's ``provider_fn`` (AriaSSE/Ollama selection + cloud→Ollama fallback) and
+    tool executor (the LOCAL_TOOLS registry), then hands them to the neutral
+    gateway, which drives ``run_agent`` and streams via the callbacks. Returns
+    the assistant response text ("" if none).
+    """
+    from runtime.gateway import run_turn
+
+    provider_fn = make_provider_fn(
+        model=model, config=config, api_url=api_url, ollama_url=ollama_url,
+        tool_schemas=tool_schemas, thinking_mode=thinking_mode,
+        user_context=user_context, auth_token=auth_token, project_context=project_context,
+        system_override=system_override,
+    )
+    executor = build_tool_executor(local_tools, config)
+
+    result = await run_turn(
+        prompt, history,
+        provider_fn=provider_fn, tool_executor=executor,
+        tool_schemas=list(tool_schemas),
+        on_token=on_token, on_thinking=on_thinking,
+        on_tool_call=on_tool_call, on_tool_result=on_tool_result, on_status=on_status,
+        cancel_event=cancel_event, max_rounds=max_rounds,
+    )
+    return result.text