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

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__()
@@ -98,7 +112,7 @@ class McpRemoteTool(McpServerTool):
             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}")
         
-        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 +120,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 +160,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/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 """
@@ -95,8 +135,15 @@ class AlitaStreamlitCallback(BaseCallbackHandler):
             "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 +151,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 +168,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 +218,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 +246,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 +284,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)