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

alita_sdk/runtime/tools/vectorstore_base.py

@@ -155,15 +155,45 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
         if values.get('alita') and values.get('embedding_model'):
             values['embeddings'] = values.get('alita').get_embeddings(values.get('embedding_model'))
 
-        if values.get('vectorstore_type') and values.get('vectorstore_params') and values.get('embedding_model'):
-            values['vectorstore'] = get_vectorstore(values['vectorstore_type'], values['vectorstore_params'], embedding_func=values['embeddings'])
-            # Initialize the new vector adapter
-            values['vector_adapter'] = VectorStoreAdapterFactory.create_adapter(values['vectorstore_type'])
-            logger.debug(f"Vectorstore wrapper initialized: {values}")
+        # Lazy initialization: vectorstore and vector_adapter are initialized on-demand
+        # This prevents errors when using non-index tools with broken/missing vector DB
         return values
 
+    def _ensure_vectorstore_initialized(self):
+        """Lazily initialize vectorstore and vector_adapter when needed for index operations."""
+        if self.vectorstore is None:
+            if not self.vectorstore_type or not self.vectorstore_params:
+                raise ToolException(
+                    "Vector store is not configured. "
+                    "Please ensure embedding_model and pgvector_configuration are provided."
+                )
+
+            from ..langchain.interfaces.llm_processor import get_vectorstore
+            try:
+                self.vectorstore = get_vectorstore(
+                    self.vectorstore_type,
+                    self.vectorstore_params,
+                    embedding_func=self.embeddings
+                )
+                logger.debug(f"Vectorstore initialized: {self.vectorstore_type}")
+            except Exception as e:
+                raise ToolException(
+                    f"Failed to initialize vector store: {str(e)}. "
+                    "Check your vector database configuration and connection."
+                )
+
+        if self.vector_adapter is None:
+            try:
+                self.vector_adapter = VectorStoreAdapterFactory.create_adapter(self.vectorstore_type)
+                logger.debug(f"Vector adapter initialized: {self.vectorstore_type}")
+            except Exception as e:
+                raise ToolException(
+                    f"Failed to initialize vector adapter: {str(e)}"
+                )
+
     def _init_pg_helper(self, language='english'):
         """Initialize PGVector helper if needed and not already initialized"""
+        self._ensure_vectorstore_initialized()
         if self.pg_helper is None and hasattr(self.vectorstore, 'connection_string') and hasattr(self.vectorstore, 'collection_name'):
             try:
                 from .pgvector_search import PGVectorSearch
@@ -192,6 +222,7 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
         Raises:
             ToolException: When DataException occurs or other search errors
         """
+        self._ensure_vectorstore_initialized()
         try:
             return self.vectorstore.similarity_search_with_score(
                 query, filter=filter, k=k
@@ -210,19 +241,21 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
 
     def list_collections(self) -> List[str]:
         """List all collections in the vectorstore."""
-
+        self._ensure_vectorstore_initialized()
         collections = self.vector_adapter.list_collections(self)
         if not collections:
             return "No indexed collections"
         return collections
 
     def get_index_meta(self, index_name: str):
+        self._ensure_vectorstore_initialized()
         index_metas = self.vector_adapter.get_index_meta(self, index_name)
         if len(index_metas) > 1:
             raise RuntimeError(f"Multiple index_meta documents found: {index_metas}")
         return index_metas[0] if index_metas else None
 
     def get_indexed_count(self, index_name: str) -> int:
+        self._ensure_vectorstore_initialized()
         from sqlalchemy.orm import Session
         from sqlalchemy import func, or_
 
@@ -237,15 +270,16 @@ 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.
         """
+        self._ensure_vectorstore_initialized()
         self._log_tool_event(
             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"
@@ -259,6 +293,7 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
             progress_step (int): Step for progress reporting, default is 20.
             clean_index (bool): If True, clean the index before re-indexing all documents.
         """
+        self._ensure_vectorstore_initialized()
         if clean_index:
             self._clean_index(index_name)
 
@@ -268,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)