legion-data 1.7.3 → 1.8.0

data/lib/legion/data/models/identity_provider.rb

@@ -1,13 +1,21 @@
 # frozen_string_literal: true
 
+require_relative 'identity/model_helpers'
+
 return unless Legion::Data::Connection.adapter == :postgres
 
 module Legion
   module Data
     module Model
       class IdentityProvider < Sequel::Model(:identity_providers)
+        include Identity::ModelHelpers
+
         one_to_many :identities, class: 'Legion::Data::Model::Identity'
 
+        def self.lookup_columns
+          %i[id uuid name]
+        end
+
         def parsed_capabilities
           Array(capabilities)
         end

data/lib/legion/data/models/llm/conversation.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class Conversation < Sequel::Model(:llm_conversations)
+          include ModelHelpers
+
+          one_to_many :messages
+          one_to_many :message_inference_requests
+          one_to_many :conversation_compactions
+          one_to_many :policy_evaluations
+          one_to_many :security_events
+
+          def security_incident_lineage
+            SecurityEvent.lineage_for_conversation(self)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/conversation_compaction.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class ConversationCompaction < Sequel::Model(:llm_conversation_compactions)
+          include ModelHelpers
+
+          many_to_one :conversation
+          many_to_one :triggered_by_message_inference_request,
+                      class: 'Legion::Data::Models::LLM::MessageInferenceRequest',
+                      key:   :triggered_by_message_inference_request_id
+          many_to_one :replaces_message_from, class: 'Legion::Data::Models::LLM::Message', key: :replaces_message_from_id
+          many_to_one :replaces_message_to, class: 'Legion::Data::Models::LLM::Message', key: :replaces_message_to_id
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/message.rb

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class Message < Sequel::Model(:llm_messages)
+          include ModelHelpers
+
+          many_to_one :conversation
+          many_to_one :parent_message, class: 'Legion::Data::Models::LLM::Message', key: :parent_message_id
+          many_to_one :message_inference_request
+          many_to_one :message_inference_response
+          many_to_one :tool_call
+
+          one_to_many :child_messages, class: 'Legion::Data::Models::LLM::Message', key: :parent_message_id
+          one_to_many :triggered_message_inference_requests,
+                      class: 'Legion::Data::Models::LLM::MessageInferenceRequest',
+                      key:   :latest_message_id
+          one_to_many :message_inference_responses,
+                      class: 'Legion::Data::Models::LLM::MessageInferenceResponse',
+                      key:   :response_message_id
+          one_to_many :requested_tool_calls, class: 'Legion::Data::Models::LLM::ToolCall',
+                                             key:   :requested_by_message_id
+          one_to_many :result_tool_calls, class: 'Legion::Data::Models::LLM::ToolCall',
+                                          key:   :result_message_id
+          one_to_many :compactions_from, class: 'Legion::Data::Models::LLM::ConversationCompaction',
+                                         key:   :replaces_message_from_id
+          one_to_many :compactions_to, class: 'Legion::Data::Models::LLM::ConversationCompaction',
+                                       key:   :replaces_message_to_id
+
+          class << self
+            def incident_flow_from(message_or_id)
+              message = message_or_id.is_a?(self) ? message_or_id : self[message_or_id]
+              message&.incident_flow
+            end
+          end
+
+          def incident_flow
+            requests = incident_flow_requests
+            responses = incident_flow_responses(requests)
+            route_attempts = RouteAttempt.where(message_inference_request_id: requests.map(&:id))
+                                         .order(:message_inference_request_id, :attempt_no, :id)
+                                         .all
+            tool_calls = incident_flow_tool_calls(responses)
+            tool_call_attempts = ToolCallAttempt.where(tool_call_id: tool_calls.map(&:id))
+                                                .order(:tool_call_id, :attempt_no, :id)
+                                                .all
+
+            {
+              message:            self,
+              conversation:       conversation,
+              requests:           requests,
+              route_attempts:     route_attempts,
+              responses:          responses,
+              response_messages:  responses.filter_map(&:response_message),
+              tool_calls:         tool_calls,
+              tool_call_attempts: tool_call_attempts,
+              result_messages:    incident_flow_result_messages(responses, tool_calls)
+            }
+          end
+
+          private
+
+          def incident_flow_requests
+            request_ids = []
+            request_ids << message_inference_request_id if message_inference_request_id
+            request_ids.concat(MessageInferenceRequest.where(latest_message_id: id).select_map(:id))
+            if message_inference_response_id && (linked_response = MessageInferenceResponse[message_inference_response_id])
+              request_ids << linked_response.message_inference_request_id
+            end
+            if tool_call_id && (linked_tool_call = ToolCall[tool_call_id])
+              request_ids << linked_tool_call.message_inference_response.message_inference_request_id
+            end
+
+            MessageInferenceRequest.where(id: request_ids.uniq).order(:id).all
+          end
+
+          def incident_flow_responses(requests)
+            request_ids = requests.map(&:id)
+            response_scope = MessageInferenceResponse.where(message_inference_request_id: request_ids)
+            response_scope = response_scope.or(id: message_inference_response_id) if message_inference_response_id
+            response_scope.order(:id).all
+          end
+
+          def incident_flow_tool_calls(responses)
+            response_ids = responses.map(&:id)
+            scope = ToolCall.where(message_inference_response_id: response_ids)
+            scope = scope.or(requested_by_message_id: id).or(result_message_id: id)
+            scope.order(:message_inference_response_id, :tool_call_index, :id).all
+          end
+
+          def incident_flow_result_messages(responses, tool_calls)
+            message_ids = responses.filter_map(&:response_message_id) + tool_calls.filter_map(&:result_message_id)
+            scope = Message.where(id: message_ids.uniq)
+            scope = scope.or(tool_call_id: tool_calls.map(&:id)) unless tool_calls.empty?
+            scope.order(:seq, :id).all
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/message_inference_metric.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class MessageInferenceMetric < Sequel::Model(:llm_message_inference_metrics)
+          include ModelHelpers
+
+          many_to_one :message_inference_request
+          many_to_one :message_inference_response
+
+          class << self
+            def finance_usage_by_cost_center_model_day(cost_center: nil, model_key: nil, from: nil, to: nil)
+              usage_day = Sequel.function(:date, :recorded_at)
+              scope = dataset
+              scope = scope.where(cost_center: cost_center) unless cost_center.nil?
+              scope = scope.where(model_key: model_key) unless model_key.nil?
+              scope = scope.where { recorded_at >= from } unless from.nil?
+              scope = scope.where { recorded_at < to } unless to.nil?
+
+              scope
+                .select(
+                  :cost_center,
+                  :model_key,
+                  usage_day.as(:usage_day),
+                  Sequel.function(:sum, :input_tokens).as(:input_tokens),
+                  Sequel.function(:sum, :output_tokens).as(:output_tokens),
+                  Sequel.function(:sum, :thinking_tokens).as(:thinking_tokens),
+                  Sequel.function(:sum, :total_tokens).as(:total_tokens),
+                  Sequel.function(:sum, :cost_usd).as(:cost_usd),
+                  Sequel.function(:sum, :latency_ms).as(:latency_ms),
+                  Sequel.function(:sum, :wall_clock_ms).as(:wall_clock_ms)
+                )
+                .group(:cost_center, :model_key, usage_day)
+                .order(:cost_center, :model_key, usage_day)
+                .map(&:values)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/message_inference_request.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class MessageInferenceRequest < Sequel::Model(:llm_message_inference_requests)
+          include ModelHelpers
+
+          many_to_one :conversation
+          many_to_one :latest_message, class: 'Legion::Data::Models::LLM::Message', key: :latest_message_id
+          one_to_many :message_inference_responses
+          one_to_many :route_attempts
+          one_to_many :message_inference_metrics
+          one_to_many :conversation_compactions, key: :triggered_by_message_inference_request_id
+          one_to_many :policy_evaluations
+          one_to_many :security_events
+
+          class << self
+            def lookup(reference)
+              return reference if reference.is_a?(self)
+
+              value = reference.to_s
+              scope = where(uuid: value).or(request_ref: value)
+              scope = scope.or(id: value.to_i) if value.match?(/\A\d+\z/)
+              scope.first
+            end
+
+            def audit_lineage_for(reference)
+              lookup(reference)&.audit_lineage
+            end
+          end
+
+          def audit_lineage
+            responses = message_inference_responses_dataset.order(:id).all
+            response_ids = responses.map(&:id)
+            tool_calls = ToolCall.where(message_inference_response_id: response_ids).order(:tool_call_index, :id).all
+            tool_call_ids = tool_calls.map(&:id)
+
+            {
+              request:            self,
+              request_id:         id,
+              request_ref:        request_ref,
+              conversation:       conversation,
+              latest_message:     latest_message,
+              caller_principal:   caller_principal,
+              caller_identity:    caller_identity,
+              route_attempts:     route_attempts_dataset.order(:attempt_no, :id).all,
+              responses:          responses,
+              response_messages:  responses.filter_map(&:response_message),
+              metrics:            message_inference_metrics_dataset.order(:recorded_at, :id).all,
+              policy_evaluations: policy_evaluations_dataset.order(:evaluated_at, :id).all,
+              security_events:    security_events_dataset.order(:detected_at, :id).all,
+              tool_calls:         tool_calls,
+              tool_call_attempts: ToolCallAttempt.where(tool_call_id: tool_call_ids).order(:tool_call_id, :attempt_no, :id).all
+            }
+          end
+
+          def request
+            self
+          end
+
+          def caller_principal
+            return nil unless caller_principal_id && defined?(Legion::Data::Model::Identity::Principal)
+
+            Legion::Data::Model::Identity::Principal.first(id: caller_principal_id)
+          end
+
+          def caller_identity
+            return nil unless caller_identity_id && defined?(Legion::Data::Model::Identity::Identity)
+
+            Legion::Data::Model::Identity::Identity.first(id: caller_identity_id)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/message_inference_response.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class MessageInferenceResponse < Sequel::Model(:llm_message_inference_responses)
+          include ModelHelpers
+
+          many_to_one :message_inference_request
+          many_to_one :response_message, class: 'Legion::Data::Models::LLM::Message', key: :response_message_id
+          one_to_many :route_attempts
+          one_to_many :message_inference_metrics
+          one_to_many :tool_calls
+          one_to_many :policy_evaluations
+          one_to_many :security_events
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/model_helpers.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        module ModelHelpers
+          def before_create
+            self[:uuid] ||= SecureRandom.uuid if columns.include?(:uuid)
+            super
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/policy_evaluation.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class PolicyEvaluation < Sequel::Model(:llm_policy_evaluations)
+          include ModelHelpers
+
+          many_to_one :conversation
+          many_to_one :message_inference_request
+          many_to_one :message_inference_response
+          one_to_many :security_events
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/registry_event.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class RegistryEvent < Sequel::Model(:llm_registry_events)
+          include ModelHelpers
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/route_attempt.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class RouteAttempt < Sequel::Model(:llm_route_attempts)
+          include ModelHelpers
+
+          many_to_one :message_inference_request
+          many_to_one :message_inference_response
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/security_event.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class SecurityEvent < Sequel::Model(:llm_security_events)
+          include ModelHelpers
+
+          many_to_one :conversation
+          many_to_one :message_inference_request
+          many_to_one :message_inference_response
+          many_to_one :tool_call
+          many_to_one :tool_call_attempt
+          many_to_one :policy_evaluation
+
+          class << self
+            def lineage_for_conversation(conversation_or_id)
+              conversation_id = conversation_or_id.respond_to?(:id) ? conversation_or_id.id : conversation_or_id
+              requests = MessageInferenceRequest.where(conversation_id: conversation_id).order(:id).all
+              request_ids = requests.map(&:id)
+              responses = MessageInferenceResponse.where(message_inference_request_id: request_ids).order(:id).all
+              response_ids = responses.map(&:id)
+              tool_calls = ToolCall.where(message_inference_response_id: response_ids).order(:tool_call_index, :id).all
+              tool_call_ids = tool_calls.map(&:id)
+
+              {
+                conversation:            Conversation[conversation_id],
+                messages:                Message.where(conversation_id: conversation_id).order(:seq, :id).all,
+                requests:                requests,
+                route_attempts:          RouteAttempt.where(message_inference_request_id: request_ids).order(:message_inference_request_id, :attempt_no,
+                                                                                                             :id).all,
+                responses:               responses,
+                request_payload_hashes:  requests.filter_map(&:request_content_hash),
+                response_payload_hashes: responses.filter_map(&:response_content_hash),
+                policy_evaluations:      policy_evaluations_for(conversation_id, request_ids, response_ids),
+                security_events:         security_events_for(conversation_id, request_ids, response_ids, tool_call_ids),
+                tool_calls:              tool_calls,
+                tool_call_attempts:      ToolCallAttempt.where(tool_call_id: tool_call_ids).order(:tool_call_id, :attempt_no, :id).all
+              }
+            end
+
+            private
+
+            def policy_evaluations_for(conversation_id, request_ids, response_ids)
+              scope = PolicyEvaluation.where(conversation_id: conversation_id)
+              scope = scope.or(message_inference_request_id: request_ids) unless request_ids.empty?
+              scope = scope.or(message_inference_response_id: response_ids) unless response_ids.empty?
+              scope.order(:evaluated_at, :id).all
+            end
+
+            def security_events_for(conversation_id, request_ids, response_ids, tool_call_ids)
+              scope = where(conversation_id: conversation_id)
+              scope = scope.or(message_inference_request_id: request_ids) unless request_ids.empty?
+              scope = scope.or(message_inference_response_id: response_ids) unless response_ids.empty?
+              scope = scope.or(tool_call_id: tool_call_ids) unless tool_call_ids.empty?
+              scope.order(:detected_at, :id).all
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/tool_call.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class ToolCall < Sequel::Model(:llm_tool_calls)
+          include ModelHelpers
+
+          many_to_one :message_inference_response
+          many_to_one :requested_by_message, class: 'Legion::Data::Models::LLM::Message', key: :requested_by_message_id
+          many_to_one :result_message, class: 'Legion::Data::Models::LLM::Message', key: :result_message_id
+          one_to_many :tool_call_attempts
+          one_to_many :security_events
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/llm/tool_call_attempt.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Models
+      module LLM
+        class ToolCallAttempt < Sequel::Model(:llm_tool_call_attempts)
+          include ModelHelpers
+
+          many_to_one :tool_call
+          one_to_many :security_events
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/node.rb

@@ -8,7 +8,8 @@ module Legion
       class Node < Sequel::Model
         include Legion::Logging::Helper
 
-        # one_to_many :task_log
+        one_to_many :task_log
+        one_to_many :task_logs, class: 'Legion::Data::Model::TaskLog'
         many_to_one :principal, class: 'Legion::Data::Model::Principal'
 
         def parsed_metrics

data/lib/legion/data/models/principal.rb

@@ -1,13 +1,26 @@
 # frozen_string_literal: true
 
+require_relative 'identity/model_helpers'
+
 return unless Legion::Data::Connection.adapter == :postgres
 
 module Legion
   module Data
     module Model
       class Principal < Sequel::Model(:principals)
+        include Identity::ModelHelpers
+
         one_to_many :identities, class: 'Legion::Data::Model::Identity'
         one_to_many :group_memberships, class: 'Legion::Data::Model::IdentityGroupMembership'
+        many_to_many :groups,
+                     class:      'Legion::Data::Model::IdentityGroup',
+                     join_table: :identity_group_memberships,
+                     left_key:   :principal_id,
+                     right_key:  :group_id
+
+        def self.lookup_columns
+          %i[id uuid canonical_name employee_key employee_id]
+        end
 
         def active_groups
           group_memberships_dataset

data/lib/legion/data/models/rbac/cross_team_grants.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Model
+      module RBAC
+        class CrossTeamGrant < Sequel::Model(:rbac_cross_team_grants)
+          include ModelHelpers
+
+          def validate
+            super
+            errors.add(:source_team, 'cannot be empty') if source_team.nil? || source_team.empty?
+            errors.add(:target_team, 'cannot be empty') if target_team.nil? || target_team.empty?
+            errors.add(:source_team, 'cannot equal target_team') if source_team == target_team
+            errors.add(:runner_pattern, 'cannot be empty') if runner_pattern.nil? || runner_pattern.empty?
+            errors.add(:actions, 'cannot be empty') if actions.nil? || actions.empty?
+            errors.add(:granted_by, 'cannot be empty') if granted_by.nil? || granted_by.empty?
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/rbac/model_helpers.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Legion
+  module Data
+    module Model
+      module RBAC
+        module ModelHelpers
+          def expired?
+            return false if expires_at.nil?
+
+            expires_at < Time.now
+          end
+
+          def active?
+            !expired?
+          end
+
+          def actions_list
+            (actions || '').split(',').map(&:strip)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/rbac/role_assignments.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Model
+      module RBAC
+        class RoleAssignment < Sequel::Model(:rbac_role_assignments)
+          include ModelHelpers
+
+          VALID_PRINCIPAL_TYPES = %w[worker human].freeze
+
+          def validate
+            super
+            errors.add(:principal_type, 'must be worker or human') unless VALID_PRINCIPAL_TYPES.include?(principal_type)
+            errors.add(:principal_id, 'cannot be empty') if principal_id.nil? || principal_id.empty?
+            errors.add(:role, 'cannot be empty') if role.nil? || role.empty?
+            errors.add(:granted_by, 'cannot be empty') if granted_by.nil? || granted_by.empty?
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/rbac/runner_grants.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require_relative 'model_helpers'
+
+module Legion
+  module Data
+    module Model
+      module RBAC
+        class RunnerGrant < Sequel::Model(:rbac_runner_grants)
+          include ModelHelpers
+
+          def validate
+            super
+            errors.add(:team, 'cannot be empty') if team.nil? || team.empty?
+            errors.add(:runner_pattern, 'cannot be empty') if runner_pattern.nil? || runner_pattern.empty?
+            errors.add(:actions, 'cannot be empty') if actions.nil? || actions.empty?
+            errors.add(:granted_by, 'cannot be empty') if granted_by.nil? || granted_by.empty?
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/data/models/relationship.rb

@@ -6,6 +6,7 @@ module Legion
       class Relationship < Sequel::Model
         many_to_one :trigger, class: 'Legion::Data::Model::Function'
         many_to_one :action, class: 'Legion::Data::Model::Function'
+        many_to_one :chain
         one_to_many :tasks
       end
     end