daita-agents 0.1.1__tar.gz → 0.1.2__tar.gz

{daita_agents-0.1.1 → daita_agents-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: daita-agents
-Version: 0.1.1
+Version: 0.1.2
 Summary: Daita Agents - Data focused AI agent framework with free local use and premium hosted enterprise features
 Author-email: Daita <support@daita-tech.io>
 License: Proprietary

{daita_agents-0.1.1 → daita_agents-0.1.2}/daita/agents/base.py

@@ -207,6 +207,9 @@ class BaseAgent(Agent):
         context: Dict[str, Any]
     ) -> Dict[str, Any]:
         """Process task with full reliability features."""
+        # Track processing time
+        start_time = time.time()
+
         # Check backpressure first
         if self.backpressure_controller and not await self.backpressure_controller.acquire_processing_slot():
             raise BackpressureError(
@@ -214,7 +217,7 @@ class BaseAgent(Agent):
                 agent_id=self.agent_id,
                 queue_size=self.backpressure_controller.task_queue.qsize()
             )
-        
+
         # Create task in task manager
         task_id = None
         if self.task_manager:
@@ -227,7 +230,7 @@ class BaseAgent(Agent):
             context['task_id'] = task_id
             # Update task status to running
             await self.task_manager.update_status(task_id, TaskStatus.RUNNING)
-        
+
         try:
             # Automatically trace the entire operation
             async with self.trace_manager.span(
@@ -241,17 +244,22 @@ class BaseAgent(Agent):
                 retry_enabled=str(self.config.retry_enabled),
                 reliability_enabled="true"
             ) as span_id:
-                
+
                 # Execute with or without retry logic
                 if self.config.retry_enabled:
                     result = await self._process_with_retry(span_id, task, data, context)
                 else:
                     result = await self._process_fail_fast(span_id, task, data, context)
-                
+
+                # Add processing time to result
+                processing_time_ms = (time.time() - start_time) * 1000
+                if isinstance(result, dict):
+                    result['processing_time_ms'] = processing_time_ms
+
                 # Update task status to completed
                 if task_id and self.task_manager:
                     await self.task_manager.update_status(task_id, TaskStatus.COMPLETED)
-                
+
                 return result
         
         except Exception as e:
@@ -272,6 +280,9 @@ class BaseAgent(Agent):
         context: Dict[str, Any]
     ) -> Dict[str, Any]:
         """Process task without reliability features (original behavior)."""
+        # Track processing time
+        start_time = time.time()
+
         # Automatically trace the entire operation
         async with self.trace_manager.span(
             operation_name=f"agent_process_{task}",
@@ -283,13 +294,18 @@ class BaseAgent(Agent):
             retry_enabled=str(self.config.retry_enabled),
             reliability_enabled="false"
         ) as span_id:
-            
+
             # Execute with or without retry logic
             if self.config.retry_enabled:
                 result = await self._process_with_retry(span_id, task, data, context)
             else:
                 result = await self._process_fail_fast(span_id, task, data, context)
-            
+
+            # Add processing time to result
+            processing_time_ms = (time.time() - start_time) * 1000
+            if isinstance(result, dict):
+                result['processing_time_ms'] = processing_time_ms
+
             return result
     
     async def _process_with_retry(

{daita_agents-0.1.1 → daita_agents-0.1.2}/daita/agents/substrate.py

@@ -442,15 +442,20 @@ class SubstrateAgent(BaseAgent):
         return response
 
     def _register_builtin_handlers(self):
-        """Register built-in handlers for common tasks."""
-        self.handlers.update({
+        """Register built-in handlers for common tasks (without overwriting custom handlers)."""
+        builtin_handlers = {
             'analyze': self._handle_analyze,
             'transform': self._handle_transform,
             'process_data': self._handle_process_data,
             'custom': self._handle_custom,
             'relay_message': self._handle_relay_message,
             'llm_query': self._handle_llm_query
-        })
+        }
+
+        # Only add built-in handlers if custom handler doesn't already exist
+        for task, handler in builtin_handlers.items():
+            if task not in self.handlers:
+                self.handlers[task] = handler
     
     async def _process_once(
         self,
@@ -489,7 +494,7 @@ class SubstrateAgent(BaseAgent):
         else:
             # No handler and no LLM available - use default behavior
             return self._handle_default(data, context, self)
-        
+
         # Execute the handler
         try:
             if asyncio.iscoroutinefunction(handler):

{daita_agents-0.1.1 → daita_agents-0.1.2}/daita/core/relay.py

@@ -59,10 +59,11 @@ class MessageStatus(str, Enum):
 
 @dataclass
 class ReliableMessage:
-    """A message with acknowledgment tracking."""
+    """A message with acknowledgment tracking and metadata."""
     id: str
     channel: str
     data: Any
+    metadata: Dict[str, Any] = field(default_factory=dict)
     publisher: Optional[str] = None
     timestamp: float = field(default_factory=time.time)
     status: MessageStatus = MessageStatus.PENDING
@@ -101,7 +102,8 @@ class RelayManager:
         self.channels: Dict[str, deque] = {}
 
         # Subscribers: channel_name -> set of callbacks
-        self.subscribers: Dict[str, weakref.WeakSet] = {}
+        # Using regular set instead of WeakSet to prevent premature garbage collection
+        self.subscribers: Dict[str, set] = {}
 
         # Per-channel locks to prevent race conditions between publish and subscribe
         self._channel_locks: Dict[str, asyncio.Lock] = {}
@@ -143,8 +145,60 @@ class RelayManager:
         """Ensure channel exists."""
         if channel not in self.channels:
             self.channels[channel] = deque(maxlen=self.max_messages_per_channel)
-            self.subscribers[channel] = weakref.WeakSet()
+            self.subscribers[channel] = set()  # Regular set to prevent garbage collection
             self._channel_locks[channel] = asyncio.Lock()
+
+    def _extract_metadata(self, agent_response: Dict[str, Any], publisher: Optional[str]) -> Dict[str, Any]:
+        """
+        Extract metadata from agent response for context propagation.
+
+        Automatically captures metrics, token usage, confidence scores, and other
+        context that downstream agents need to make intelligent decisions.
+        """
+        metadata = {
+            'upstream_agent': publisher or agent_response.get('agent_name'),
+            'upstream_agent_id': agent_response.get('agent_id'),
+            'timestamp': agent_response.get('timestamp', time.time()),
+        }
+
+        # Extract processing metrics if available
+        if 'processing_time_ms' in agent_response:
+            metadata['processing_time_ms'] = agent_response['processing_time_ms']
+
+        # Extract token usage if available
+        if 'token_usage' in agent_response:
+            metadata['token_usage'] = agent_response['token_usage']
+
+        # Extract confidence score if available
+        if 'confidence_score' in agent_response:
+            metadata['confidence_score'] = agent_response['confidence_score']
+        elif 'confidence' in agent_response:
+            metadata['confidence_score'] = agent_response['confidence']
+
+        # Extract record count if available
+        if 'record_count' in agent_response:
+            metadata['record_count'] = agent_response['record_count']
+
+        # Extract error information if available
+        if 'error_count' in agent_response:
+            metadata['error_count'] = agent_response['error_count']
+
+        # Extract correlation ID for distributed tracing
+        if 'correlation_id' in agent_response:
+            metadata['correlation_id'] = agent_response['correlation_id']
+        elif 'context' in agent_response and isinstance(agent_response['context'], dict):
+            metadata['correlation_id'] = agent_response['context'].get('correlation_id', str(uuid.uuid4()))
+        else:
+            metadata['correlation_id'] = str(uuid.uuid4())
+
+        # Extract retry information if available
+        if 'retry_info' in agent_response:
+            metadata['retry_info'] = agent_response['retry_info']
+
+        # Extract status (success/error)
+        metadata['status'] = agent_response.get('status', 'unknown')
+
+        return metadata
     
     async def publish(
         self,
@@ -154,11 +208,11 @@ class RelayManager:
         require_ack: Optional[bool] = None
     ) -> Optional[str]:
         """
-        Publish data to a channel.
+        Publish data to a channel with automatic metadata propagation.
 
         Args:
             channel: Channel name
-            agent_response: Full agent response (we extract 'result' field)
+            agent_response: Full agent response (we extract 'result' field and metadata)
             publisher: Optional publisher identifier
             require_ack: Whether this message requires acknowledgment (overrides global setting)
 
@@ -175,29 +229,33 @@ class RelayManager:
             logger.warning(f"No 'result' field found in agent response for channel '{channel}'")
             result_data = agent_response  # Fallback to full response if no result field
 
+        # Extract metadata for context propagation
+        metadata = self._extract_metadata(agent_response, publisher)
+
         self._ensure_channel(channel)
 
         # Determine if we need reliability
         needs_reliability = require_ack if require_ack is not None else self.enable_reliability
 
         if needs_reliability:
-            return await self._publish_reliable(channel, result_data, agent_response, publisher)
+            return await self._publish_reliable(channel, result_data, metadata, publisher)
         else:
-            return await self._publish_fire_and_forget(channel, result_data, agent_response, publisher)
+            return await self._publish_fire_and_forget(channel, result_data, metadata, publisher)
     
     async def _publish_fire_and_forget(
         self,
         channel: str,
         result_data: Any,
-        agent_response: Dict[str, Any],
+        metadata: Dict[str, Any],
         publisher: Optional[str]
     ) -> None:
-        """Publish message without reliability features (original behavior)."""
+        """Publish message without reliability features with metadata propagation."""
         # Use per-channel lock to make publish atomic
         async with self._channel_locks[channel]:
-            # Create message with result data only
+            # Create message with result data and metadata
             message = {
                 'data': result_data,
+                'metadata': metadata,
                 'publisher': publisher,
                 'timestamp': time.time()
             }
@@ -205,8 +263,8 @@ class RelayManager:
             # Store result message
             self.channels[channel].append(message)
 
-            # Notify subscribers with just the result data (while holding lock)
-            await self._notify_subscribers(channel, result_data)
+            # Notify subscribers with result data and metadata (while holding lock)
+            await self._notify_subscribers(channel, result_data, metadata)
 
         logger.debug(f"Published result to channel '{channel}' from {publisher}")
         return None
@@ -215,16 +273,17 @@ class RelayManager:
         self,
         channel: str,
         result_data: Any,
-        agent_response: Dict[str, Any],
+        metadata: Dict[str, Any],
         publisher: Optional[str]
     ) -> str:
-        """Publish message with reliability features."""
+        """Publish message with reliability features and metadata propagation."""
         # Create reliable message
         message_id = uuid.uuid4().hex
         reliable_message = ReliableMessage(
             id=message_id,
             channel=channel,
             data=result_data,
+            metadata=metadata,
             publisher=publisher,
             ack_timeout=self.default_ack_timeout
         )
@@ -234,10 +293,11 @@ class RelayManager:
 
         # Use per-channel lock to make publish atomic
         async with self._channel_locks[channel]:
-            # Create message for channel storage
+            # Create message for channel storage with metadata
             message = {
                 'id': message_id,
                 'data': result_data,
+                'metadata': metadata,
                 'publisher': publisher,
                 'timestamp': time.time(),
                 'requires_ack': True
@@ -254,66 +314,73 @@ class RelayManager:
             self.message_timeouts[message_id] = timeout_task
 
             # Notify subscribers with message ID for acknowledgment (while holding lock)
-            await self._notify_subscribers_reliable(channel, result_data, message_id)
+            await self._notify_subscribers_reliable(channel, result_data, metadata, message_id)
 
         logger.debug(f"Published reliable message {message_id} to channel '{channel}' from {publisher}")
         return message_id
 
-    async def _notify_subscribers(self, channel: str, result_data: Any) -> None:
-        """Notify all subscribers of a channel with result data only."""
+    async def _notify_subscribers(self, channel: str, result_data: Any, metadata: Dict[str, Any]) -> None:
+        """Notify all subscribers of a channel with result data and metadata."""
         if channel not in self.subscribers:
             return
-        
+
         # Get snapshot of subscribers to avoid modification during iteration
         subscriber_list = list(self.subscribers[channel])
-        
+
         if not subscriber_list:
             return
-        
+
         # Notify all subscribers concurrently
         tasks = []
         for subscriber in subscriber_list:
-            task = asyncio.create_task(self._call_subscriber(subscriber, result_data))
+            task = asyncio.create_task(self._call_subscriber(subscriber, result_data, metadata))
             task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
             tasks.append(task)
-        
+
         # Wait for all notifications to complete
         if tasks:
             await asyncio.gather(*tasks, return_exceptions=True)
     
-    async def _notify_subscribers_reliable(self, channel: str, result_data: Any, message_id: str) -> None:
-        """Notify all subscribers of a reliable message."""
+    async def _notify_subscribers_reliable(self, channel: str, result_data: Any, metadata: Dict[str, Any], message_id: str) -> None:
+        """Notify all subscribers of a reliable message with metadata."""
         if channel not in self.subscribers:
             return
-        
+
         # Get snapshot of subscribers to avoid modification during iteration
         subscriber_list = list(self.subscribers[channel])
-        
+
         if not subscriber_list:
             return
-        
-        # Notify all subscribers concurrently with message ID
+
+        # Notify all subscribers concurrently with message ID and metadata
         tasks = []
         for subscriber in subscriber_list:
             task = asyncio.create_task(
-                self._call_subscriber_reliable(subscriber, result_data, message_id)
+                self._call_subscriber_reliable(subscriber, result_data, metadata, message_id)
             )
             task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
             tasks.append(task)
-        
+
         # Wait for all notifications to complete
         if tasks:
             await asyncio.gather(*tasks, return_exceptions=True)
     
-    async def _call_subscriber_reliable(self, callback: Callable, result_data: Any, message_id: str) -> None:
-        """Safely call a subscriber callback for reliable message."""
+    async def _call_subscriber_reliable(self, callback: Callable, result_data: Any, metadata: Dict[str, Any], message_id: str) -> None:
+        """Safely call a subscriber callback for reliable message with metadata."""
         try:
             if asyncio.iscoroutinefunction(callback):
-                # Check if callback supports message_id parameter
+                # Check callback signature to determine what parameters to pass
                 import inspect
                 sig = inspect.signature(callback)
-                if 'message_id' in sig.parameters:
+                params = list(sig.parameters.keys())
+
+                # Call with appropriate parameters based on signature
+                if 'message_id' in params and 'metadata' in params:
+                    await callback(result_data, metadata=metadata, message_id=message_id)
+                elif 'message_id' in params:
                     await callback(result_data, message_id=message_id)
+                elif 'metadata' in params:
+                    await callback(result_data, metadata=metadata)
                 else:
                     await callback(result_data)
             else:
@@ -325,11 +392,20 @@ class RelayManager:
             # NACK the message on callback error
             await self.nack_message(message_id, str(e))
     
-    async def _call_subscriber(self, callback: Callable, result_data: Any) -> None:
-        """Safely call a subscriber callback with result data."""
+    async def _call_subscriber(self, callback: Callable, result_data: Any, metadata: Dict[str, Any]) -> None:
+        """Safely call a subscriber callback with result data and metadata."""
         try:
             if asyncio.iscoroutinefunction(callback):
-                await callback(result_data)
+                # Check callback signature to determine what parameters to pass
+                import inspect
+                sig = inspect.signature(callback)
+                params = list(sig.parameters.keys())
+
+                # Call with metadata if callback accepts it
+                if 'metadata' in params:
+                    await callback(result_data, metadata=metadata)
+                else:
+                    await callback(result_data)
             else:
                 # Run sync callback in thread pool
                 loop = asyncio.get_event_loop()
@@ -606,23 +682,24 @@ class RelayManager:
                 retry_message = {
                     'id': message.id,
                     'data': message.data,
+                    'metadata': message.metadata,
                     'publisher': message.publisher,
                     'timestamp': message.timestamp,
                     'requires_ack': True
                 }
-                
+
                 # Add to channel
                 self.channels[message.channel].append(retry_message)
-                
+
                 # Set up new timeout task for the retry
                 timeout_task = asyncio.create_task(
                     self._handle_message_timeout(message.id, message.ack_timeout)
                 )
                 timeout_task.add_done_callback(lambda t: t.exception() if not t.cancelled() else None)
                 self.message_timeouts[message.id] = timeout_task
-                
+
                 # Notify subscribers again
-                await self._notify_subscribers_reliable(message.channel, message.data, message.id)
+                await self._notify_subscribers_reliable(message.channel, message.data, message.metadata, message.id)
             
             # Clean up retry task reference
             self.message_timeouts.pop(f"{message.id}_retry", None)

{daita_agents-0.1.1 → daita_agents-0.1.2}/daita/core/workflow.py

@@ -499,9 +499,9 @@ class Workflow:
                 raise WorkflowError(f"Failed to set up connection {connection}: {str(e)}")
     
     def _create_traced_callback(self, connection: Connection):
-        """Create a callback that automatically traces workflow communication."""
-        async def traced_callback(data: Any):
-            """Callback that processes relay data and automatically traces communication."""
+        """Create a callback that automatically traces workflow communication and propagates metadata."""
+        async def traced_callback(data: Any, metadata: Optional[Dict[str, Any]] = None):
+            """Callback that processes relay data with automatic metadata propagation."""
             try:
                 # Automatically trace the workflow communication
                 await self._trace_workflow_communication(
@@ -511,28 +511,31 @@ class Workflow:
                     data=data,
                     success=True
                 )
-                
+
+                # Build enriched context with metadata propagation
+                enriched_context = {
+                    'source_agent': connection.from_agent,
+                    'channel': connection.channel,
+                    'workflow': self.name
+                }
+
+                # Add all upstream metadata to context
+                if metadata:
+                    enriched_context.update(metadata)
+
                 # Process the data with the destination agent or agent pool
                 if connection.to_agent in self.agents:
                     # Single agent
                     dest_agent = self.agents[connection.to_agent]
                     if hasattr(dest_agent, 'process'):
-                        await dest_agent.process(connection.task, data, {
-                            'source_agent': connection.from_agent,
-                            'channel': connection.channel,
-                            'workflow': self.name
-                        })
+                        await dest_agent.process(connection.task, data, enriched_context)
                     else:
                         logger.warning(f"Agent '{connection.to_agent}' has no process method")
                 
                 elif connection.to_agent in self.agent_pools:
-                    # Agent pool - submit task to pool
+                    # Agent pool - submit task to pool with enriched context
                     pool = self.agent_pools[connection.to_agent]
-                    await pool.submit_task(connection.task, data, {
-                        'source_agent': connection.from_agent,
-                        'channel': connection.channel,
-                        'workflow': self.name
-                    })
+                    await pool.submit_task(connection.task, data, enriched_context)
                 
                 else:
                     logger.error(f"Destination '{connection.to_agent}' not found in agents or agent pools")
@@ -579,9 +582,9 @@ class Workflow:
                 logger.warning(f"Failed to configure reliability for agent '{agent_name}': {e}")
     
     def _create_reliable_callback(self, connection: Connection):
-        """Create a callback with reliability features enabled."""
-        async def reliable_callback(data: Any, message_id: Optional[str] = None):
-            """Callback that processes relay data with reliability features."""
+        """Create a callback with reliability features and metadata propagation."""
+        async def reliable_callback(data: Any, metadata: Optional[Dict[str, Any]] = None, message_id: Optional[str] = None):
+            """Callback that processes relay data with reliability features and metadata."""
             start_time = time.time()
 
             # Deduplication check (only if message_id provided)
@@ -607,16 +610,23 @@ class Workflow:
                     message_id=message_id
                 )
 
+                # Build enriched context with metadata propagation
+                enriched_context = {
+                    'source_agent': connection.from_agent,
+                    'channel': connection.channel,
+                    'workflow': self.name,
+                    'message_id': message_id,
+                    'reliability_enabled': True
+                }
+
+                # Add all upstream metadata to context
+                if metadata:
+                    enriched_context.update(metadata)
+
                 # Process the data with the destination agent
                 dest_agent = self.agents[connection.to_agent]
                 if hasattr(dest_agent, 'process'):
-                    context = {
-                        'source_agent': connection.from_agent,
-                        'channel': connection.channel,
-                        'workflow': self.name,
-                        'message_id': message_id,
-                        'reliability_enabled': True
-                    }
+                    context = enriched_context
                     
                     # Handle backpressure if enabled
                     if (self.reliability_config and 

{daita_agents-0.1.1 → daita_agents-0.1.2}/daita_agents.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: daita-agents
-Version: 0.1.1
+Version: 0.1.2
 Summary: Daita Agents - Data focused AI agent framework with free local use and premium hosted enterprise features
 Author-email: Daita <support@daita-tech.io>
 License: Proprietary

{daita_agents-0.1.1 → daita_agents-0.1.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "daita-agents"
-version = "0.1.1"
+version = "0.1.2"
 description = "Daita Agents - Data focused AI agent framework with free local use and premium hosted enterprise features"
 readme = "README.md"
 authors = [{name = "Daita", email = "support@daita-tech.io"}]