File size: 5,430 Bytes
c52954a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import asyncio
import json
import logging
from datetime import datetime
from typing import AsyncGenerator, List, Dict, Any, Tuple
from pinchtab_browser import PinchTabBrowser
from web_agent_config import WEB_AGENT_SYSTEM_PROMPT, WEB_TOOL_CONTENT
import shared_utils
import app  # Still needed for generate_response and MAX_NEW_TOKENS for now

logger = logging.getLogger(__name__)

async def run_web_agent_streaming(
    goal: str,
    max_rounds: int = 15,
    inference_mode: str = "Auto (local + cloud backup)"
) -> AsyncGenerator[str, None]:
    mode = app._INFERENCE_MODE_MAP.get(inference_mode, "auto")
    browser = PinchTabBrowser()
    html_parts = [shared_utils.render_user_message(goal)]
    yield ''.join(html_parts)

    system_prompt = WEB_AGENT_SYSTEM_PROMPT.format(current_date=datetime.now().strftime('%Y-%m-%d'))
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": goal}
    ]

    tools = WEB_TOOL_CONTENT
    context_budget = await app.get_context_budget(mode)

    try:
        for round_num in range(1, max_rounds + 1):
            html_parts.append(shared_utils.render_round_badge(round_num, max_rounds))
            yield ''.join(html_parts)

            # Context management
            if shared_utils.estimate_tokens(messages) > context_budget:
                messages = shared_utils.compress_messages(messages, keep_last_rounds=3)

            html_parts.append('<div class="thinking-streaming">Processing...</div>')
            yield ''.join(html_parts)

            try:
                api_result = await app.generate_response(messages, tools, max_new_tokens=app.MAX_NEW_TOKENS, mode=mode)
                html_parts.pop() # Remove placeholder
            except Exception as e:
                html_parts.pop()
                html_parts.append(f"<p style='color:#dc2626;'>Generation Error: {str(e)}</p>")
                yield ''.join(html_parts)
                return

            generated = api_result["content"]
            reasoning, content = shared_utils.extract_thinking(generated)
            tool_call, clean_content = shared_utils.parse_tool_call(content)

            if reasoning:
                html_parts.append(shared_utils.render_thinking_collapsed(reasoning))
                yield ''.join(html_parts)

            if tool_call:
                fn_name = tool_call.get("name", "unknown")
                args = tool_call.get("arguments", {})
                html_parts.append(shared_utils.render_tool_call(fn_name, args))
                yield ''.join(html_parts)

            if clean_content.strip() and not tool_call:
                html_parts.append(f'<div class="answer-section">{clean_content}</div>')
                yield ''.join(html_parts)

            messages.append({
                "role": "assistant",
                "content": clean_content if tool_call is None else "",
                "reasoning_content": reasoning,
                "tool_calls": [{
                    "id": str(round_num),
                    "type": "function",
                    "function": {
                        "name": tool_call.get("name", ""),
                        "arguments": tool_call.get("arguments", {})
                    }
                }] if tool_call else None
            })

            if tool_call:
                fn_name = tool_call.get("name", "")
                actual_fn = fn_name.split(".", 1)[1] if "." in fn_name else fn_name
                args = tool_call.get("arguments", {})

                result = ""
                try:
                    if actual_fn == "navigate":
                        res = await browser.navigate(args.get("url"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "snapshot":
                        res = await browser.snapshot(args.get("filter", "interactive"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "click":
                        res = await browser.action("click", ref=args.get("ref"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "type":
                        res = await browser.action("type", ref=args.get("ref"), text=args.get("text"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "press":
                        res = await browser.action("press", key=args.get("key"))
                        result = json.dumps(res, indent=2)
                    elif actual_fn == "get_text":
                        res = await browser.get_text()
                        result = json.dumps(res, indent=2)
                    else:
                        result = f"Unknown tool: {fn_name}"
                except Exception as e:
                    result = f"Tool error: {str(e)}"

                html_parts.append(shared_utils.render_tool_result(result, fn_name))
                yield ''.join(html_parts)

                messages.append({
                    "role": "tool",
                    "tool_call_id": str(round_num),
                    "content": result
                })
                continue

            if shared_utils.is_final_answer(generated):
                html_parts.append(shared_utils.render_completion())
                yield ''.join(html_parts)
                break
    finally:
        await browser.close()