alita-sdk 0.3.379__py3-none-any.whl → 0.3.627__py3-none-any.whl

alita_sdk/runtime/tools/tool.py

@@ -86,7 +86,9 @@ Answer must be JSON only extractable by JSON.LOADS."""
         else:
             input_[-1].content += self.unstructured_output
             completion = self.client.invoke(input_, config=config)
-            result = _extract_json(completion.content.strip())
+            from ..langchain.utils import extract_text_from_completion
+            content_text = extract_text_from_completion(completion)
+            result = _extract_json(content_text.strip())
             logger.info(f"ToolNode tool params: {result}")
         try:
             # handler for application added as a tool

alita_sdk/runtime/tools/vectorstore.py

@@ -12,9 +12,11 @@ from alita_sdk.tools.vector_adapters.VectorStoreAdapter import VectorStoreAdapte
 from logging import getLogger
 
 from ..utils.logging import dispatch_custom_event
+from ..langchain.utils import extract_text_from_completion
 
 logger = getLogger(__name__)
 
+
 class IndexDocumentsModel(BaseModel):
     documents: Any = Field(description="Generator of documents to index")
 
@@ -414,7 +416,8 @@ class VectorStoreWrapper(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 new documents to index."}
 
     def search_documents(self, query:str, doctype: str = 'code', 
                          filter:dict|str={}, cut_off: float=0.5,
@@ -683,8 +686,10 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 ]
             )
         ])
+        # Extract text content safely (handles both string and list content from thinking models)
+        search_query = extract_text_from_completion(result)
         search_results = self.search_documents(
-            result.content, doctype, filter, cut_off, search_top, 
+            search_query, doctype, filter, cut_off, search_top, 
             full_text_search=full_text_search,
             reranking_config=reranking_config,
             extended_search=extended_search
@@ -713,7 +718,8 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                 ]
             )
         ])
-        return result.content
+        # Extract text content safely (handles both string and list content from thinking models)
+        return extract_text_from_completion(result)
 
     def _log_data(self, message: str, tool_name: str = "index_data"):
         """Log data and dispatch custom event for indexing progress"""

alita_sdk/runtime/tools/vectorstore_base.py

@@ -1,9 +1,9 @@
 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
@@ -12,10 +12,12 @@ 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 ..utils.logging import dispatch_custom_event
+from ...runtime.utils.utils import IndexerKeywords
+from ...runtime.langchain.utils import extract_text_from_completion
 
 logger = getLogger(__name__)
 
+
 class IndexDocumentsModel(BaseModel):
     documents: Any = Field(description="Generator of documents to index")
 
@@ -155,15 +157,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 +224,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,27 +243,45 @@ 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 _clean_collection(self, index_name: str = ''):
+    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()
+
+    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"
@@ -244,6 +295,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)
 
@@ -253,7 +305,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:
@@ -308,7 +360,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,
@@ -574,8 +627,10 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                 ]
             )
         ])
+        # Extract text content safely (handles both string and list content from thinking models)
+        search_query = extract_text_from_completion(result)
         search_results = self.search_documents(
-            result.content, doctype, filter, cut_off, search_top, 
+            search_query, doctype, filter, cut_off, search_top, 
             full_text_search=full_text_search,
             reranking_config=reranking_config,
             extended_search=extended_search
@@ -604,7 +659,8 @@ class VectorStoreWrapperBase(BaseToolApiWrapper):
                 ]
             )
         ])
-        return result.content
+        # Extract text content safely (handles both string and list content from thinking models)
+        return extract_text_from_completion(result)
 
     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 """
@@ -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,13 @@ 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)
+        # Check if callback_state exists and is not None before accessing
+        if self.callback_state is not None and 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)

alita_sdk/runtime/utils/constants.py

@@ -436,4 +436,8 @@ STYLES = r"""
             filter: hue-rotate(180deg) brightness(1.2);
         }
         </style>
-        """
+        """
+
+TOOLKIT_NAME_META = "toolkit_name"
+TOOL_NAME_META = "tool_name"
+TOOLKIT_TYPE_META = "toolkit_type"