alita-sdk 0.3.263__py3-none-any.whl → 0.3.499__py3-none-any.whl

alita_sdk/runtime/tools/vectorstore_base.py

@@ -1,16 +1,18 @@
 import json
-import math
+from collections import OrderedDict
 from logging import getLogger
 from typing import Any, Optional, List, Dict, Generator
 
+import math
 from langchain_core.documents import Document
 from langchain_core.messages import HumanMessage
+from langchain_core.tools import ToolException
+from psycopg.errors import DataException
 from pydantic import BaseModel, model_validator, Field
 
 from alita_sdk.tools.elitea_base import BaseToolApiWrapper
 from alita_sdk.tools.vector_adapters.VectorStoreAdapter import VectorStoreAdapterFactory
-from ..langchain.tools.vector import VectorAdapter
-from ..utils.logging import dispatch_custom_event
+from ...runtime.utils.utils import IndexerKeywords
 
 logger = getLogger(__name__)
 
@@ -132,15 +134,12 @@ How did you come up with the answer?
 
 class VectorStoreWrapperBase(BaseToolApiWrapper):
     llm: Any
-    embedding_model: str
-    vectorstore_type: str
-    vectorstore_params: dict
-    max_docs_per_add: int = 100
-    dataset: str = None
-    embedding: Any = None
+    embedding_model: Optional[str] = None
+    vectorstore_type: Optional[str]  = None
+    vectorstore_params: Optional[dict]  = None
+    max_docs_per_add: int = 20
+    dataset: Optional[str] = None
     vectorstore: Any = None
-    # Review usage of old adapter
-    vectoradapter: Any = None
     pg_helper: Any = None
     embeddings: Any = None
     # New adapter for vector database operations
@@ -149,32 +148,52 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
     @model_validator(mode='before')
     @classmethod
     def validate_toolkit(cls, values):
-        from ..langchain.interfaces.llm_processor import get_embeddings, get_vectorstore
+        from ..langchain.interfaces.llm_processor import get_vectorstore
         logger.debug(f"Validating toolkit: {values}")
-        if not values.get('vectorstore_type'):
-            raise ValueError("Vectorstore type is required.")
-        if not values.get('embedding_model'):
-            raise ValueError("Embedding model is required.")
-        if not values.get('vectorstore_params'):
-            raise ValueError("Vectorstore parameters are required.")
-        values["dataset"] = values.get('vectorstore_params').get('collection_name')
-        if not values["dataset"]:
-            raise ValueError("Collection name is required.")
-        if not values.get('embeddings'):
-            values['embeddings'] = values['alita'].get_embeddings(values['embedding_model'])
-        values['vectorstore'] = get_vectorstore(values['vectorstore_type'], values['vectorstore_params'], embedding_func=values['embeddings'])
-        values['vectoradapter'] = VectorAdapter(
-            vectorstore=values['vectorstore'],
-            embeddings=values['embeddings'],
-            quota_params=None,
-        )
-        # Initialize the new vector adapter
-        values['vector_adapter'] = VectorStoreAdapterFactory.create_adapter(values['vectorstore_type'])
-        logger.debug(f"Vectorstore wrapper initialized: {values}")
+        values["dataset"] = values.get('collection_name')
+
+        if values.get('alita') and values.get('embedding_model'):
+            values['embeddings'] = values.get('alita').get_embeddings(values.get('embedding_model'))
+
+        # 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
@@ -188,26 +207,85 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
             except Exception as e:
                 logger.error(f"Failed to initialize PGVectorSearch: {str(e)}")
 
+    def _similarity_search_with_score(self, query: str, filter: dict = None, k: int = 10):
+        """
+        Perform similarity search with proper exception handling for DataException.
+
+        Args:
+            query: Search query string
+            filter: Optional filter dictionary
+            k: Number of results to return
+
+        Returns:
+            List of (Document, score) tuples
+
+        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
+            )
+        except DataException as dimException:
+            exception_str = str(dimException)
+            if 'different vector dimensions' in exception_str:
+                logger.error(f"Data exception: {exception_str}")
+                raise ToolException(f"Global search cannot be completed since collections were indexed using "
+                                    f"different embedding models. Use search within a single collection."
+                                    f"\nDetails: {exception_str}")
+            raise ToolException(f"Data exception during search. Possibly invalid filter: {exception_str}")
+        except Exception as e:
+            logger.error(f"Error during similarity search: {str(e)}")
+            raise ToolException(f"Search failed: {str(e)}")
+
     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_
+
+        with Session(self.vectorstore.session_maker.bind) as session:
+            return session.query(
+                self.vectorstore.EmbeddingStore.id,
+            ).filter(
+                func.jsonb_extract_path_text(self.vectorstore.EmbeddingStore.cmetadata, 'collection') == index_name,
+                or_(
+                    func.jsonb_extract_path_text(self.vectorstore.EmbeddingStore.cmetadata, 'type').is_(None),
+                    func.jsonb_extract_path_text(self.vectorstore.EmbeddingStore.cmetadata, 'type') != IndexerKeywords.INDEX_META_TYPE.value
+                )
+            ).count()
 
-        return self.vector_adapter.list_collections(self)
-
-    def _clean_collection(self, collection_suffix: str = ''):
+    def _clean_collection(self, index_name: str = '', including_index_meta: bool = False):
         """
         Clean the vectorstore collection by deleting all indexed data.
         """
-        self._log_data(
+        self._ensure_vectorstore_initialized()
+        self._log_tool_event(
             f"Cleaning collection '{self.dataset}'",
             tool_name="_clean_collection"
         )
-        self.vector_adapter.clean_collection(self, collection_suffix)
-        self._log_data(
+        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"
         )
 
-    def index_documents(self, documents: Generator[Document, None, None], collection_suffix: str, progress_step: int = 20, clean_index: bool = True):
+    def index_documents(self, documents: Generator[Document, None, None], index_name: str, progress_step: int = 20, clean_index: bool = True):
         """ Index documents in the vectorstore.
 
         Args:
@@ -215,24 +293,23 @@ 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(collection_suffix)
+            self._clean_index(index_name)
 
-        return self._save_index(list(documents), collection_suffix, progress_step)
+        return self._save_index(list(documents), index_name, progress_step)
 
-    def _clean_index(self, collection_suffix: str):
+    def _clean_index(self, index_name: str):
         logger.info("Cleaning index before re-indexing all documents.")
-        self._log_data("Cleaning index before re-indexing all documents. Previous index will be removed", tool_name="index_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(collection_suffix)
-            self.vectoradapter.persist()
-            self.vectoradapter.vacuum()
-            self._log_data("Previous index has been removed",
+            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:
             logger.warning(f"Failed to clean index: {str(e)}. Continuing with re-indexing.")
 
-    def _save_index(self, documents: list[Document], collection_suffix: Optional[str] = None, progress_step: int = 20):
+    def _save_index(self, documents: list[Document], index_name: Optional[str] = None, progress_step: int = 20):
         from ..langchain.interfaces.llm_processor import add_documents
         #
         for doc in documents:
@@ -240,15 +317,14 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                 logger.warning(f"Document is missing required metadata field 'id' or 'updated_on': {doc.metadata}")
 
         logger.debug(f"Indexing documents: {documents}")
-        logger.debug(self.vectoradapter)
 
-        # if collection_suffix is provided, add it to metadata of each document
-        if collection_suffix:
+        # if index_name is provided, add it to metadata of each document
+        if index_name:
             for doc in documents:
                 if not doc.metadata.get('collection'):
-                    doc.metadata['collection'] = collection_suffix
+                    doc.metadata['collection'] = index_name
                 else:
-                    doc.metadata['collection'] += f";{collection_suffix}"
+                    doc.metadata['collection'] += f";{index_name}"
 
         total_docs = len(documents)
         documents_count = 0
@@ -258,6 +334,10 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
         progress_step = 20 if progress_step not in range(0, 100) else progress_step
         next_progress_point = progress_step
         for document in documents:
+            if not document.page_content:
+                # To avoid case when all documents have empty content
+                # See llm_processor.add_documents which exclude metadata of docs with empty content
+                continue
             documents_count += 1
             # logger.debug(f"Indexing document: {document}")
             try:
@@ -270,7 +350,7 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                 if percent >= next_progress_point:
                     msg = f"Indexing progress: {percent}%. Processed {documents_count} of {total_docs} documents."
                     logger.debug(msg)
-                    self._log_data(msg)
+                    self._log_tool_event(msg)
                     next_progress_point += progress_step
             except Exception:
                 from traceback import format_exc
@@ -278,7 +358,8 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                 return {"status": "error", "message": f"Error: {format_exc()}"}
         if _documents:
             add_documents(vectorstore=self.vectorstore, documents=_documents)
-        return {"status": "ok", "message": f"successfully indexed {documents_count} documents"}
+        return {"status": "ok", "message": f"successfully indexed {documents_count} documents" if documents_count > 0
+        else "no documents to index"}
 
     def search_documents(self, query:str, doctype: str = 'code', 
                          filter:dict|str={}, cut_off: float=0.5,
@@ -312,7 +393,7 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                 }
                 
             try:
-                document_items = self.vectorstore.similarity_search_with_score(
+                document_items = self._similarity_search_with_score(
                     query, filter=document_filter, k=search_top
                 )                
                 # Add document results to unique docs
@@ -345,18 +426,16 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                     }
                 
                 try:
-                    chunk_items = self.vectorstore.similarity_search_with_score(
+                    chunk_items = self._similarity_search_with_score(
                         query, filter=chunk_filter, k=search_top
                     )
-                    
-                    logger.debug(f"Chunk items for {chunk_type}: {chunk_items[0]}")
-                    
+
                     for doc, score in chunk_items:
                         # Create unique identifier for document
                         source = doc.metadata.get('source')
                         chunk_id = doc.metadata.get('chunk_id')
                         doc_id = f"{source}_{chunk_id}" if source and chunk_id else str(doc.metadata.get('id', id(doc)))
-                        
+
                         # Store document and its score
                         if doc_id not in unique_docs:
                             unique_docs[doc_id] = doc
@@ -376,9 +455,9 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                                 doc_filter = {
                                     "$and": doc_filter_parts
                                 }
-                                
+
                             try:
-                                fetch_items = self.vectorstore.similarity_search_with_score(
+                                fetch_items = self._similarity_search_with_score(
                                     query, filter=doc_filter, k=1
                                 )
                                 if fetch_items:
@@ -392,18 +471,25 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
         else:
             # Default search behavior (unchanged)
             max_search_results = 30 if search_top * 3 > 30 else search_top * 3
-            vector_items = self.vectorstore.similarity_search_with_score(
+            vector_items = self._similarity_search_with_score(
                 query, filter=filter, k=max_search_results
             )
             
         # Initialize document map for tracking by ID
         doc_map = {
-            f"{doc.metadata.get('id', f'idx_{i}')}_{doc.metadata['chunk_id']}"
-            if 'chunk_id' in doc.metadata
-            else doc.metadata.get('id', f"idx_{i}"): (doc, score)
+            (
+                f"{doc.metadata.get('id', f'idx_{i}')}_{doc.metadata['chunk_id']}"
+                if 'chunk_id' in doc.metadata
+                else doc.metadata.get('id', f"idx_{i}")
+            ): (doc, 1 - score)
             for i, (doc, score) in enumerate(vector_items)
         }
-        
+
+        # Sort the items by the new score in descending order
+        doc_map = OrderedDict(
+            sorted(doc_map.items(), key=lambda x: x[1][1], reverse=True)
+        )
+
         # Process full-text search if configured
         if full_text_search and full_text_search.get('enabled') and full_text_search.get('fields'):
             language = full_text_search.get('language', 'english')
@@ -416,7 +502,7 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                 for field_name in full_text_search.get('fields', []):
                     try:
                         text_results = self.pg_helper.full_text_search(field_name, query)
-                        
+
                         # Combine text search results with vector results
                         for result in text_results:
                             doc_id = result['id']
@@ -452,7 +538,7 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
         
         # Apply cutoff filter
         if cut_off:
-            combined_items = [item for item in combined_items if abs(item[1]) <= cut_off]
+            combined_items = [item for item in combined_items if abs(item[1]) >= cut_off]
         
         # Sort by score and limit results
         # DISABLED: for chroma we want ascending order (lower score is better), for others descending
@@ -571,21 +657,6 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
         ])
         return result.content
 
-    def _log_data(self, message: str, tool_name: str = "index_data"):
-        """Log data and dispatch custom event for indexing progress"""
-
-        try:
-            dispatch_custom_event(
-                name="thinking_step",
-                data={
-                    "message": message,
-                    "tool_name": tool_name,
-                    "toolkit": "vectorstore",
-                },
-            )
-        except Exception as e:
-            logger.warning(f"Failed to dispatch progress event: {str(e)}")
-
     def get_available_tools(self):
         return [
             {

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)