File size: 8,842 Bytes
399b80c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
"""
OpenSpace UI Integration

Integrates the UI system with OpenSpace core components.
Provides hooks and callbacks to update UI in real-time.
"""

import asyncio
from typing import Optional

from openspace.utils.ui import OpenSpaceUI, AgentStatus
from openspace.utils.logging import Logger

logger = Logger.get_logger(__name__)


class UIIntegration:
    """
    UI Integration for OpenSpace
    
    Connects OpenSpace components with the UI system to provide real-time
    visualization of agent activities and execution flow.
    """
    
    def __init__(self, ui: OpenSpaceUI):
        """
        Initialize UI integration
        
        Args:
            ui: OpenSpaceUI instance
        """
        self.ui = ui
        self._update_task: Optional[asyncio.Task] = None
        self._running = False
        
        # Tracked components
        self._llm_client = None
        self._grounding_client = None
    
    def attach_llm_client(self, llm_client):
        """
        Attach LLM client
        
        Args:
            llm_client: LLMClient instance
        """
        self._llm_client = llm_client
        logger.debug("UI attached to LLMClient")
    
    def attach_grounding_client(self, grounding_client):
        """
        Attach grounding client
        
        Args:
            grounding_client: GroundingClient instance
        """
        self._grounding_client = grounding_client
        logger.debug("UI attached to GroundingClient")
    
    async def start_monitoring(self, poll_interval: float = 0.5):
        """
        Start monitoring and updating UI
        
        Args:
            poll_interval: Update interval in seconds
        """
        if self._running:
            logger.warning("UI monitoring already running")
            return
        
        self._running = True
        
        # Immediately update UI once before starting the loop
        await self._update_ui()
        
        self._update_task = asyncio.create_task(
            self._monitor_loop(poll_interval)
        )
        logger.info("UI monitoring started")
    
    async def stop_monitoring(self):
        """Stop monitoring"""
        if not self._running:
            return
        
        self._running = False
        
        if self._update_task:
            self._update_task.cancel()
            try:
                await self._update_task
            except asyncio.CancelledError:
                pass
        
        logger.info("UI monitoring stopped")
    
    async def _monitor_loop(self, poll_interval: float):
        """
        Main monitoring loop
        
        Args:
            poll_interval: Update interval in seconds
        """
        while self._running:
            try:
                await self._update_ui()
                await asyncio.sleep(poll_interval)
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"UI update error: {e}", exc_info=True)
    
    async def _update_ui(self):
        """Update UI with current state"""
        # Update grounding backends info
        if self._grounding_client:
            backends = []
            try:
                # Get list of providers
                providers = self._grounding_client.list_providers()
                
                for backend_type, provider in providers.items():
                    backend_name = backend_type.value if hasattr(backend_type, 'value') else str(backend_type)
                    
                    backend_info = {
                        "name": backend_name,
                        "type": backend_name,  # gui, shell, mcp, system, web
                        "servers": []
                    }
                    
                    # For MCP provider, get server names
                    if backend_name == "mcp":
                        try:
                            # Try to get MCP sessions from provider
                            if hasattr(provider, '_sessions'):
                                backend_info["servers"] = list(provider._sessions.keys())
                        except Exception:
                            pass
                    
                    backends.append(backend_info)
                
                self.ui.update_grounding_backends(backends)
            except Exception as e:
                logger.debug(f"Failed to update grounding backends: {e}")
        
        # Refresh display
        self.ui.update_display()
    
    # Event handlers - to be called by agents
    
    def on_agent_start(self, agent_name: str, activity: str):
        """
        Called when agent starts an activity
        
        Args:
            agent_name: Agent name
            activity: Activity description
        """
        self.ui.update_agent_status(agent_name, AgentStatus.EXECUTING)
        self.ui.add_agent_activity(agent_name, activity)
        self.ui.add_log(f"{agent_name}: {activity}", level="info")
    
    def on_agent_thinking(self, agent_name: str):
        """
        Called when agent is thinking
        
        Args:
            agent_name: Agent name
        """
        self.ui.update_agent_status(agent_name, AgentStatus.THINKING)
    
    def on_agent_complete(self, agent_name: str, result: str = ""):
        """
        Called when agent completes an activity
        
        Args:
            agent_name: Agent name
            result: Result description
        """
        self.ui.update_agent_status(agent_name, AgentStatus.IDLE)
        if result:
            self.ui.add_log(f"{agent_name}: {result}", level="success")
    
    def on_llm_call(self, model: str, prompt_length: int):
        """
        Called when LLM is called
        
        Args:
            model: Model name
            prompt_length: Prompt length
        """
        self.ui.update_metrics(
            llm_calls=self.ui.metrics.get("llm_calls", 0) + 1
        )
        self.ui.add_log(f"LLM call: {model} (prompt: {prompt_length} chars)", level="debug")
    
    def on_grounding_call(self, backend: str, action: str):
        """
        Called when grounding backend is called
        
        Args:
            backend: Backend name
            action: Action description
        """
        self.ui.add_grounding_operation(backend, action, status="pending")
        self.ui.add_log(f"Grounding [{backend}]: {action}", level="info")
    
    def on_grounding_complete(self, backend: str, action: str, success: bool):
        """
        Called when grounding operation completes
        
        Args:
            backend: Backend name
            action: Action description
            success: Whether operation succeeded
        """
        status = "success" if success else "error"
        
        # Update last operation status
        for op in reversed(self.ui.grounding_operations):
            if op["backend"] == backend and op["action"] == action and op["status"] == "pending":
                op["status"] = status
                break
        
        level = "success" if success else "error"
        result = "✓" if success else "✗"
        self.ui.add_log(f"Grounding [{backend}]: {action} {result}", level=level)
    
    
    def on_iteration(self, iteration: int):
        """
        Called on each iteration
        
        Args:
            iteration: Iteration number
        """
        self.ui.update_metrics(iterations=iteration)
    
    def on_error(self, message: str):
        """
        Called when an error occurs
        
        Args:
            message: Error message
        """
        self.ui.add_log(f"ERROR: {message}", level="error")


class UILoggingHandler:
    """
    Logging handler that forwards logs to UI
    """
    
    def __init__(self, ui: OpenSpaceUI):
        """
        Initialize logging handler
        
        Args:
            ui: OpenSpaceUI instance
        """
        self.ui = ui
    
    def emit(self, record):
        """
        Emit a log record to UI
        
        Args:
            record: Log record
        """
        level_map = {
            "DEBUG": "debug",
            "INFO": "info",
            "WARNING": "warning",
            "ERROR": "error",
            "CRITICAL": "error",
        }
        
        level = level_map.get(record.levelname, "info")
        message = record.getMessage()
        
        # Filter out noisy logs
        if any(skip in message.lower() for skip in ["processing card", "workflow poll"]):
            return
        
        self.ui.add_log(message, level=level)


def create_integration(ui: OpenSpaceUI) -> UIIntegration:
    """
    Create UI integration instance
    
    Args:
        ui: OpenSpaceUI instance
        
    Returns:
        UIIntegration instance
    """
    return UIIntegration(ui)