turnkit 0.2.8 → 0.2.10

data/lib/turnkit/agent.rb

@@ -3,13 +3,14 @@
 module TurnKit
   class Agent
     attr_reader :name, :description, :model, :instructions, :tools, :skills, :available_skills, :sub_agents
-    attr_reader :client, :store, :max_iterations, :timeout, :cost_limit, :max_depth, :max_tool_executions
+    attr_reader :client, :store, :max_iterations, :timeout, :cost_limit, :max_depth, :max_tool_executions, :max_tool_executions_by_name
     attr_reader :prompt_sections, :system_prompt, :prompt_mode, :thinking, :compaction, :output_schema, :on_event
+    attr_reader :output_audit, :output_audit_mode, :output_policy_model
 
     def initialize(name:, description: "", model: nil, instructions: "", tools: [], skills: [], available_skills: [], sub_agents: [],
       system_prompt: nil, prompt_sections: nil, prompt_mode: nil, client: nil, store: nil,
-      max_iterations: nil, timeout: nil, cost_limit: nil, max_depth: nil, max_tool_executions: nil, thinking: nil, compaction: nil,
-      output_schema: nil, on_event: nil)
+      max_iterations: nil, timeout: nil, cost_limit: nil, max_depth: nil, max_tool_executions: nil, max_tool_executions_by_name: nil, thinking: nil, compaction: nil,
+      output_schema: nil, output_audit: nil, output_audit_mode: nil, output_policy: nil, output_policy_mode: nil, output_policy_model: nil, output_policy_thinking: nil, on_event: nil)
       @name = name.to_s
       @description = description.to_s
       @model = model
@@ -28,9 +29,13 @@ module TurnKit
       @cost_limit = cost_limit
       @max_depth = max_depth
       @max_tool_executions = max_tool_executions
+      @max_tool_executions_by_name = max_tool_executions_by_name
       @thinking = self.class.normalize_thinking(thinking)
       @compaction = compaction
       @output_schema = output_schema
+      @output_policy_model = output_policy_model
+      @output_audit = normalize_output_policy_options(output_audit: output_audit, output_policy: output_policy, output_policy_model: output_policy_model, output_policy_thinking: output_policy_thinking)
+      @output_audit_mode = normalize_output_policy_mode(output_audit_mode: output_audit_mode, output_policy_mode: output_policy_mode)
       @on_event = on_event
       raise ArgumentError, "name is required" if @name.empty?
       validate_tools!
@@ -62,7 +67,7 @@ module TurnKit
       Conversation.new(agent: self, record: record, store: store, model: model || effective_model, subject: subject, metadata: metadata)
     end
 
-    def run(prompt = nil, task: nil, input: nil, async: false, subject: nil, metadata: {}, parent_run: nil, root_turn_id: nil, **options)
+    def run(prompt = nil, task: nil, input: nil, async: false, subject: nil, metadata: {}, parent_run: nil, root_turn_id: nil, prompt_mode: :task, **options)
       task = task || prompt
       raise ArgumentError, "task is required" if task.to_s.empty?
 
@@ -71,6 +76,7 @@ module TurnKit
       turn = conversation.build_turn(
         trigger_message_id: message.id,
         root_turn_id: root_turn_id || parent_run_root_turn_id(parent_run),
+        prompt_mode: prompt_mode,
         **options
       )
       run = Run.new(turn)
@@ -93,6 +99,18 @@ module TurnKit
       thinking
     end
 
+    def effective_output_audit
+      Array(output_audit).compact
+    end
+
+    def output_policy
+      output_audit
+    end
+
+    def output_policy_mode
+      output_audit_mode
+    end
+
     def effective_client
       client || TurnKit.client
     end
@@ -142,6 +160,7 @@ module TurnKit
         timeout: timeout || TurnKit.timeout,
         max_depth: max_depth || TurnKit.max_depth,
         max_tool_executions: max_tool_executions || TurnKit.max_tool_executions,
+        max_tool_executions_by_name: max_tool_executions_by_name || TurnKit.max_tool_executions_by_name,
         cost_limit: cost_limit || TurnKit.cost_limit,
         root_started_at: root_started_at
       )
@@ -169,6 +188,53 @@ module TurnKit
         effective_tools.each(&:validate_definition!)
       end
 
+      def normalize_output_policy_options(output_audit:, output_policy:, output_policy_model:, output_policy_thinking:)
+        raise ArgumentError, "use output_policy: or output_audit:, not both" if output_audit && output_policy
+
+        output_policy.nil? ? output_audit : normalize_output_policy(output_policy, model: output_policy_model, thinking: output_policy_thinking)
+      end
+
+      def normalize_output_policy(value, model: nil, thinking: nil)
+        case value
+        when nil
+          nil
+        when Array
+          value.map { |item| normalize_output_policy(item, model: model, thinking: thinking) }.compact
+        when String
+          output_policy_from_path(value, model: model, thinking: thinking)
+        when Pathname
+          output_policy_from_path(value.to_s, model: model, thinking: thinking)
+        else
+          return value if value.respond_to?(:call) || value.respond_to?(:check)
+
+          raise ArgumentError, "output_policy must be a policy file path, a #call/#check object, or an array of those"
+        end
+      end
+
+      def output_policy_from_path(path, model: nil, thinking: nil)
+        unless path.match?(/\.(md|markdown|txt)\z/i)
+          raise ArgumentError, "output_policy string must be a .md, .markdown, or .txt file path"
+        end
+
+        TurnKit::OutputPolicy.from_file(
+          path,
+          model: model || TurnKit.output_policy_model,
+          thinking: thinking || TurnKit.output_policy_thinking
+        )
+      end
+
+      def normalize_output_policy_mode(output_audit_mode:, output_policy_mode:)
+        if output_audit_mode && output_policy_mode && output_audit_mode.to_sym != output_policy_mode.to_sym
+          raise ArgumentError, "use output_policy_mode: or output_audit_mode:, not both"
+        end
+
+        value = output_policy_mode || output_audit_mode || :report
+        mode = value.to_sym
+        raise ArgumentError, "unknown output_policy_mode: #{value}" unless %i[report fail].include?(mode)
+
+        mode
+      end
+
       def task_message(task, input)
         text = task.to_s
         return text if input.nil?

data/lib/turnkit/budget.rb

@@ -2,32 +2,40 @@
 
 module TurnKit
   class Budget
-    attr_reader :root_started_at, :max_iterations, :timeout, :max_depth, :max_tool_executions, :cost_limit
+    attr_reader :root_started_at, :max_iterations, :timeout, :max_depth, :max_tool_executions, :max_tool_executions_by_name, :cost_limit
 
-    def initialize(max_iterations:, timeout:, max_depth:, max_tool_executions:, cost_limit: nil, root_started_at: Clock.now)
+    def initialize(max_iterations:, timeout:, max_depth:, max_tool_executions:, max_tool_executions_by_name: {}, cost_limit: nil, root_started_at: Clock.now)
       @root_started_at = root_started_at
       @max_iterations = max_iterations
       @timeout = timeout
       @max_depth = max_depth
       @max_tool_executions = max_tool_executions
+      @max_tool_executions_by_name = normalize_tool_limits(max_tool_executions_by_name)
       @cost_limit = cost_limit
       @iterations = 0
       @tool_executions = 0
+      @tool_executions_by_name = Hash.new(0)
       @cost = 0
       @mutex = Mutex.new
     end
 
     def count_iteration!
       @mutex.synchronize do
+        raise BudgetError, "maximum iterations reached" if max_iterations && @iterations >= max_iterations
+
         @iterations += 1
-        raise Error, "maximum iterations reached" if max_iterations && @iterations > max_iterations
       end
     end
 
-    def count_tool_execution!
+    def count_tool_execution!(name = nil)
       @mutex.synchronize do
+        key = name.to_s if name
+        limit = max_tool_executions_by_name[key] if key
+        raise BudgetError, "maximum tool executions reached" if max_tool_executions && @tool_executions >= max_tool_executions
+        raise BudgetError, "maximum executions reached for tool #{key}" if limit && @tool_executions_by_name[key] >= limit
+
         @tool_executions += 1
-        raise Error, "maximum tool executions reached" if max_tool_executions && @tool_executions > max_tool_executions
+        @tool_executions_by_name[key] += 1 if key
       end
     end
 
@@ -40,13 +48,20 @@ module TurnKit
 
       @mutex.synchronize do
         @cost += cost.to_f
-        raise Error, "cost limit reached" if @cost > cost_limit
+        raise BudgetError, "cost limit reached" if @cost > cost_limit
       end
     end
 
     def check!(depth:)
-      raise Error, "maximum sub-agent depth reached" if max_depth && depth > max_depth
-      raise Error, "turn timed out" if timeout && Clock.now >= root_started_at + timeout
+      raise BudgetError, "maximum sub-agent depth reached" if max_depth && depth > max_depth
+      raise BudgetError, "turn timed out" if timeout && Clock.now >= root_started_at + timeout
     end
+
+    private
+      def normalize_tool_limits(value)
+        value.to_h.transform_keys(&:to_s).transform_values do |limit|
+          limit.nil? ? nil : Integer(limit)
+        end
+      end
   end
 end

data/lib/turnkit/compaction.rb

@@ -117,6 +117,8 @@ module TurnKit
       return unless force || over_threshold?(messages, policy)
 
       compact!(turn.conversation, agent: turn.agent, turn: turn, focus: focus, auto: true, overrides: policy, force: true)
+    rescue BudgetError
+      raise
     rescue StandardError => error
       TurnKit.logger&.warn("TurnKit compaction failed: #{error.class}: #{error.message}")
       nil
@@ -144,12 +146,15 @@ module TurnKit
         target_tokens: summary_budget(selected_tokens, policy),
         fallback_model: turn&.model || conversation.model || agent.effective_model,
         conversation_id: conversation.id,
-        turn_id: turn&.id
+        turn_id: turn&.id,
+        turn: turn
       )
 
       append_summary(conversation, turn: turn, summary: summary, selected: selected, policy: policy, focus: focus, auto: auto, input_tokens: selected_tokens)
     rescue CompactionError
       raise
+    rescue BudgetError
+      raise
     rescue StandardError => error
       raise CompactionError, "#{error.class}: #{error.message}"
     end
@@ -350,18 +355,24 @@ module TurnKit
       index
     end
 
-    def generate_summary(agent:, policy:, messages:, previous_summary:, focus:, target_tokens:, fallback_model:, conversation_id:, turn_id:)
+    def generate_summary(agent:, policy:, messages:, previous_summary:, focus:, target_tokens:, fallback_model:, conversation_id:, turn_id:, turn: nil)
       client = policy["client"] || agent.effective_client
       model = policy["model"] || fallback_model
       safe_messages = messages.map { |message| sanitize_message(message, policy) }
       prompt = build_prompt(previous_summary: previous_summary, focus: focus, target_tokens: target_tokens)
-      result = client.chat(
+      attrs = {
         model: model,
         messages: MessageProjection.for(safe_messages) + [ { role: :user, content: prompt } ],
         tools: [],
         instructions: COMPACTION_SYSTEM_PROMPT,
         metadata: { compaction: true, conversation_id: conversation_id, turn_id: turn_id }
-      )
+      }
+      result = if turn
+        turn.internal_model_call(**attrs, purpose: "compaction", client: policy["client"])
+      else
+        client.validate!(model: model)
+        client.chat(**attrs)
+      end
       text = result.text.to_s.strip
       raise CompactionError, "compaction model returned an empty summary" if text.empty?
 

data/lib/turnkit/conversation.rb

@@ -26,17 +26,18 @@ module TurnKit
       async ? turn : turn.run!
     end
 
-    def run!(trigger_message_id: nil, model: nil, budget: nil, parent_turn: nil, parent_tool_execution: nil, root_turn_id: nil, depth: 0, agent: self.agent, thinking: THINKING_UNSET, compact: nil, output_schema: nil, on_event: nil)
-      build_turn(trigger_message_id: trigger_message_id, model: model, budget: budget, parent_turn: parent_turn, parent_tool_execution: parent_tool_execution, root_turn_id: root_turn_id, depth: depth, agent: agent, thinking: thinking, compact: compact, output_schema: output_schema, on_event: on_event).run!
+    def run!(trigger_message_id: nil, model: nil, budget: nil, parent_turn: nil, parent_tool_execution: nil, root_turn_id: nil, depth: 0, agent: self.agent, thinking: THINKING_UNSET, compact: nil, output_schema: nil, prompt_mode: nil, on_event: nil)
+      build_turn(trigger_message_id: trigger_message_id, model: model, budget: budget, parent_turn: parent_turn, parent_tool_execution: parent_tool_execution, root_turn_id: root_turn_id, depth: depth, agent: agent, thinking: thinking, compact: compact, output_schema: output_schema, prompt_mode: prompt_mode, on_event: on_event).run!
     end
 
-    def build_turn(trigger_message_id: nil, model: nil, budget: nil, parent_turn: nil, parent_tool_execution: nil, root_turn_id: nil, depth: 0, agent: self.agent, thinking: THINKING_UNSET, compact: nil, output_schema: nil, on_event: nil)
+    def build_turn(trigger_message_id: nil, model: nil, budget: nil, parent_turn: nil, parent_tool_execution: nil, root_turn_id: nil, depth: 0, agent: self.agent, thinking: THINKING_UNSET, compact: nil, output_schema: nil, prompt_mode: nil, on_event: nil)
       snapshot = latest_message_sequence
       effective_thinking = thinking.equal?(THINKING_UNSET) ? agent.effective_thinking : Agent.normalize_thinking(thinking)
       options = { "trigger_message_id" => trigger_message_id }.compact
       options["thinking"] = effective_thinking
       options["compact"] = compact unless compact.nil?
       options["output_schema"] = output_schema || agent.output_schema if output_schema || agent.output_schema
+      options["prompt_mode"] = prompt_mode.to_sym if prompt_mode
       record = store.create_turn(
         "conversation_id" => id,
         "agent_name" => agent.name,

data/lib/turnkit/error.rb

@@ -2,6 +2,7 @@
 
 module TurnKit
   class Error < StandardError; end
+  class BudgetError < Error; end
   class ConfigError < Error; end
   class CompactionError < Error; end
   class ModelAccessError < ConfigError; end

data/lib/turnkit/output_audit.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module TurnKit
+  class OutputAudit
+    Violation = Struct.new(:rule, :message, :metadata, keyword_init: true) do
+      def to_h
+        { "rule" => rule.to_s, "message" => message.to_s, "metadata" => metadata || {} }
+      end
+    end
+
+    Result = Struct.new(:violations, keyword_init: true) do
+      def clean?
+        violations.empty?
+      end
+
+      def messages
+        violations.map(&:message)
+      end
+
+      def to_h
+        { "clean" => clean?, "violations" => violations.map(&:to_h) }
+      end
+    end
+
+    def self.check(output, constraints: [], context: {})
+      new(output, constraints: constraints, context: context).check
+    end
+
+    def initialize(output, constraints: [], context: {})
+      @output = output
+      @constraints = Array(constraints)
+      @context = context || {}
+    end
+
+    def check
+      Result.new(violations: constraints.flat_map { |constraint| normalize(check_constraint(constraint)) })
+    end
+
+    private
+      attr_reader :output, :constraints, :context
+
+      def check_constraint(constraint)
+        if constraint.respond_to?(:check)
+          call_with_optional_context(constraint.method(:check))
+        elsif constraint.respond_to?(:call)
+          callable = constraint.is_a?(Proc) ? constraint : constraint.method(:call)
+          call_with_optional_context(callable)
+        else
+          raise ArgumentError, "output constraints must respond to #call or #check"
+        end
+      end
+
+      def call_with_optional_context(method)
+        parameters = method.parameters
+        return method.call(output) unless parameters.any? { |kind, _| %i[key keyreq keyrest].include?(kind) }
+        return method.call(output, **context) if parameters.any? { |kind, _| kind == :keyrest }
+
+        accepted = parameters.filter_map { |kind, name| name if %i[key keyreq].include?(kind) }
+        method.call(output, **context.slice(*accepted))
+      end
+
+      def normalize(value)
+        case value
+        when nil, false, true
+          []
+        when Violation
+          [ value ]
+        when Result
+          value.violations
+        when String
+          [ Violation.new(rule: "output_constraint", message: value, metadata: {}) ]
+        when Hash
+          [ violation_from_hash(value) ]
+        else
+          if value.respond_to?(:to_ary)
+            value.to_ary.flat_map { |item| normalize(item) }
+          else
+            raise ArgumentError, "output constraint returned unsupported value: #{value.class}"
+          end
+        end
+      end
+
+      def violation_from_hash(value)
+        attrs = value.transform_keys(&:to_s)
+        Violation.new(
+          rule: attrs["rule"] || "output_constraint",
+          message: attrs["message"] || attrs["error"] || "output constraint failed",
+          metadata: attrs["metadata"] || attrs.reject { |key, _| %w[rule message error].include?(key) }
+        )
+      end
+  end
+end

data/lib/turnkit/output_policy.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+module TurnKit
+  class OutputPolicy
+    DEFAULT_SCHEMA = {
+      type: "object",
+      properties: {
+        approved: { type: "boolean" },
+        violations: {
+          type: "array",
+          items: {
+            type: "object",
+            properties: {
+              rule: { type: "string" },
+              message: { type: "string" }
+            },
+            required: [ "rule", "message" ]
+          }
+        }
+      },
+      required: [ "approved", "violations" ]
+    }.freeze
+
+    attr_reader :name, :content, :model, :thinking, :client
+
+    def self.from_file(path, name: nil, **options)
+      new(name: name || File.basename(path, File.extname(path)), content: File.read(path), **options)
+    end
+
+    def initialize(content:, name: "output_policy", model: nil, thinking: nil, client: nil)
+      @name = name.to_s
+      @content = content.to_s
+      @model = model
+      @thinking = Agent.normalize_thinking(thinking)
+      @client = client
+      raise ArgumentError, "content is required" if @content.empty?
+    end
+
+    def call(output, run: nil, turn: nil)
+      model_name = model || turn&.model || run&.turn&.model || TurnKit.default_model
+      result = if turn
+        turn.internal_model_call(
+          model: model_name,
+          messages: audit_messages(output),
+          tools: [],
+          instructions: audit_instructions,
+          thinking: thinking,
+          output_schema: DEFAULT_SCHEMA,
+          metadata: { output_policy: name },
+          purpose: "output_policy",
+          client: client
+        )
+      else
+        audit_client = client || TurnKit.client
+        audit_client.validate!(model: model_name)
+        chat(audit_client, model: model_name, messages: audit_messages(output), tools: [], instructions: audit_instructions, thinking: thinking, output_schema: DEFAULT_SCHEMA, metadata: { output_policy: name })
+      end
+      data = result.output_data || parse_json(result.text)
+      return if data.fetch("approved", false)
+
+      Array(data["violations"]).map do |violation|
+        attrs = violation.transform_keys(&:to_s)
+        OutputAudit::Violation.new(
+          rule: attrs["rule"] || name,
+          message: attrs["message"] || "output policy failed",
+          metadata: attrs.reject { |key, _| %w[rule message].include?(key) }
+        )
+      end
+    end
+
+    private
+      def audit_instructions
+        <<~TEXT
+          You audit model outputs against the policy below.
+
+          Return only a JSON object matching this shape:
+          {"approved":true,"violations":[]}
+
+          Set approved to true only when the output satisfies the policy. For each violation, include a concise rule and message. Do not repair the output. Do not wrap the JSON in Markdown. Do not include commentary before or after the JSON.
+
+          Policy:
+          #{content}
+        TEXT
+      end
+
+      def audit_messages(output)
+        [ { role: :user, content: JSON.generate(output: output) } ]
+      end
+
+      def chat(client, **kwargs)
+        accepted = chat_keyword_names(client)
+        kwargs = kwargs.slice(*accepted) unless accepted.include?(:keyrest)
+        client.chat(**kwargs)
+      end
+
+      def chat_keyword_names(client)
+        client.method(:chat).parameters.filter_map do |kind, name|
+          return [ :keyrest ] if kind == :keyrest
+
+          name if %i[key keyreq].include?(kind)
+        end
+      end
+
+      def parse_json(value)
+        JSON.parse(extract_json(value.to_s))
+      rescue JSON::ParserError
+        { "approved" => false, "violations" => [ { "rule" => name, "message" => "output policy returned invalid JSON" } ] }
+      end
+
+      def extract_json(value)
+        text = value.strip
+        return text if text.start_with?("{") && text.end_with?("}")
+
+        fenced = text[/```(?:json)?\s*(\{.*?\})\s*```/m, 1]
+        return fenced if fenced
+
+        object = text[/\{.*\}/m]
+        object || text
+      end
+  end
+end

data/lib/turnkit/run.rb

@@ -14,6 +14,8 @@ module TurnKit
     def output = output_text
     def output_text = turn.output_text
     def output_data = turn.output_data
+    def output_audit = turn.output_audit
+    def output_audit_clean? = output_audit.nil? || output_audit.fetch("clean", false)
     def usage = Usage.from_records(turn_records)
     def cost = Cost.from_records(turn_records)
     def steps = turn_records.length

data/lib/turnkit/tool_runner.rb

@@ -23,10 +23,17 @@ module TurnKit
       attr_reader :turn
 
       def run(tool_call)
-        turn.budget.count_tool_execution!
-        tool = tool_for(tool_call.name)
         execution = ToolExecution.new(create_execution(tool_call))
 
+        begin
+          turn.budget.count_tool_execution!(tool_call.name)
+        rescue BudgetError => error
+          finish_error(execution, tool_call, error.message, details: { "class" => error.class.name, "budget_denied" => true })
+          raise
+        end
+
+        tool = tool_for(tool_call.name)
+
         unless tool
           return finish_error(execution, tool_call, "unknown tool: #{tool_call.name}")
         end
@@ -58,7 +65,7 @@ module TurnKit
       def finish_success(execution, tool_call, payload)
         attrs = turn.store.update_tool_execution(execution.id, "status" => "completed", "result" => payload, "completed_at" => Clock.now)
         append_result(execution, tool_call, payload)
-        turn.emit("tool_call.completed", id: tool_call.id, name: tool_call.name)
+        turn.emit("tool_call.completed", id: tool_call.id, name: tool_call.name, result_chars: payload.to_json.length)
         ToolExecution.new(attrs)
       end
 
@@ -66,7 +73,7 @@ module TurnKit
         error = { "message" => message.to_s, "details" => details }.compact
         attrs = turn.store.update_tool_execution(execution.id, "status" => "failed", "error" => error, "completed_at" => Clock.now)
         append_result(execution, tool_call, error)
-        turn.emit("tool_call.failed", id: tool_call.id, name: tool_call.name, error: error)
+        turn.emit("tool_call.failed", id: tool_call.id, name: tool_call.name, error: error, result_chars: error.to_json.length)
         ToolExecution.new(attrs)
       end