Module: Legion::Extensions::Apollo::Runners::EntityExtractor

Included in:
Actor::EntityWatchdog
Defined in:
lib/legion/extensions/apollo/runners/entity_extractor.rb

Overview

rubocop:disable Legion/Extension/RunnerIncludeHelpers

Constant Summary collapse

DEFAULT_ENTITY_TYPES = 
%w[person service repository concept].freeze
DEFAULT_MIN_CONFIDENCE = 
0.7

Instance Method Summary collapse

Instance Method Details

#entity_extraction_prompt(text:, entity_types:) ⇒ Object 



38
39
40
41
42
43
44
45
46
47
48
49
50
51
 
# File 'lib/legion/extensions/apollo/runners/entity_extractor.rb', line 38

def entity_extraction_prompt(text:, entity_types:, **)
  type_list = Array(entity_types).join(', ')
  <<~PROMPT.strip
    Extract named entities from the following text. Return only entities of these types: #{type_list}.

    For each entity provide:
    - name: the canonical name as it appears (string)
    - type: one of #{type_list} (string)
    - confidence: your confidence this is a real entity of that type (float 0.0-1.0)

    Text:
    #{text}
  PROMPT
end

#entity_schema ⇒ Object 



53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
 
# File 'lib/legion/extensions/apollo/runners/entity_extractor.rb', line 53

def entity_schema
  {
    type:       'object',
    properties: {
      entities: {
        type:  'array',
        items: {
          type:       'object',
          properties: {
            name:       { type: 'string' },
            type:       { type: 'string' },
            confidence: { type: 'number' }
          },
          required:   %w[name type confidence]
        }
      }
    },
    required:   ['entities']
  }
end

#extract_entities(text:, entity_types: nil, min_confidence: Helpers::Confidence.apollo_setting(:entity_extractor, :min_confidence, default: DEFAULT_MIN_CONFIDENCE)) ⇒ Object

rubocop:disable Layout/LineLength



11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
 
# File 'lib/legion/extensions/apollo/runners/entity_extractor.rb', line 11

def extract_entities(text:, entity_types: nil, min_confidence: Helpers::Confidence.apollo_setting(:entity_extractor, :min_confidence, default: DEFAULT_MIN_CONFIDENCE), **) # rubocop:disable Layout/LineLength
  return { success: true, entities: [], source: :empty } if text.to_s.strip.empty?

  return { success: true, entities: [], source: :unavailable } unless defined?(Legion::LLM) && Legion::LLM.started?

  types = Array(entity_types).map(&:to_s)
  types = DEFAULT_ENTITY_TYPES if types.empty?

  result = Legion::LLM.structured(
    messages: [
      { role: 'user', content: entity_extraction_prompt(text: text, entity_types: types) }
    ],
    schema:   entity_schema,
    caller:   { extension: 'lex-apollo', runner: 'entity_extractor' }
  )

  raw_entities = result.dig(:data, :entities) || []
  filtered = raw_entities.select do |entity|
    (entity[:confidence] || 0.0) >= min_confidence &&
      (types.empty? || types.include?(entity[:type].to_s))
  end

  { success: true, entities: filtered, source: :llm }
rescue StandardError => e
  { success: false, entities: [], error: e.message, source: :error }
end