fix: coderabbit comments on the pr

Author: pira998

examples/observability/cross_workflow_tracking/cross_workflow_tracking_example.py

@@ -21,15 +21,14 @@ async def create_simple_config() -> str:
 
     config_content = """
 llms:
-  nvidia_llm:
-    _type: nim
-    model_name: meta/llama-3.1-8b-instruct
-    temperature: 0.7
-    max_tokens: 1024
-
+  demo_llm:
+    _type: nat_test_llm
+    response_seq:
+      - "Stubbed workflow reply."
+    delay_ms: 0
 workflow:
   _type: chat_completion
-  llm_name: nvidia_llm
+  llm_name: demo_llm
   system_prompt: "You are a helpful customer support assistant. Provide clear, concise, and helpful responses."
 """
 

examples/observability/cross_workflow_tracking/example.py

@@ -105,7 +105,7 @@ async def simulate_workflow_execution(workflow_name: str,
         result = await transform_data(input_data)
     elif workflow_name == "storage" or "store" in workflow_name:
         result = await store_data(input_data)
-    elif workflow_name == "report" in workflow_name:
+    elif "report" in workflow_name:
         result = await generate_report(input_data)
     else:
         # Generic processing

src/nat/observability/processor/cross_workflow_processor.py

@@ -104,7 +104,7 @@ async def process(self, item: Span) -> Span:
 
         except (AttributeError, KeyError, TypeError, ValueError) as e:
             # If there's any error in processing, log it but don't fail the span
-            logger.warning(f"Error processing cross-workflow observability data: {e}", exc_info=True)
+            logger.exception("Error processing cross-workflow observability data: %s", e)
             item.set_attribute("observability.processing_error", str(e))
 
         return item
@@ -165,7 +165,7 @@ async def process(self, item: Span) -> Span:
 
         except (AttributeError, IndexError, TypeError) as e:
             # If there's any error in processing, log it but don't fail the span
-            logger.warning(f"Error processing workflow relationship data: {e}", exc_info=True)
+            logger.exception("Error processing workflow relationship data: %s", e)
             item.set_attribute("relationship.processing_error", str(e))
 
         return item

src/nat/observability/workflow_utils.py

@@ -88,27 +88,19 @@ async def invoke_workflow_with_context(workflow: "Workflow[InputT, StreamingOutp
         try:
             async with workflow.run(message, observability_context=obs_context) as runner:
                 result = await runner.result(to_type=to_type)
-
-                # Update workflow metadata on completion
-                if obs_context:
-                    current_workflow = obs_context.get_current_workflow()
-                    if current_workflow:
-                        current_workflow.end_time = time.time()
-                        current_workflow.status = "completed"
-
                 return result
-
-        except Exception as e:
-            # Update workflow metadata on failure and log error
+        finally:
+            exc = sys.exc_info()[1]
             if obs_context:
                 current_workflow = obs_context.get_current_workflow()
                 if current_workflow:
                     current_workflow.end_time = time.time()
-                    current_workflow.status = "failed"
-                    current_workflow.tags["error"] = str(e)
-
-            logger.error(f"Workflow '{workflow_name}' failed with error: {e}", exc_info=True)
-            raise
+                    if exc is None:
+                        current_workflow.status = "completed"
+                    else:
+                        current_workflow.status = "failed"
+                        current_workflow.tags["error"] = str(exc)
+                        logger.error("Workflow '%s' failed: %s", workflow_name, exc)
 
     @staticmethod
     async def invoke_workflow_stream_with_context(
@@ -158,25 +150,18 @@ async def invoke_workflow_stream_with_context(
             async with workflow.run(message, observability_context=obs_context) as runner:
                 async for item in runner.result_stream(to_type=to_type):
                     yield item
-
-                # Update workflow metadata on completion
-                if obs_context:
-                    current_workflow = obs_context.get_current_workflow()
-                    if current_workflow:
-                        current_workflow.end_time = time.time()
-                        current_workflow.status = "completed"
-
-        except Exception as e:
-            # Update workflow metadata on failure and log error
+        finally:
+            exc = sys.exc_info()[1]
             if obs_context:
                 current_workflow = obs_context.get_current_workflow()
                 if current_workflow:
                     current_workflow.end_time = time.time()
-                    current_workflow.status = "failed"
-                    current_workflow.tags["error"] = str(e)
-
-            logger.error(f"Streaming workflow '{workflow_name}' failed with error: {e}", exc_info=True)
-            raise
+                    if exc is None:
+                        current_workflow.status = "completed"
+                    else:
+                        current_workflow.status = "failed"
+                        current_workflow.tags["error"] = str(exc)
+                        logger.error("Streaming workflow '%s' failed: %s", workflow_name, exc)
 
     @staticmethod
     def get_current_observability_context() -> Optional[ObservabilityContext]:
@@ -268,24 +253,16 @@ async def invoke_with_steps_and_context(workflow: "Workflow[InputT, StreamingOut
                 to_type=to_type,
                 observability_context=obs_context
             )
-
-            # Update workflow metadata on completion
-            if obs_context:
-                current_workflow = obs_context.get_current_workflow()
-                if current_workflow:
-                    current_workflow.end_time = time.time()
-                    current_workflow.status = "completed"
-
             return result, steps
-
-        except Exception as e:
-            # Update workflow metadata on failure and log error
+        finally:
+            exc = sys.exc_info()[1]
             if obs_context:
                 current_workflow = obs_context.get_current_workflow()
                 if current_workflow:
                     current_workflow.end_time = time.time()
-                    current_workflow.status = "failed"
-                    current_workflow.tags["error"] = str(e)
-
-            logger.error(f"Workflow with steps '{workflow_name}' failed with error: {e}", exc_info=True)
-            raise
+                    if exc is None:
+                        current_workflow.status = "completed"
+                    else:
+                        current_workflow.status = "failed"
+                        current_workflow.tags["error"] = str(exc)
+                        logger.error("Workflow with steps '%s' failed: %s", workflow_name, exc)