alita-sdk 0.3.486__py3-none-any.whl → 0.3.515__py3-none-any.whl

alita_sdk/runtime/tools/llm.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from traceback import format_exc
-from typing import Any, Optional, List, Union
+from typing import Any, Optional, List, Union, Literal
 
 from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
 from langchain_core.runnables import RunnableConfig
@@ -34,6 +34,7 @@ class LLMNode(BaseTool):
     available_tools: Optional[List[BaseTool]] = Field(default=None, description='Available tools for binding')
     tool_names: Optional[List[str]] = Field(default=None, description='Specific tool names to filter')
     steps_limit: Optional[int] = Field(default=25, description='Maximum steps for tool execution')
+    tool_execution_timeout: Optional[int] = Field(default=900, description='Timeout (seconds) for tool execution. Default is 15 minutes.')
 
     def get_filtered_tools(self) -> List[BaseTool]:
         """
@@ -129,7 +130,9 @@ class LLMNode(BaseTool):
         # or standalone LLM node for chat (with messages only)
         if 'system' in func_args.keys():
             # Flow for LLM node with prompt/task from pipeline
-            if not func_args.get('system') or not func_args.get('task'):
+            if func_args.get('system') is None or func_args.get('task') is None:
+                raise ToolException(f"LLMNode requires 'system' and 'task' parameters in input mapping. "
+                                    f"Actual params: {func_args}")
                 raise ToolException(f"LLMNode requires 'system' and 'task' parameters in input mapping. "
                                     f"Actual params: {func_args}")
             # cast to str in case user passes variable different from str
@@ -171,26 +174,36 @@ class LLMNode(BaseTool):
                     for key, value in (self.structured_output_dict or {}).items()
                 }
                 # Add default output field for proper response to user
-                struct_params['elitea_response'] = {'description': 'final output to user', 'type': 'str'}
+                struct_params['elitea_response'] = {
+                    'description': 'final output to user (summarized output from LLM)', 'type': 'str',
+                    "default": None}
                 struct_model = create_pydantic_model(f"LLMOutput", struct_params)
-                completion = llm_client.invoke(messages, config=config)
-                if hasattr(completion, 'tool_calls') and completion.tool_calls:
+                initial_completion = llm_client.invoke(messages, config=config)
+                if hasattr(initial_completion, 'tool_calls') and initial_completion.tool_calls:
                     new_messages, _ = self._run_async_in_sync_context(
-                        self.__perform_tool_calling(completion, messages, llm_client, config)
+                        self.__perform_tool_calling(initial_completion, messages, llm_client, config)
                     )
                     llm = self.__get_struct_output_model(llm_client, struct_model)
                     completion = llm.invoke(new_messages, config=config)
                     result = completion.model_dump()
                 else:
-                    llm = self.__get_struct_output_model(llm_client, struct_model)
-                    completion = llm.invoke(messages, config=config)
+                    try:
+                        llm = self.__get_struct_output_model(llm_client, struct_model)
+                        completion = llm.invoke(messages, config=config)
+                    except ValueError as e:
+                        logger.error(f"Error invoking structured output model: {format_exc()}")
+                        logger.info("Attemping to fall back to json mode")
+                        # Fallback to regular LLM with JSON extraction
+                        completion = self.__get_struct_output_model(llm_client, struct_model,
+                                                                    method="json_mode").invoke(messages, config=config)
                     result = completion.model_dump()
 
                 # Ensure messages are properly formatted
                 if result.get('messages') and isinstance(result['messages'], list):
                     result['messages'] = [{'role': 'assistant', 'content': '\n'.join(result['messages'])}]
                 else:
-                    result['messages'] = messages + [AIMessage(content=result.get(ELITEA_RS, ''))]
+                    result['messages'] = messages + [
+                        AIMessage(content=result.get(ELITEA_RS, '') or initial_completion.content)]
 
                 return result
             else:
@@ -243,40 +256,146 @@ class LLMNode(BaseTool):
         
         For MCP tools with persistent sessions, we reuse the same event loop
         that was used to create the MCP client and sessions (set by CLI).
+
+        When called from within a running event loop (e.g., nested LLM nodes),
+        we need to handle this carefully to avoid "event loop already running" errors.
+
+        This method handles three scenarios:
+        1. Called from async context (event loop running) - creates new thread with new loop
+        2. Called from sync context with persistent loop - reuses persistent loop
+        3. Called from sync context without loop - creates new persistent loop
         """
+        import threading
+
+        # Check if there's a running loop
         try:
-            loop = asyncio.get_running_loop()
-            # Already in async context - run in thread with new loop
-            import threading
-            
+            running_loop = asyncio.get_running_loop()
+            loop_is_running = True
+            logger.debug(f"Detected running event loop (id: {id(running_loop)}), executing tool calls in separate thread")
+        except RuntimeError:
+            loop_is_running = False
+
+        # Scenario 1: Loop is currently running - MUST use thread
+        if loop_is_running:
             result_container = []
-            
+            exception_container = []
+
+            # Try to capture Streamlit context from current thread for propagation
+            streamlit_ctx = None
+            try:
+                from streamlit.runtime.scriptrunner import get_script_run_ctx, add_script_run_ctx
+                streamlit_ctx = get_script_run_ctx()
+                if streamlit_ctx:
+                    logger.debug("Captured Streamlit context for propagation to worker thread")
+            except (ImportError, Exception) as e:
+                logger.debug(f"Streamlit context not available or failed to capture: {e}")
+
             def run_in_thread():
+                """Run coroutine in a new thread with its own event loop."""
                 new_loop = asyncio.new_event_loop()
                 asyncio.set_event_loop(new_loop)
                 try:
-                    result_container.append(new_loop.run_until_complete(coro))
+                    result = new_loop.run_until_complete(coro)
+                    result_container.append(result)
+                except Exception as e:
+                    logger.debug(f"Exception in async thread: {e}")
+                    exception_container.append(e)
                 finally:
                     new_loop.close()
-            
-            thread = threading.Thread(target=run_in_thread)
+                    asyncio.set_event_loop(None)
+
+            thread = threading.Thread(target=run_in_thread, daemon=False)
+
+            # Propagate Streamlit context to the worker thread if available
+            if streamlit_ctx is not None:
+                try:
+                    add_script_run_ctx(thread, streamlit_ctx)
+                    logger.debug("Successfully propagated Streamlit context to worker thread")
+                except Exception as e:
+                    logger.warning(f"Failed to propagate Streamlit context to worker thread: {e}")
+
             thread.start()
-            thread.join()
+            thread.join(timeout=self.tool_execution_timeout)  # 15 minute timeout for safety
+
+            if thread.is_alive():
+                logger.error("Async operation timed out after 5 minutes")
+                raise TimeoutError("Async operation in thread timed out")
+
+            # Re-raise exception if one occurred
+            if exception_container:
+                raise exception_container[0]
+
             return result_container[0] if result_container else None
-            
-        except RuntimeError:
-            # No event loop running - use/create persistent loop
-            # This loop is shared with MCP session creation for stateful tools
+
+        # Scenario 2 & 3: No loop running - use or create persistent loop
+        else:
+            # Get or create persistent loop
             if not hasattr(self.__class__, '_persistent_loop') or \
                self.__class__._persistent_loop is None or \
                self.__class__._persistent_loop.is_closed():
                 self.__class__._persistent_loop = asyncio.new_event_loop()
                 logger.debug("Created persistent event loop for async tools")
-            
+
             loop = self.__class__._persistent_loop
-            asyncio.set_event_loop(loop)
-            return loop.run_until_complete(coro)
-    
+
+            # Double-check the loop is not running (safety check)
+            if loop.is_running():
+                logger.debug("Persistent loop is unexpectedly running, using thread execution")
+
+                result_container = []
+                exception_container = []
+
+                # Try to capture Streamlit context from current thread for propagation
+                streamlit_ctx = None
+                try:
+                    from streamlit.runtime.scriptrunner import get_script_run_ctx, add_script_run_ctx
+                    streamlit_ctx = get_script_run_ctx()
+                    if streamlit_ctx:
+                        logger.debug("Captured Streamlit context for propagation to worker thread")
+                except (ImportError, Exception) as e:
+                    logger.debug(f"Streamlit context not available or failed to capture: {e}")
+
+                def run_in_thread():
+                    """Run coroutine in a new thread with its own event loop."""
+                    new_loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(new_loop)
+                    try:
+                        result = new_loop.run_until_complete(coro)
+                        result_container.append(result)
+                    except Exception as ex:
+                        logger.debug(f"Exception in async thread: {ex}")
+                        exception_container.append(ex)
+                    finally:
+                        new_loop.close()
+                        asyncio.set_event_loop(None)
+
+                thread = threading.Thread(target=run_in_thread, daemon=False)
+
+                # Propagate Streamlit context to the worker thread if available
+                if streamlit_ctx is not None:
+                    try:
+                        add_script_run_ctx(thread, streamlit_ctx)
+                        logger.debug("Successfully propagated Streamlit context to worker thread")
+                    except Exception as e:
+                        logger.warning(f"Failed to propagate Streamlit context to worker thread: {e}")
+
+                thread.start()
+                thread.join(timeout=self.tool_execution_timeout)
+
+                if thread.is_alive():
+                    logger.error("Async operation timed out after 15 minutes")
+                    raise TimeoutError("Async operation in thread timed out")
+
+                if exception_container:
+                    raise exception_container[0]
+
+                return result_container[0] if result_container else None
+            else:
+                # Loop exists but not running - safe to use run_until_complete
+                logger.debug(f"Using persistent loop (id: {id(loop)}) with run_until_complete")
+                asyncio.set_event_loop(loop)
+                return loop.run_until_complete(coro)
+
     async def _arun(self, *args, **kwargs):
         # Legacy async support
         return self.invoke(kwargs, **kwargs)
@@ -324,12 +443,14 @@ class LLMNode(BaseTool):
                         
                         # Try async invoke first (for MCP tools), fallback to sync
                         tool_result = None
-                        try:
-                            # Try async invocation first
-                            tool_result = await tool_to_execute.ainvoke(tool_args, config=config)
-                        except NotImplementedError:
-                            # Tool doesn't support async, use sync invoke
-                            logger.debug(f"Tool '{tool_name}' doesn't support async, using sync invoke")
+                        if hasattr(tool_to_execute, 'ainvoke'):
+                            try:
+                                tool_result = await tool_to_execute.ainvoke(tool_args, config=config)
+                            except (NotImplementedError, AttributeError):
+                                logger.debug(f"Tool '{tool_name}' ainvoke failed, falling back to sync invoke")
+                                tool_result = tool_to_execute.invoke(tool_args, config=config)
+                        else:
+                            # Sync-only tool
                             tool_result = tool_to_execute.invoke(tool_args, config=config)
 
                         # Create tool message with result - preserve structured content
@@ -539,5 +660,5 @@ class LLMNode(BaseTool):
 
         return new_messages, current_completion
 
-    def __get_struct_output_model(self, llm_client, pydantic_model):
-        return llm_client.with_structured_output(pydantic_model)
+    def __get_struct_output_model(self, llm_client, pydantic_model, method: Literal["function_calling", "json_mode", "json_schema"] = "json_schema"):
+        return llm_client.with_structured_output(pydantic_model, method=method)

alita_sdk/runtime/tools/mcp_remote_tool.py

@@ -20,10 +20,14 @@ from ..utils.mcp_oauth import (
     fetch_resource_metadata_async,
     infer_authorization_servers_from_realm,
 )
-from ..utils.mcp_sse_client import McpSseClient
+from ..utils.mcp_client import McpClient
 
 logger = logging.getLogger(__name__)
 
+# Global registry to store MCP tool session metadata by tool name
+# This is used to pass session info to callbacks since LangChain's serialization doesn't include all fields
+MCP_TOOL_SESSION_REGISTRY: Dict[str, Dict[str, Any]] = {}
+
 
 class McpRemoteTool(McpServerTool):
     """
@@ -43,6 +47,7 @@ class McpRemoteTool(McpServerTool):
         """Update metadata with session info after model initialization."""
         super().model_post_init(__context)
         self._update_metadata_with_session()
+        self._register_session_metadata()
     
     def _update_metadata_with_session(self):
         """Update the metadata dict with current session information."""
@@ -54,6 +59,15 @@ class McpRemoteTool(McpServerTool):
                 'mcp_server_url': canonical_resource(self.server_url)
             })
     
+    def _register_session_metadata(self):
+        """Register session metadata in global registry for callback access."""
+        if self.session_id and self.server_url:
+            MCP_TOOL_SESSION_REGISTRY[self.name] = {
+                'mcp_session_id': self.session_id,
+                'mcp_server_url': canonical_resource(self.server_url)
+            }
+            logger.debug(f"[MCP] Registered session metadata for tool '{self.name}': session={self.session_id}")
+    
     def __getstate__(self):
         """Custom serialization for pickle compatibility."""
         state = super().__getstate__()
@@ -85,7 +99,6 @@ class McpRemoteTool(McpServerTool):
 
     async def _execute_remote_tool(self, kwargs: Dict[str, Any]) -> str:
         """Execute the actual remote MCP tool call using SSE client."""
-        from ...tools.utils import TOOLKIT_SPLITTER
         
         # Check for session_id requirement
         if not self.session_id:
@@ -95,10 +108,10 @@ class McpRemoteTool(McpServerTool):
         # Use the original tool name from discovery for MCP server invocation
         tool_name_for_server = self.original_tool_name
         if not tool_name_for_server:
-            tool_name_for_server = self.name.rsplit(TOOLKIT_SPLITTER, 1)[-1] if TOOLKIT_SPLITTER in self.name else self.name
-            logger.warning(f"original_tool_name not set for '{self.name}', using extracted: {tool_name_for_server}")
+            tool_name_for_server = self.name
+            logger.warning(f"original_tool_name not set for '{self.name}', using: {tool_name_for_server}")
         
-        logger.info(f"[MCP SSE] Executing tool '{tool_name_for_server}' with session {self.session_id}")
+        logger.info(f"[MCP] Executing tool '{tool_name_for_server}' with session {self.session_id}")
         
         try:
             # Prepare headers
@@ -106,16 +119,18 @@ class McpRemoteTool(McpServerTool):
             if self.server_headers:
                 headers.update(self.server_headers)
             
-            # Create SSE client
-            client = McpSseClient(
+            # Create unified MCP client (auto-detects transport)
+            client = McpClient(
                 url=self.server_url,
                 session_id=self.session_id,
                 headers=headers,
                 timeout=self.tool_timeout_sec
             )
             
-            # Execute tool call via SSE
-            result = await client.call_tool(tool_name_for_server, kwargs)
+            # Execute tool call (client auto-detects SSE vs Streamable HTTP)
+            async with client:
+                await client.initialize()
+                result = await client.call_tool(tool_name_for_server, kwargs)
             
             # Format the result
             if isinstance(result, dict):
@@ -144,7 +159,7 @@ class McpRemoteTool(McpServerTool):
             return str(result)
             
         except Exception as e:
-            logger.error(f"[MCP SSE] Tool execution failed: {e}", exc_info=True)
+            logger.error(f"[MCP] Tool execution failed: {e}", exc_info=True)
             raise
 
     def _parse_sse(self, text: str) -> Dict[str, Any]:

alita_sdk/runtime/tools/mcp_server_tool.py

@@ -5,8 +5,6 @@ from typing import Any, Type, Literal, Optional, Union, List
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel, Field, create_model, EmailStr, constr, ConfigDict
 
-from ...tools.utils import TOOLKIT_SPLITTER
-
 logger = getLogger(__name__)
 
 
@@ -91,13 +89,13 @@ class McpServerTool(BaseTool):
         return create_model(model_name, **fields)
 
     def _run(self, *args, **kwargs):
-        # Extract the actual tool/prompt name (remove toolkit prefix)
+        # Use the tool name directly (no prefix extraction needed)
         call_data = {
             "server": self.server,
             "tool_timeout_sec": self.tool_timeout_sec,
             "tool_call_id": str(uuid.uuid4()),
             "params": {
-                "name": self.name.rsplit(TOOLKIT_SPLITTER)[1] if TOOLKIT_SPLITTER in self.name else self.name,
+                "name": self.name,
                 "arguments": kwargs
             }
         }

alita_sdk/runtime/tools/vectorstore_base.py

@@ -270,7 +270,7 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                 )
             ).count()
 
-    def _clean_collection(self, index_name: str = ''):
+    def _clean_collection(self, index_name: str = '', including_index_meta: bool = False):
         """
         Clean the vectorstore collection by deleting all indexed data.
         """
@@ -279,7 +279,7 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
             f"Cleaning collection '{self.dataset}'",
             tool_name="_clean_collection"
         )
-        self.vector_adapter.clean_collection(self, index_name)
+        self.vector_adapter.clean_collection(self, index_name, including_index_meta)
         self._log_tool_event(
             f"Collection '{self.dataset}' has been cleaned. ",
             tool_name="_clean_collection"
@@ -303,7 +303,7 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
         logger.info("Cleaning index before re-indexing all documents.")
         self._log_tool_event("Cleaning index before re-indexing all documents. Previous index will be removed", tool_name="index_documents")
         try:
-            self._clean_collection(index_name)
+            self._clean_collection(index_name, including_index_meta=False)
             self._log_tool_event("Previous index has been removed",
                            tool_name="index_documents")
         except Exception as e:

alita_sdk/runtime/utils/AlitaCallback.py

@@ -23,9 +23,45 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
         self.tokens_out = 0
         self.pending_llm_requests = defaultdict(int)
         self.current_model_name = 'gpt-4'
+        self._event_queue = []  # Queue for events when context is unavailable
         #
         super().__init__()
 
+    def _has_streamlit_context(self) -> bool:
+        """Check if Streamlit context is available in the current thread."""
+        try:
+            # Try to import streamlit runtime context checker
+            from streamlit.runtime.scriptrunner import get_script_run_ctx
+            ctx = get_script_run_ctx()
+            return ctx is not None
+        except (ImportError, Exception) as e:
+            if self.debug:
+                log.debug(f"Streamlit context check failed: {e}")
+            return False
+
+    def _safe_streamlit_call(self, func, *args, **kwargs):
+        """Safely execute a Streamlit UI operation, handling missing context gracefully."""
+        if not self._has_streamlit_context():
+            func_name = getattr(func, '__name__', str(func))
+            if self.debug:
+                log.warning(f"Streamlit context not available for {func_name}, queueing event")
+            # Store the event for potential replay when context is available
+            self._event_queue.append({
+                'func': func_name,
+                'args': args,
+                'kwargs': kwargs,
+                'timestamp': datetime.now(tz=timezone.utc)
+            })
+            return None
+
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            func_name = getattr(func, '__name__', str(func))
+            # Handle any Streamlit-specific exceptions gracefully
+            log.warning(f"Streamlit operation {func_name} failed: {e}")
+            return None
+
     #
     # Chain
     #
@@ -76,10 +112,14 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
             json.dumps(payload, ensure_ascii=False, default=lambda o: str(o))
         )
 
-        self.callback_state[str(run_id)] = self.st.status(
-            f"Running {payload.get('tool_name')}...", expanded=True
+        status_widget = self._safe_streamlit_call(
+            self.st.status,
+            f"Running {payload.get('tool_name')}...",
+            expanded=True
         )
-        self.callback_state[str(run_id)].write(f"Tool inputs: {payload}")
+        if status_widget:
+            self.callback_state[str(run_id)] = status_widget
+            self._safe_streamlit_call(status_widget.write, f"Tool inputs: {payload}")
 
     def on_tool_start(self, *args, run_id: UUID, **kwargs):
         """ Callback """
@@ -88,15 +128,51 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
 
         tool_name = args[0].get("name")
         tool_run_id = str(run_id)
+        
+        # Extract metadata from tool if available (from BaseAction.metadata)
+        # Try multiple sources for metadata with toolkit_name
+        tool_meta = args[0].copy()
+        
+        # Source 1: kwargs['serialized']['metadata'] - LangChain's full tool serialization
+        if 'serialized' in kwargs and 'metadata' in kwargs['serialized']:
+            tool_meta['metadata'] = kwargs['serialized']['metadata']
+            log.info(f"[METADATA] Extracted from serialized: {kwargs['serialized']['metadata']}")
+        # Source 2: Check if metadata is directly in args[0] (some LangChain versions)
+        elif 'metadata' in args[0]:
+            tool_meta['metadata'] = args[0]['metadata']
+            log.info(f"[METADATA] Extracted from args[0]: {args[0]['metadata']}")
+        else:
+            log.info(f"[METADATA] No metadata found. args[0] keys: {list(args[0].keys())}, kwargs keys: {list(kwargs.keys())}")
+            # Fallback: Try to extract toolkit_name from description
+            description = args[0].get('description', '')
+            if description:
+                import re
+                # Try pattern 1: [Toolkit: name]
+                match = re.search(r'\[Toolkit:\s*([^\]]+)\]', description)
+                if not match:
+                    # Try pattern 2: Toolkit: name at start or end
+                    match = re.search(r'(?:^|\n)Toolkit:\s*([^\n]+)', description)
+                if match:
+                    toolkit_name = match.group(1).strip()
+                    tool_meta['metadata'] = {'toolkit_name': toolkit_name}
+                    log.info(f"[METADATA] Extracted toolkit_name from description: {toolkit_name}")
+        
         payload = {
             "tool_name": tool_name,
             "tool_run_id": tool_run_id,
-            "tool_meta": args[0],
+            "tool_meta": tool_meta,
             "tool_inputs": kwargs.get('inputs')
         }
         payload = json.loads(json.dumps(payload, ensure_ascii=False, default=lambda o: str(o)))
-        self.callback_state[tool_run_id] = self.st.status(f"Running {tool_name}...", expanded=True)
-        self.callback_state[tool_run_id].write(f"Tool inputs: {kwargs.get('inputs')}")
+
+        status_widget = self._safe_streamlit_call(
+            self.st.status,
+            f"Running {tool_name}...",
+            expanded=True
+        )
+        if status_widget:
+            self.callback_state[tool_run_id] = status_widget
+            self._safe_streamlit_call(status_widget.write, f"Tool inputs: {kwargs.get('inputs')}")
 
     def on_tool_end(self, *args, run_id: UUID, **kwargs):
         """ Callback """
@@ -104,11 +180,16 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
             log.info("on_tool_end(%s, %s)", args, kwargs)
         tool_run_id = str(run_id)
         tool_output = args[0]
-        if self.callback_state[tool_run_id]:
-            self.callback_state[tool_run_id].write(f"Tool output: {tool_output}")
-            self.callback_state[tool_run_id].update(label=f"Completed {kwargs.get('name')}", state="complete", expanded=False)
+        if self.callback_state.get(tool_run_id):
+            status_widget = self.callback_state[tool_run_id]
+            self._safe_streamlit_call(status_widget.write, f"Tool output: {tool_output}")
+            self._safe_streamlit_call(
+                status_widget.update,
+                label=f"Completed {kwargs.get('name')}",
+                state="complete",
+                expanded=False
+            )
             self.callback_state.pop(tool_run_id, None)
-            del self.callback_state[run_id]
 
     def on_tool_error(self, *args, run_id: UUID, **kwargs):
         """ Callback """
@@ -116,9 +197,19 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
             log.info("on_tool_error(%s, %s)", args, kwargs)
         tool_run_id = str(run_id)
         tool_exception = args[0]
-        self.callback_state[tool_run_id].write(f"{traceback.format_exception(tool_exception)}")
-        self.callback_state[tool_run_id].update(label=f"Error {kwargs.get('name')}", state="error", expanded=False)
-        self.callback_state.pop(tool_run_id, None)
+        if self.callback_state.get(tool_run_id):
+            status_widget = self.callback_state[tool_run_id]
+            self._safe_streamlit_call(
+                status_widget.write,
+                f"{traceback.format_exception(tool_exception)}"
+            )
+            self._safe_streamlit_call(
+                status_widget.update,
+                label=f"Error {kwargs.get('name')}",
+                state="error",
+                expanded=False
+            )
+            self.callback_state.pop(tool_run_id, None)
 
     #
     # Agent
@@ -156,8 +247,14 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
         self.current_model_name = metadata.get('ls_model_name', self.current_model_name)
         llm_run_id = str(run_id)
 
-        self.callback_state[llm_run_id] = self.st.status(f"Running LLM ...", expanded=True)
-        self.callback_state[llm_run_id].write(f"LLM inputs: {messages}")
+        status_widget = self._safe_streamlit_call(
+            self.st.status,
+            f"Running LLM ...",
+            expanded=True
+        )
+        if status_widget:
+            self.callback_state[llm_run_id] = status_widget
+            self._safe_streamlit_call(status_widget.write, f"LLM inputs: {messages}")
 
     def on_llm_start(self, *args, **kwargs):
         """ Callback """
@@ -178,16 +275,27 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
         content = None
         if chunk:
             content = chunk.text
-        self.callback_state[str(run_id)].write(content)
+
+        llm_run_id = str(run_id)
+        if self.callback_state.get(llm_run_id):
+            status_widget = self.callback_state[llm_run_id]
+            self._safe_streamlit_call(status_widget.write, content)
 
     def on_llm_error(self, *args, run_id: UUID, **kwargs):
         """ Callback """
         if self.debug:
             log.error("on_llm_error(%s, %s)", args, kwargs)
         llm_run_id = str(run_id)
-        self.callback_state[llm_run_id].write(f"on_llm_error({args}, {kwargs})")
-        self.callback_state[llm_run_id].update(label=f"Error {kwargs.get('name')}", state="error", expanded=False)
-        self.callback_state.pop(llm_run_id, None)
+        if self.callback_state.get(llm_run_id):
+            status_widget = self.callback_state[llm_run_id]
+            self._safe_streamlit_call(status_widget.write, f"on_llm_error({args}, {kwargs})")
+            self._safe_streamlit_call(
+                status_widget.update,
+                label=f"Error {kwargs.get('name')}",
+                state="error",
+                expanded=False
+            )
+            self.callback_state.pop(llm_run_id, None)
         #
         # exception = args[0]
         # FIXME: should we emit an error here too?
@@ -205,5 +313,12 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
         if self.debug:
             log.debug("on_llm_end(%s, %s)", response, kwargs)
         llm_run_id = str(run_id)
-        self.callback_state[llm_run_id].update(label=f"Completed LLM call", state="complete", expanded=False)
-        self.callback_state.pop(llm_run_id, None)
+        if self.callback_state.get(llm_run_id):
+            status_widget = self.callback_state[llm_run_id]
+            self._safe_streamlit_call(
+                status_widget.update,
+                label=f"Completed LLM call",
+                state="complete",
+                expanded=False
+            )
+            self.callback_state.pop(llm_run_id, None)