agno 2.3.8__py3-none-any.whl → 2.3.10__py3-none-any.whl

agno/vectordb/pgvector/pgvector.py

@@ -28,7 +28,7 @@ from agno.filters import FilterExpr
 from agno.knowledge.document import Document
 from agno.knowledge.embedder import Embedder
 from agno.knowledge.reranker.base import Reranker
-from agno.utils.log import log_debug, log_info, logger
+from agno.utils.log import log_debug, log_error, log_info, log_warning
 from agno.vectordb.base import VectorDb
 from agno.vectordb.distance import Distance
 from agno.vectordb.pgvector.index import HNSW, Ivfflat
@@ -62,6 +62,7 @@ class PgVector(VectorDb):
         schema_version: int = 1,
         auto_upgrade_schema: bool = False,
         reranker: Optional[Reranker] = None,
+        create_schema: bool = True,
     ):
         """
         Initialize the PgVector instance.
@@ -82,6 +83,8 @@ class PgVector(VectorDb):
             content_language (str): Language for full-text search.
             schema_version (int): Version of the database schema.
             auto_upgrade_schema (bool): Automatically upgrade schema if True.
+            create_schema (bool): Whether to automatically create the database schema if it doesn't exist.
+                Set to False if schema is managed externally (e.g., via migrations). Defaults to True.
         """
         if not table_name:
             raise ValueError("Table name must be provided.")
@@ -104,7 +107,7 @@ class PgVector(VectorDb):
             try:
                 db_engine = create_engine(db_url)
             except Exception as e:
-                logger.error(f"Failed to create engine from 'db_url': {e}")
+                log_error(f"Failed to create engine from 'db_url': {e}")
                 raise
 
         # Database settings
@@ -147,6 +150,9 @@ class PgVector(VectorDb):
         # Reranker instance
         self.reranker: Optional[Reranker] = reranker
 
+        # Schema creation flag
+        self.create_schema: bool = create_schema
+
         # Database session
         self.Session: scoped_session = scoped_session(sessionmaker(bind=self.db_engine))
         # Database table
@@ -209,7 +215,7 @@ class PgVector(VectorDb):
         try:
             return inspect(self.db_engine).has_table(self.table_name, schema=self.schema)
         except Exception as e:
-            logger.error(f"Error checking if table exists: {e}")
+            log_error(f"Error checking if table exists: {e}")
             return False
 
     def create(self) -> None:
@@ -220,7 +226,7 @@ class PgVector(VectorDb):
             with self.Session() as sess, sess.begin():
                 log_debug("Creating extension: vector")
                 sess.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
-                if self.schema is not None:
+                if self.create_schema and self.schema is not None:
                     log_debug(f"Creating schema: {self.schema}")
                     sess.execute(text(f"CREATE SCHEMA IF NOT EXISTS {self.schema};"))
             log_debug(f"Creating table: {self.table_name}")
@@ -247,7 +253,7 @@ class PgVector(VectorDb):
                 result = sess.execute(stmt).first()
                 return result is not None
         except Exception as e:
-            logger.error(f"Error checking if record exists: {e}")
+            log_error(f"Error checking if record exists: {e}")
             return False
 
     def name_exists(self, name: str) -> bool:
@@ -324,7 +330,7 @@ class PgVector(VectorDb):
                             try:
                                 batch_records.append(self._get_document_record(doc, filters, content_hash))
                             except Exception as e:
-                                logger.error(f"Error processing document '{doc.name}': {e}")
+                                log_error(f"Error processing document '{doc.name}': {e}")
 
                         # Insert the batch of records
                         insert_stmt = postgresql.insert(self.table)
@@ -332,11 +338,11 @@ class PgVector(VectorDb):
                         sess.commit()  # Commit batch independently
                         log_info(f"Inserted batch of {len(batch_records)} documents.")
                     except Exception as e:
-                        logger.error(f"Error with batch starting at index {i}: {e}")
+                        log_error(f"Error with batch starting at index {i}: {e}")
                         sess.rollback()  # Rollback the current batch if there's an error
                         raise
         except Exception as e:
-            logger.error(f"Error inserting documents: {e}")
+            log_error(f"Error inserting documents: {e}")
             raise
 
     async def async_insert(
@@ -380,7 +386,7 @@ class PgVector(VectorDb):
                                 }
                                 batch_records.append(record)
                             except Exception as e:
-                                logger.error(f"Error processing document '{doc.name}': {e}")
+                                log_error(f"Error processing document '{doc.name}': {e}")
 
                         # Insert the batch of records
                         if batch_records:
@@ -389,11 +395,11 @@ class PgVector(VectorDb):
                             sess.commit()  # Commit batch independently
                             log_info(f"Inserted batch of {len(batch_records)} documents.")
                     except Exception as e:
-                        logger.error(f"Error with batch starting at index {i}: {e}")
+                        log_error(f"Error with batch starting at index {i}: {e}")
                         sess.rollback()  # Rollback the current batch if there's an error
                         raise
         except Exception as e:
-            logger.error(f"Error inserting documents: {e}")
+            log_error(f"Error inserting documents: {e}")
             raise
 
     def upsert_available(self) -> bool:
@@ -422,7 +428,7 @@ class PgVector(VectorDb):
                 self._delete_by_content_hash(content_hash)
             self._upsert(content_hash, documents, filters, batch_size)
         except Exception as e:
-            logger.error(f"Error upserting documents by content hash: {e}")
+            log_error(f"Error upserting documents by content hash: {e}")
             raise
 
     def _upsert(
@@ -452,7 +458,7 @@ class PgVector(VectorDb):
                             try:
                                 batch_records_dict[doc.id] = self._get_document_record(doc, filters, content_hash)  # type: ignore
                             except Exception as e:
-                                logger.error(f"Error processing document '{doc.name}': {e}")
+                                log_error(f"Error processing document '{doc.name}': {e}")
 
                         # Convert dict to list for upsert
                         batch_records = list(batch_records_dict.values())
@@ -479,11 +485,11 @@ class PgVector(VectorDb):
                         sess.commit()  # Commit batch independently
                         log_info(f"Upserted batch of {len(batch_records)} documents.")
                     except Exception as e:
-                        logger.error(f"Error with batch starting at index {i}: {e}")
+                        log_error(f"Error with batch starting at index {i}: {e}")
                         sess.rollback()  # Rollback the current batch if there's an error
                         raise
         except Exception as e:
-            logger.error(f"Error upserting documents: {e}")
+            log_error(f"Error upserting documents: {e}")
             raise
 
     def _get_document_record(
@@ -532,7 +538,7 @@ class PgVector(VectorDb):
                             doc.embedding = embeddings[j]
                             doc.usage = usages[j] if j < len(usages) else None
                     except Exception as e:
-                        logger.error(f"Error assigning batch embedding to document '{doc.name}': {e}")
+                        log_error(f"Error assigning batch embedding to document '{doc.name}': {e}")
 
             except Exception as e:
                 # Check if this is a rate limit error - don't fall back as it would make things worse
@@ -543,17 +549,41 @@ class PgVector(VectorDb):
                 )
 
                 if is_rate_limit:
-                    logger.error(f"Rate limit detected during batch embedding.  {e}")
+                    log_error(f"Rate limit detected during batch embedding.  {e}")
                     raise e
                 else:
-                    logger.warning(f"Async batch embedding failed, falling back to individual embeddings: {e}")
+                    log_warning(f"Async batch embedding failed, falling back to individual embeddings: {e}")
                     # Fall back to individual embedding
                     embed_tasks = [doc.async_embed(embedder=self.embedder) for doc in batch_docs]
-                    await asyncio.gather(*embed_tasks, return_exceptions=True)
+                    results = await asyncio.gather(*embed_tasks, return_exceptions=True)
+
+                    # Check for exceptions and handle them
+                    for i, result in enumerate(results):
+                        if isinstance(result, Exception):
+                            error_msg = str(result)
+                            # If it's an event loop closure error, log it but don't fail
+                            if "Event loop is closed" in error_msg or "RuntimeError" in type(result).__name__:
+                                log_warning(
+                                    f"Event loop closure during embedding for document {i}, but operation may have succeeded: {result}"
+                                )
+                            else:
+                                log_error(f"Error embedding document {i}: {result}")
         else:
             # Use individual embedding
             embed_tasks = [doc.async_embed(embedder=self.embedder) for doc in batch_docs]
-            await asyncio.gather(*embed_tasks, return_exceptions=True)
+            results = await asyncio.gather(*embed_tasks, return_exceptions=True)
+
+            # Check for exceptions and handle them
+            for i, result in enumerate(results):
+                if isinstance(result, Exception):
+                    error_msg = str(result)
+                    # If it's an event loop closure error, log it but don't fail
+                    if "Event loop is closed" in error_msg or "RuntimeError" in type(result).__name__:
+                        log_warning(
+                            f"Event loop closure during embedding for document {i}, but operation may have succeeded: {result}"
+                        )
+                    else:
+                        log_error(f"Error embedding document {i}: {result}")
 
     async def async_upsert(
         self,
@@ -568,7 +598,7 @@ class PgVector(VectorDb):
                 self._delete_by_content_hash(content_hash)
             await self._async_upsert(content_hash, documents, filters, batch_size)
         except Exception as e:
-            logger.error(f"Error upserting documents by content hash: {e}")
+            log_error(f"Error upserting documents by content hash: {e}")
             raise
 
     async def _async_upsert(
@@ -597,11 +627,18 @@ class PgVector(VectorDb):
 
                         # Prepare documents for upserting
                         batch_records_dict = {}  # Use dict to deduplicate by ID
-                        for doc in batch_docs:
+                        for idx, doc in enumerate(batch_docs):
                             try:
                                 cleaned_content = self._clean_content(doc.content)
                                 record_id = md5(cleaned_content.encode()).hexdigest()
 
+                                if (
+                                    doc.embedding is not None
+                                    and isinstance(doc.embedding, list)
+                                    and len(doc.embedding) == 0
+                                ):
+                                    log_warning(f"Document {idx} '{doc.name}' has empty embedding (length 0)")
+
                                 meta_data = doc.meta_data or {}
                                 if filters:
                                     meta_data.update(filters)
@@ -619,7 +656,7 @@ class PgVector(VectorDb):
                                 }
                                 batch_records_dict[record_id] = record  # This deduplicates by ID
                             except Exception as e:
-                                logger.error(f"Error processing document '{doc.name}': {e}")
+                                log_error(f"Error processing document '{doc.name}': {e}")
 
                         # Convert dict to list for upsert
                         batch_records = list(batch_records_dict.values())
@@ -646,11 +683,11 @@ class PgVector(VectorDb):
                         sess.commit()  # Commit batch independently
                         log_info(f"Upserted batch of {len(batch_records)} documents.")
                     except Exception as e:
-                        logger.error(f"Error with batch starting at index {i}: {e}")
+                        log_error(f"Error with batch starting at index {i}: {e}")
                         sess.rollback()  # Rollback the current batch if there's an error
                         raise
         except Exception as e:
-            logger.error(f"Error upserting documents: {e}")
+            log_error(f"Error upserting documents: {e}")
             raise
 
     def update_metadata(self, content_id: str, metadata: Dict[str, Any]) -> None:
@@ -679,7 +716,7 @@ class PgVector(VectorDb):
                 sess.execute(stmt)
                 sess.commit()
         except Exception as e:
-            logger.error(f"Error updating metadata for document {content_id}: {e}")
+            log_error(f"Error updating metadata for document {content_id}: {e}")
             raise
 
     def search(
@@ -703,7 +740,7 @@ class PgVector(VectorDb):
         elif self.search_type == SearchType.hybrid:
             return self.hybrid_search(query=query, limit=limit, filters=filters)
         else:
-            logger.error(f"Invalid search type '{self.search_type}'.")
+            log_error(f"Invalid search type '{self.search_type}'.")
             return []
 
     async def async_search(
@@ -751,7 +788,7 @@ class PgVector(VectorDb):
             # Get the embedding for the query string
             query_embedding = self.embedder.get_embedding(query)
             if query_embedding is None:
-                logger.error(f"Error getting embedding for Query: {query}")
+                log_error(f"Error getting embedding for Query: {query}")
                 return []
 
             # Define the columns to select
@@ -789,7 +826,7 @@ class PgVector(VectorDb):
             elif self.distance == Distance.max_inner_product:
                 stmt = stmt.order_by(self.table.c.embedding.max_inner_product(query_embedding))
             else:
-                logger.error(f"Unknown distance metric: {self.distance}")
+                log_error(f"Unknown distance metric: {self.distance}")
                 return []
 
             # Limit the number of results
@@ -808,8 +845,8 @@ class PgVector(VectorDb):
                             sess.execute(text(f"SET LOCAL hnsw.ef_search = {self.vector_index.ef_search}"))
                     results = sess.execute(stmt).fetchall()
             except Exception as e:
-                logger.error(f"Error performing semantic search: {e}")
-                logger.error("Table might not exist, creating for future use")
+                log_error(f"Error performing semantic search: {e}")
+                log_error("Table might not exist, creating for future use")
                 self.create()
                 return []
 
@@ -834,7 +871,7 @@ class PgVector(VectorDb):
             log_info(f"Found {len(search_results)} documents")
             return search_results
         except Exception as e:
-            logger.error(f"Error during vector search: {e}")
+            log_error(f"Error during vector search: {e}")
             return []
 
     def enable_prefix_matching(self, query: str) -> str:
@@ -916,8 +953,8 @@ class PgVector(VectorDb):
                 with self.Session() as sess, sess.begin():
                     results = sess.execute(stmt).fetchall()
             except Exception as e:
-                logger.error(f"Error performing keyword search: {e}")
-                logger.error("Table might not exist, creating for future use")
+                log_error(f"Error performing keyword search: {e}")
+                log_error("Table might not exist, creating for future use")
                 self.create()
                 return []
 
@@ -939,7 +976,7 @@ class PgVector(VectorDb):
             log_info(f"Found {len(search_results)} documents")
             return search_results
         except Exception as e:
-            logger.error(f"Error during keyword search: {e}")
+            log_error(f"Error during keyword search: {e}")
             return []
 
     def hybrid_search(
@@ -963,7 +1000,7 @@ class PgVector(VectorDb):
             # Get the embedding for the query string
             query_embedding = self.embedder.get_embedding(query)
             if query_embedding is None:
-                logger.error(f"Error getting embedding for Query: {query}")
+                log_error(f"Error getting embedding for Query: {query}")
                 return []
 
             # Define the columns to select
@@ -1001,7 +1038,7 @@ class PgVector(VectorDb):
                 # Normalize to range [0, 1]
                 vector_score = (raw_vector_score + 1) / 2
             else:
-                logger.error(f"Unknown distance metric: {self.distance}")
+                log_error(f"Unknown distance metric: {self.distance}")
                 return []
 
             # Apply weights to control the influence of each score
@@ -1052,7 +1089,7 @@ class PgVector(VectorDb):
                             sess.execute(text(f"SET LOCAL hnsw.ef_search = {self.vector_index.ef_search}"))
                     results = sess.execute(stmt).fetchall()
             except Exception as e:
-                logger.error(f"Error performing hybrid search: {e}")
+                log_error(f"Error performing hybrid search: {e}")
                 return []
 
             # Process the results and convert to Document objects
@@ -1076,7 +1113,7 @@ class PgVector(VectorDb):
             log_info(f"Found {len(search_results)} documents")
             return search_results
         except Exception as e:
-            logger.error(f"Error during hybrid search: {e}")
+            log_error(f"Error during hybrid search: {e}")
             return []
 
     def drop(self) -> None:
@@ -1089,7 +1126,7 @@ class PgVector(VectorDb):
                 self.table.drop(self.db_engine)
                 log_info(f"Table '{self.table.fullname}' dropped successfully.")
             except Exception as e:
-                logger.error(f"Error dropping table '{self.table.fullname}': {e}")
+                log_error(f"Error dropping table '{self.table.fullname}': {e}")
                 raise
         else:
             log_info(f"Table '{self.table.fullname}' does not exist.")
@@ -1124,7 +1161,7 @@ class PgVector(VectorDb):
                 result = sess.execute(stmt).scalar()
                 return int(result) if result is not None else 0
         except Exception as e:
-            logger.error(f"Error getting count from table '{self.table.fullname}': {e}")
+            log_error(f"Error getting count from table '{self.table.fullname}': {e}")
             return 0
 
     def optimize(self, force_recreate: bool = False) -> None:
@@ -1165,7 +1202,7 @@ class PgVector(VectorDb):
                 drop_index_sql = f'DROP INDEX IF EXISTS "{self.schema}"."{index_name}";'
                 sess.execute(text(drop_index_sql))
         except Exception as e:
-            logger.error(f"Error dropping index '{index_name}': {e}")
+            log_error(f"Error dropping index '{index_name}': {e}")
             raise
 
     def _create_vector_index(self, force_recreate: bool = False) -> None:
@@ -1220,10 +1257,10 @@ class PgVector(VectorDb):
                 elif isinstance(self.vector_index, HNSW):
                     self._create_hnsw_index(sess, table_fullname, index_distance)
                 else:
-                    logger.error(f"Unknown index type: {type(self.vector_index)}")
+                    log_error(f"Unknown index type: {type(self.vector_index)}")
                     return
         except Exception as e:
-            logger.error(f"Error creating vector index '{self.vector_index.name}': {e}")
+            log_error(f"Error creating vector index '{self.vector_index.name}': {e}")
             raise
 
     def _create_ivfflat_index(self, sess: Session, table_fullname: str, index_distance: str) -> None:
@@ -1322,7 +1359,7 @@ class PgVector(VectorDb):
                 )
                 sess.execute(create_gin_index_sql)
         except Exception as e:
-            logger.error(f"Error creating GIN index '{gin_index_name}': {e}")
+            log_error(f"Error creating GIN index '{gin_index_name}': {e}")
             raise
 
     def delete(self) -> bool:
@@ -1341,7 +1378,7 @@ class PgVector(VectorDb):
                 log_info(f"Deleted all records from table '{self.table.fullname}'.")
                 return True
         except Exception as e:
-            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            log_error(f"Error deleting rows from table '{self.table.fullname}': {e}")
             sess.rollback()
             return False
 
@@ -1357,7 +1394,7 @@ class PgVector(VectorDb):
                 log_info(f"Deleted records with id '{id}' from table '{self.table.fullname}'.")
                 return True
         except Exception as e:
-            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            log_error(f"Error deleting rows from table '{self.table.fullname}': {e}")
             sess.rollback()
             return False
 
@@ -1373,7 +1410,7 @@ class PgVector(VectorDb):
                 log_info(f"Deleted records with name '{name}' from table '{self.table.fullname}'.")
                 return True
         except Exception as e:
-            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            log_error(f"Error deleting rows from table '{self.table.fullname}': {e}")
             sess.rollback()
             return False
 
@@ -1389,7 +1426,7 @@ class PgVector(VectorDb):
                 log_info(f"Deleted records with metadata '{metadata}' from table '{self.table.fullname}'.")
                 return True
         except Exception as e:
-            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            log_error(f"Error deleting rows from table '{self.table.fullname}': {e}")
             sess.rollback()
             return False
 
@@ -1405,7 +1442,7 @@ class PgVector(VectorDb):
                 log_info(f"Deleted records with content ID '{content_id}' from table '{self.table.fullname}'.")
                 return True
         except Exception as e:
-            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            log_error(f"Error deleting rows from table '{self.table.fullname}': {e}")
             sess.rollback()
             return False
 
@@ -1421,7 +1458,7 @@ class PgVector(VectorDb):
                 log_info(f"Deleted records with content hash '{content_hash}' from table '{self.table.fullname}'.")
                 return True
         except Exception as e:
-            logger.error(f"Error deleting rows from table '{self.table.fullname}': {e}")
+            log_error(f"Error deleting rows from table '{self.table.fullname}': {e}")
             sess.rollback()
             return False
 

agno/workflow/parallel.py

@@ -102,7 +102,7 @@ class Parallel:
                 step_name=self.name or "Parallel",
                 step_id=str(uuid4()),
                 step_type=StepType.PARALLEL,
-                content=f"Parallel {self.name or 'execution'} completed with 1 result",
+                content=self._build_aggregated_content(step_outputs),
                 executor_name=self.name or "Parallel",
                 images=single_result.images,
                 videos=single_result.videos,
@@ -116,8 +116,8 @@ class Parallel:
 
         early_termination_requested = any(output.stop for output in step_outputs if hasattr(output, "stop"))
 
-        # Multiple results - aggregate them
-        aggregated_content = f"Parallel {self.name or 'execution'} completed with {len(step_outputs)} results"
+        # Multiple results - aggregate them with actual content from all steps
+        aggregated_content = self._build_aggregated_content(step_outputs)
 
         # Combine all media from parallel steps
         all_images = []
@@ -419,6 +419,7 @@ class Parallel:
                     step_index=sub_step_index,
                     store_executor_outputs=store_executor_outputs,
                     session_state=step_session_state,
+                    run_context=run_context,
                     parent_step_id=parallel_step_id,
                     workflow_session=workflow_session,
                     add_workflow_history_to_steps=add_workflow_history_to_steps,
@@ -580,6 +581,7 @@ class Parallel:
                     add_workflow_history_to_steps=add_workflow_history_to_steps,
                     num_history_runs=num_history_runs,
                     session_state=step_session_state,
+                    run_context=run_context,
                     background_tasks=background_tasks,
                 )  # type: ignore[union-attr]
                 return idx, inner_step_result, step_session_state

agno/workflow/step.py

@@ -1363,10 +1363,22 @@ class Step:
 
         For container steps (Steps, Router, Loop, etc.), this will recursively find the content from the
         last actual step rather than using the generic container message.
+
+        For Parallel steps, aggregates content from ALL inner steps (not just the last one).
         """
-        # If this step has nested steps (like Steps, Condition, Router, Loop, etc.)
+        # If this step has nested steps (like Steps, Condition, Router, Loop, Parallel, etc.)
         if hasattr(step_output, "steps") and step_output.steps and len(step_output.steps) > 0:
-            # Recursively get content from the last nested step
+            # For Parallel steps, aggregate content from ALL inner steps
+            if step_output.step_type == StepType.PARALLEL:
+                aggregated_parts = []
+                for i, inner_step in enumerate(step_output.steps):
+                    inner_content = self._get_deepest_content_from_step_output(inner_step)
+                    if inner_content:
+                        step_name = inner_step.step_name or f"Step {i + 1}"
+                        aggregated_parts.append(f"=== {step_name} ===\n{inner_content}")
+                return "\n\n".join(aggregated_parts) if aggregated_parts else step_output.content  # type: ignore
+
+            # For other nested step types, recursively get content from the last nested step
             return self._get_deepest_content_from_step_output(step_output.steps[-1])
 
         # For regular steps, return their content

agno/workflow/types.py

@@ -107,10 +107,46 @@ class StepInput:
             return str(self.input)
 
     def get_step_output(self, step_name: str) -> Optional["StepOutput"]:
-        """Get output from a specific previous step by name"""
+        """Get output from a specific previous step by name
+
+        Searches recursively through nested steps (Parallel, Condition, Router, Loop, Steps)
+        to find step outputs at any depth.
+        """
         if not self.previous_step_outputs:
             return None
-        return self.previous_step_outputs.get(step_name)
+
+        # First try direct lookup
+        direct = self.previous_step_outputs.get(step_name)
+        if direct:
+            return direct
+
+        # Search recursively in nested steps
+        return self._search_nested_steps(step_name)
+
+    def _search_nested_steps(self, step_name: str) -> Optional["StepOutput"]:
+        """Recursively search for a step output in nested steps (Parallel, Condition, etc.)"""
+        if not self.previous_step_outputs:
+            return None
+
+        for step_output in self.previous_step_outputs.values():
+            result = self._search_in_step_output(step_output, step_name)
+            if result:
+                return result
+        return None
+
+    def _search_in_step_output(self, step_output: "StepOutput", step_name: str) -> Optional["StepOutput"]:
+        """Helper to recursively search within a single StepOutput"""
+        if not step_output.steps:
+            return None
+
+        for nested_step in step_output.steps:
+            if nested_step.step_name == step_name:
+                return nested_step
+            # Recursively search deeper
+            result = self._search_in_step_output(nested_step, step_name)
+            if result:
+                return result
+        return None
 
     def get_step_content(self, step_name: str) -> Optional[Union[str, Dict[str, str]]]:
         """Get content from a specific previous step by name

agno/workflow/workflow.py

@@ -2569,6 +2569,8 @@ class Workflow:
                 else:
                     # Update status to RUNNING and save
                     workflow_run_response.status = RunStatus.running
+
+                    workflow_session.upsert_run(run=workflow_run_response)
                     if self._has_async_db():
                         await self.asave_session(session=workflow_session)
                     else:
@@ -3449,6 +3451,7 @@ class Workflow:
         input: Optional[Union[str, Dict[str, Any], List[Any], BaseModel]] = None,
         additional_data: Optional[Dict[str, Any]] = None,
         user_id: Optional[str] = None,
+        run_id: Optional[str] = None,
         session_id: Optional[str] = None,
         session_state: Optional[Dict[str, Any]] = None,
         audio: Optional[List[Audio]] = None,
@@ -3468,6 +3471,7 @@ class Workflow:
         input: Optional[Union[str, Dict[str, Any], List[Any], BaseModel]] = None,
         additional_data: Optional[Dict[str, Any]] = None,
         user_id: Optional[str] = None,
+        run_id: Optional[str] = None,
         session_id: Optional[str] = None,
         session_state: Optional[Dict[str, Any]] = None,
         audio: Optional[List[Audio]] = None,
@@ -3486,6 +3490,7 @@ class Workflow:
         input: Optional[Union[str, Dict[str, Any], List[Any], BaseModel]] = None,
         additional_data: Optional[Dict[str, Any]] = None,
         user_id: Optional[str] = None,
+        run_id: Optional[str] = None,
         session_id: Optional[str] = None,
         session_state: Optional[Dict[str, Any]] = None,
         audio: Optional[List[Audio]] = None,
@@ -3503,8 +3508,8 @@ class Workflow:
         if self._has_async_db():
             raise Exception("`run()` is not supported with an async DB. Please use `arun()`.")
 
-        # Create a run_id for this specific run and register immediately for cancellation tracking
-        run_id = str(uuid4())
+        # Set the id for the run and register it immediately for cancellation tracking
+        run_id = run_id or str(uuid4())
         register_run(run_id)
 
         input = self._validate_input(input)
@@ -3621,6 +3626,7 @@ class Workflow:
         input: Optional[Union[str, Dict[str, Any], List[Any], BaseModel, List[Message]]] = None,
         additional_data: Optional[Dict[str, Any]] = None,
         user_id: Optional[str] = None,
+        run_id: Optional[str] = None,
         session_id: Optional[str] = None,
         session_state: Optional[Dict[str, Any]] = None,
         audio: Optional[List[Audio]] = None,
@@ -3641,6 +3647,7 @@ class Workflow:
         input: Optional[Union[str, Dict[str, Any], List[Any], BaseModel, List[Message]]] = None,
         additional_data: Optional[Dict[str, Any]] = None,
         user_id: Optional[str] = None,
+        run_id: Optional[str] = None,
         session_id: Optional[str] = None,
         session_state: Optional[Dict[str, Any]] = None,
         audio: Optional[List[Audio]] = None,
@@ -3660,6 +3667,7 @@ class Workflow:
         input: Optional[Union[str, Dict[str, Any], List[Any], BaseModel, List[Message]]] = None,
         additional_data: Optional[Dict[str, Any]] = None,
         user_id: Optional[str] = None,
+        run_id: Optional[str] = None,
         session_id: Optional[str] = None,
         session_state: Optional[Dict[str, Any]] = None,
         audio: Optional[List[Audio]] = None,
@@ -3730,8 +3738,8 @@ class Workflow:
 
         self._set_debug()
 
-        # Create a run_id for this specific run and register immediately for cancellation tracking
-        run_id = str(uuid4())
+        # Set the id for the run and register it immediately for cancellation tracking
+        run_id = run_id or str(uuid4())
         register_run(run_id)
 
         self.initialize_workflow()