llmemory 0.2.6 → 0.2.7

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (55) hide show
  1. checksums.yaml +4 -4
  2. data/README.md +22 -5
  3. data/app/controllers/llmemory/dashboard/application_controller.rb +12 -0
  4. data/app/views/llmemory/dashboard/short_term/show.html.erb +1 -1
  5. data/lib/llmemory/active_record_helpers.rb +21 -0
  6. data/lib/llmemory/cli/commands/mcp.rb +5 -3
  7. data/lib/llmemory/configuration.rb +41 -6
  8. data/lib/llmemory/crypto/field_helpers.rb +5 -0
  9. data/lib/llmemory/extractors/entity_relation_extractor.rb +9 -2
  10. data/lib/llmemory/extractors/fact_extractor.rb +67 -1
  11. data/lib/llmemory/forget_log.rb +7 -3
  12. data/lib/llmemory/llm/anthropic.rb +4 -6
  13. data/lib/llmemory/llm/http_client.rb +55 -0
  14. data/lib/llmemory/llm/openai.rb +5 -7
  15. data/lib/llmemory/llm/tracking_client.rb +22 -12
  16. data/lib/llmemory/llm/usage.rb +4 -0
  17. data/lib/llmemory/llm/usage_ledger.rb +25 -24
  18. data/lib/llmemory/long_term/episodic/memory.rb +5 -2
  19. data/lib/llmemory/long_term/episodic/storages/active_record_storage.rb +3 -1
  20. data/lib/llmemory/long_term/file_based/memory.rb +9 -4
  21. data/lib/llmemory/long_term/file_based/storage.rb +15 -2
  22. data/lib/llmemory/long_term/file_based/storages/active_record_storage.rb +3 -1
  23. data/lib/llmemory/long_term/graph_based/memory.rb +39 -5
  24. data/lib/llmemory/long_term/graph_based/storage.rb +15 -2
  25. data/lib/llmemory/long_term/procedural/memory.rb +5 -2
  26. data/lib/llmemory/long_term/procedural/storages/active_record_storage.rb +3 -1
  27. data/lib/llmemory/maintenance/consolidator.rb +6 -2
  28. data/lib/llmemory/mcp/authentication.rb +5 -3
  29. data/lib/llmemory/mcp/server.rb +21 -6
  30. data/lib/llmemory/mcp/store_helpers.rb +35 -0
  31. data/lib/llmemory/mcp/tools/memory_save.rb +15 -18
  32. data/lib/llmemory/mcp/tools/memory_search.rb +16 -19
  33. data/lib/llmemory/memory.rb +29 -11
  34. data/lib/llmemory/memory_module.rb +6 -2
  35. data/lib/llmemory/provenance.rb +23 -1
  36. data/lib/llmemory/retrieval/feedback_store.rb +11 -5
  37. data/lib/llmemory/retrieval/multi_source.rb +26 -0
  38. data/lib/llmemory/retrieval/temporal_ranker.rb +4 -2
  39. data/lib/llmemory/retrieval.rb +1 -0
  40. data/lib/llmemory/short_term/pruner.rb +16 -4
  41. data/lib/llmemory/short_term/session_lifecycle.rb +18 -9
  42. data/lib/llmemory/short_term/stores/active_record_store.rb +37 -8
  43. data/lib/llmemory/short_term/stores/base.rb +9 -0
  44. data/lib/llmemory/short_term/stores/key_codec.rb +33 -0
  45. data/lib/llmemory/short_term/stores/memory_store.rb +38 -7
  46. data/lib/llmemory/short_term/stores/postgres_store.rb +81 -33
  47. data/lib/llmemory/short_term/stores/redis_store.rb +85 -9
  48. data/lib/llmemory/short_term/stores.rb +17 -2
  49. data/lib/llmemory/vector_store/active_record_store.rb +23 -8
  50. data/lib/llmemory/vector_store/base.rb +4 -0
  51. data/lib/llmemory/vector_store/memory_store.rb +9 -0
  52. data/lib/llmemory/vector_store/openai_embeddings.rb +6 -7
  53. data/lib/llmemory/version.rb +1 -1
  54. data/lib/llmemory/working_memory.rb +10 -7
  55. metadata +6 -15
checksums.yaml CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: bc40425c3adccdf3b900f4421cb642feb9e816df10934fe5fe360f2bb2a16706
4
- data.tar.gz: b1267ceae049b4c69655987b15fda8517616bde9f2048301bdc030f4d92b1259
3
+ metadata.gz: ecc37d3f00fe0f347b1cf37d17665a84b92fa8a0effdf6f238178651f44f58a9
4
+ data.tar.gz: 35dc741e76e12814e6abae4ede8cace697802359892cbb9c419a5e224c5b16f9
5
5
  SHA512:
6
- metadata.gz: f81850bc29cdeeec37d1d590442376bba2e473e895d0dda7f5be2c0c41d86b7e9bac0b0e2e1e731895a6ae978da78ca68be038ae5d9976bfcd4a1ec5702adec6
7
- data.tar.gz: 12d5f1736a795336d3ef951b5fde2ef1e89b1d0a193e00d00fc2de88e5337beee309198dd016ed7e4980598a85c342387713652b66e19fadb2455c853a7a5d94
6
+ metadata.gz: 7161647ed2bd7ac8b48a4695cd915fb26e80aee867238f8b40a70d8d2b014af72f3206c0b738738d2429a17760927b486e101ce3e82dae09efa59469bb8b54c8
7
+ data.tar.gz: 310bc288d0e6570280dfd78cfe8f002bb2e071275baf47521649ab2ba56d40ce182cdce2b93cc0bc21ba3c2ccebce7122c4e99c4e481a004790fb030634f1a43
data/README.md CHANGED
@@ -649,6 +649,19 @@ end
649
649
 
650
650
  The dashboard uses your existing `Llmemory.configuration` (short-term store, long-term store/type, etc.) and does not add any gem dependency; it only runs when Rails is present and you require `llmemory/dashboard`.
651
651
 
652
+ **Dashboard authentication**
653
+
654
+ By default, the dashboard requires HTTP Basic auth in non-development Rails environments (`dashboard_require_auth?` is `true` outside `Rails.env.development?`). Configure credentials in `Llmemory.configure`:
655
+
656
+ ```ruby
657
+ Llmemory.configure do |config|
658
+ config.dashboard_auth = { username: ENV["LLMEMORY_DASHBOARD_USER"], password: ENV["LLMEMORY_DASHBOARD_PASSWORD"] }
659
+ config.dashboard_require_auth = true # set false to disable (e.g. local dev behind a VPN)
660
+ end
661
+ ```
662
+
663
+ Requests without valid credentials receive `401 Unauthorized`.
664
+
652
665
  ## MCP Server (Model Context Protocol)
653
666
 
654
667
  llmemory includes an MCP server that allows LLM agents (like Claude Code) to interact directly with the memory system. This gives agents "agency" over their own memory—they can search, save, and retrieve memories autonomously.
@@ -685,16 +698,18 @@ llmemory mcp serve --name my-memory
685
698
  **HTTP mode** (for remote access or web integrations):
686
699
 
687
700
  ```bash
688
- # Start HTTP server on default port 3100
701
+ # Start HTTP server on default port 3100 (binds to 127.0.0.1 by default)
689
702
  llmemory mcp serve --http
690
703
 
691
704
  # Custom port and host
692
705
  llmemory mcp serve --http --port 8080 --host 127.0.0.1
693
706
 
694
- # With authentication (recommended for HTTP/HTTPS)
695
- MCP_TOKEN=your-secret-token llmemory mcp serve --http
707
+ # With authentication (required when binding outside loopback)
708
+ MCP_TOKEN=your-secret-token llmemory mcp serve --http --host 0.0.0.0
696
709
  ```
697
710
 
711
+ When the HTTP server listens on a non-loopback address (`0.0.0.0`, a public IP, etc.), **`MCP_TOKEN` is required** — the server refuses to start without it. Loopback binds (`127.0.0.1`, `localhost`) allow unauthenticated access for local development; set `MCP_TOKEN` anyway if other processes on the machine should not reach the server.
712
+
698
713
  **HTTPS mode** (secure remote access):
699
714
 
700
715
  ```bash
@@ -762,7 +777,7 @@ Or with the standalone executable:
762
777
 
763
778
  | Variable | Description |
764
779
  |----------|-------------|
765
- | `MCP_TOKEN` | Token for HTTP authentication (if set, requests must include valid token) |
780
+ | `MCP_TOKEN` | Token for HTTP authentication (**required** when the HTTP server binds outside loopback; optional on `127.0.0.1`) |
766
781
  | `LLMEMORY_DEBUG` | Set to `1` to enable debug output on stderr |
767
782
  | `OPENAI_API_KEY` | API key for LLM/embeddings |
768
783
  | `REDIS_URL` | Redis URL for short-term store |
@@ -770,7 +785,9 @@ Or with the standalone executable:
770
785
 
771
786
  ### HTTP Authentication
772
787
 
773
- When `MCP_TOKEN` is set, the HTTP server requires authentication. Requests must include the token via:
788
+ When binding to a non-loopback host, the HTTP server **requires** `MCP_TOKEN` at startup. On loopback, authentication is optional but recommended if untrusted local processes may connect.
789
+
790
+ When a token is configured (or mandatory), requests must include it via:
774
791
 
775
792
  - **Authorization header**: `Authorization: Bearer <token>` or `Authorization: <token>`
776
793
  - **Query parameter**: `?token=<token>`
@@ -3,12 +3,24 @@
3
3
  module Llmemory
4
4
  module Dashboard
5
5
  class ApplicationController < ActionController::Base
6
+ before_action :ensure_dashboard_authorized!
7
+
6
8
  helper_method :short_term_store, :file_based_storage, :graph_based_storage,
7
9
  :episodic_storage, :procedural_storage, :forget_log,
8
10
  :file_based?, :graph_based?
9
11
 
10
12
  protected
11
13
 
14
+ def ensure_dashboard_authorized!
15
+ return unless Llmemory.configuration.dashboard_require_auth?
16
+ auth = Llmemory.configuration.dashboard_auth
17
+ return if auth.respond_to?(:call) && auth.call(self)
18
+
19
+ head :forbidden
20
+ end
21
+
22
+ protected
23
+
12
24
  def short_term_store
13
25
  @short_term_store ||= build_short_term_store
14
26
  end
@@ -18,7 +18,7 @@
18
18
  <% content = (m[:content] || m["content"]).to_s %>
19
19
  <div class="message <%= role %>">
20
20
  <strong><%= role %>:</strong>
21
- <%= content.gsub(/\n/, "<br>").html_safe %>
21
+ <%= simple_format(content) %>
22
22
  </div>
23
23
  <% end %>
24
24
  <% end %>
@@ -0,0 +1,21 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Llmemory
4
+ module ActiveRecordHelpers
5
+ UNIQUE_RETRY_LIMIT = 3
6
+
7
+ private
8
+
9
+ def with_unique_retry(limit: UNIQUE_RETRY_LIMIT)
10
+ attempts = 0
11
+ begin
12
+ yield
13
+ rescue ActiveRecord::RecordNotUnique
14
+ attempts += 1
15
+ raise if attempts >= limit
16
+
17
+ retry
18
+ end
19
+ end
20
+ end
21
+ end
@@ -10,7 +10,7 @@ module Llmemory
10
10
  @server_name = "llmemory"
11
11
  @http_mode = false
12
12
  @port = 3100
13
- @host = "0.0.0.0"
13
+ @host = "127.0.0.1"
14
14
  @ssl_cert = nil
15
15
  @ssl_key = nil
16
16
 
@@ -18,7 +18,7 @@ module Llmemory
18
18
  parser.on("--name NAME", "Server name (default: llmemory)") { |v| @server_name = v }
19
19
  parser.on("--http", "Run as HTTP server instead of stdio") { @http_mode = true }
20
20
  parser.on("--port PORT", Integer, "HTTP port (default: 3100)") { |v| @port = v }
21
- parser.on("--host HOST", "HTTP host (default: 0.0.0.0)") { |v| @host = v }
21
+ parser.on("--host HOST", "HTTP host (default: 127.0.0.1)") { |v| @host = v }
22
22
  parser.on("--ssl-cert FILE", "SSL certificate file for HTTPS") { |v| @ssl_cert = v }
23
23
  parser.on("--ssl-key FILE", "SSL private key file for HTTPS") { |v| @ssl_key = v }
24
24
  parser.on("-h", "--help", "Show this help") do
@@ -51,7 +51,7 @@ module Llmemory
51
51
  puts " --name NAME Server name (default: llmemory)"
52
52
  puts " --http Run as HTTP server instead of stdio"
53
53
  puts " --port PORT HTTP port (default: 3100)"
54
- puts " --host HOST HTTP host (default: 0.0.0.0)"
54
+ puts " --host HOST HTTP host (default: 127.0.0.1)"
55
55
  puts " --ssl-cert FILE SSL certificate file for HTTPS"
56
56
  puts " --ssl-key FILE SSL private key file for HTTPS"
57
57
  puts ""
@@ -99,6 +99,8 @@ module Llmemory
99
99
  exit 1
100
100
  end
101
101
 
102
+ Llmemory.configure { |c| c.shared_memory_stores = true }
103
+
102
104
  server = Llmemory::MCP::Server.new(name: @server_name)
103
105
 
104
106
  if @http_mode
@@ -6,6 +6,10 @@ module Llmemory
6
6
  :llm_api_key,
7
7
  :llm_model,
8
8
  :llm_base_url,
9
+ :llm_timeout_seconds,
10
+ :llm_open_timeout_seconds,
11
+ :llm_http_retries,
12
+ :redis_session_ttl_override,
9
13
  :short_term_store,
10
14
  :redis_url,
11
15
  :long_term_type,
@@ -50,13 +54,20 @@ module Llmemory
50
54
  :ttl_procedural_days,
51
55
  :skill_mining_enabled,
52
56
  :encryption_enabled,
53
- :encryption_key
57
+ :encryption_key,
58
+ :shared_memory_stores,
59
+ :dashboard_auth,
60
+ :dashboard_require_auth
54
61
 
55
62
  def initialize
56
63
  @llm_provider = :openai
57
64
  @llm_api_key = ENV["OPENAI_API_KEY"]
58
65
  @llm_model = nil # falls back to the active provider's DEFAULT_MODEL
59
66
  @llm_base_url = nil
67
+ @llm_timeout_seconds = 60
68
+ @llm_open_timeout_seconds = 10
69
+ @llm_http_retries = 2
70
+ @redis_session_ttl_override = nil
60
71
  @short_term_store = :memory
61
72
  @redis_url = ENV["REDIS_URL"] || "redis://localhost:6379/0"
62
73
  @long_term_type = :file_based
@@ -102,6 +113,21 @@ module Llmemory
102
113
  @message_sanitizer_enabled = false
103
114
  @encryption_enabled = false
104
115
  @encryption_key = ENV["LLMEMORY_ENCRYPTION_KEY"]
116
+ @shared_memory_stores = false
117
+ @dashboard_auth = nil
118
+ @dashboard_require_auth = nil
119
+ end
120
+
121
+ def dashboard_require_auth?
122
+ return @dashboard_require_auth unless @dashboard_require_auth.nil?
123
+
124
+ defined?(Rails) && Rails.respond_to?(:env) && Rails.env.development? ? false : true
125
+ end
126
+
127
+ def redis_session_ttl_seconds
128
+ return @redis_session_ttl_override if @redis_session_ttl_override
129
+
130
+ session_prune_after_days.to_i * 86_400
105
131
  end
106
132
  end
107
133
 
@@ -116,20 +142,29 @@ module Llmemory
116
142
 
117
143
  def reset_configuration!
118
144
  @configuration = Configuration.new
145
+ ShortTerm::Stores.reset_shared_singletons! if defined?(ShortTerm::Stores)
146
+ if defined?(LongTerm::FileBased::Storages)
147
+ LongTerm::FileBased::Storages.reset_shared_singletons!
148
+ end
149
+ if defined?(LongTerm::GraphBased::Storages)
150
+ LongTerm::GraphBased::Storages.reset_shared_singletons!
151
+ end
119
152
  end
120
153
 
121
154
  # Builds a Crypto::Cipher when encryption is enabled and a key is present;
122
155
  # otherwise returns Crypto::NullCipher. An explicit non-empty instance key
123
156
  # enables encryption even when the global flag is off.
124
157
  def build_cipher(key = nil)
125
- explicit_key = !key.nil? && !key.to_s.empty?
126
158
  resolved = key.nil? ? configuration.encryption_key : key
127
- enabled = configuration.encryption_enabled || explicit_key
128
- if enabled && !resolved.to_s.empty?
129
- require_relative "crypto/cipher"
159
+ enabled = configuration.encryption_enabled || (!key.nil? && !key.to_s.empty?)
160
+ if enabled && resolved.to_s.empty?
161
+ raise ConfigurationError, "encryption_key cannot be empty when encryption is enabled"
162
+ end
163
+
164
+ require_relative "crypto/cipher"
165
+ if enabled
130
166
  Crypto::Cipher.new(resolved)
131
167
  else
132
- require_relative "crypto/cipher"
133
168
  Crypto::NullCipher.new
134
169
  end
135
170
  end
@@ -92,6 +92,11 @@ module Llmemory
92
92
  end
93
93
 
94
94
  str = data.to_s
95
+ if !cipher.enabled? && str.start_with?(Crypto::Cipher::MARKER)
96
+ raise DecryptionError,
97
+ "Data is encrypted but encryption is disabled; enable encryption with the correct encryption_key"
98
+ end
99
+
95
100
  json = cipher.enabled? && cipher.encrypted?(str) ? cipher.decrypt(str) : str
96
101
  JSON.parse(json, symbolize_names: true)
97
102
  end
@@ -53,8 +53,7 @@ module Llmemory
53
53
  end
54
54
 
55
55
  def extract(conversation_text)
56
- text = conversation_text.to_s.strip
57
- text = text[0, MAX_CONVERSATION_CHARS] + "\n[...]" if text.length > MAX_CONVERSATION_CHARS
56
+ text = truncate_conversation(conversation_text.to_s.strip)
58
57
  prompt = <<~PROMPT
59
58
  Infer entities and relations from this user-assistant conversation. Build a knowledge graph from what the user says, even when they don't state facts in formal language.
60
59
  - Entities: people, places, companies, concepts (type and name).
@@ -93,6 +92,14 @@ module Llmemory
93
92
 
94
93
  private
95
94
 
95
+ def truncate_conversation(text)
96
+ return text if text.length <= MAX_CONVERSATION_CHARS
97
+
98
+ head_len = (MAX_CONVERSATION_CHARS * 0.7).to_i
99
+ tail_len = MAX_CONVERSATION_CHARS - head_len
100
+ "#{text[0, head_len]}\n[...]\n#{text[-tail_len, tail_len]}"
101
+ end
102
+
96
103
  def parse_response(response)
97
104
  json = response.is_a?(Hash) ? response : extract_json(response)
98
105
  return { entities: [], relations: [] } unless json.is_a?(Hash)
@@ -43,6 +43,62 @@ module Llmemory
43
43
  end
44
44
 
45
45
  def classify_item(content)
46
+ classify_items([content])[content.to_s] || "general"
47
+ end
48
+
49
+ CLASSIFY_JSON_SCHEMA = {
50
+ name: "fact_classification",
51
+ schema: {
52
+ type: "object",
53
+ properties: {
54
+ items: {
55
+ type: "array",
56
+ items: {
57
+ type: "object",
58
+ properties: {
59
+ content: { type: "string" },
60
+ category: { type: "string" }
61
+ },
62
+ required: ["content", "category"],
63
+ additionalProperties: false
64
+ }
65
+ }
66
+ },
67
+ required: ["items"],
68
+ additionalProperties: false
69
+ }
70
+ }.freeze
71
+
72
+ def classify_items(contents)
73
+ list = Array(contents).map(&:to_s).reject(&:empty?)
74
+ return {} if list.empty?
75
+
76
+ if @llm.respond_to?(:invoke_with_json_schema)
77
+ begin
78
+ numbered = list.map.with_index { |c, i| "#{i + 1}. #{c}" }.join("\n")
79
+ prompt = <<~PROMPT
80
+ Classify each fact into ONE category. Use lowercase with underscores.
81
+ Examples: work_life, personal_life, preferences, general.
82
+
83
+ Facts:
84
+ #{numbered}
85
+
86
+ Return one category per fact in the same order.
87
+ PROMPT
88
+ parsed = @llm.invoke_with_json_schema(prompt.strip, CLASSIFY_JSON_SCHEMA)
89
+ mapped = map_classifications(parsed, list)
90
+ return mapped if mapped.size == list.size
91
+ rescue Llmemory::LLMError
92
+ # fall back to per-item classification
93
+ end
94
+ end
95
+
96
+ list.each_with_object({}) { |content, acc| acc[content] = classify_item_fallback(content) }
97
+ end
98
+
99
+ private
100
+
101
+ def classify_item_fallback(content)
46
102
  return "general" if content.to_s.strip.empty?
47
103
  prompt = <<~PROMPT
48
104
  Classify this fact into ONE category. Use lowercase with underscores. Examples: work_life, personal_life, preferences, general.
@@ -53,7 +109,17 @@ module Llmemory
53
109
  result.empty? ? "general" : result
54
110
  end
55
111
 
56
- private
112
+ def map_classifications(parsed, contents)
113
+ items = Array(parsed["items"] || parsed[:items])
114
+ by_content = items.each_with_object({}) do |row, acc|
115
+ content = (row["content"] || row[:content]).to_s
116
+ category = (row["category"] || row[:category]).to_s.strip.downcase.gsub(/\s+/, "_")
117
+ acc[content] = category.empty? ? "general" : category
118
+ end
119
+ contents.each_with_object({}) do |content, acc|
120
+ acc[content] = by_content[content] || classify_item_fallback(content)
121
+ end
122
+ end
57
123
 
58
124
  def parse_items_response(response)
59
125
  json = extract_json_array(response)
@@ -27,9 +27,13 @@ module Llmemory
27
27
  reason: reason,
28
28
  at: Time.now.iso8601
29
29
  }
30
- log = entries(user_id)
31
- log << entry
32
- @store.save(user_id, SESSION_KEY, { "entries" => log })
30
+ @store.update(user_id, SESSION_KEY) do |state|
31
+ state = symbolize(state || {})
32
+ log = state[:entries]
33
+ log = log.is_a?(Array) ? log.map { |e| symbolize(e) } : []
34
+ log << entry
35
+ state.merge(entries: log)
36
+ end
33
37
  entry
34
38
  end
35
39
 
@@ -3,10 +3,12 @@
3
3
  require "faraday"
4
4
  require "json"
5
5
  require_relative "base"
6
+ require_relative "http_client"
6
7
 
7
8
  module Llmemory
8
9
  module LLM
9
10
  class Anthropic < Base
11
+ include HttpClient
10
12
  DEFAULT_BASE_URL = "https://api.anthropic.com"
11
13
  DEFAULT_MODEL = "claude-sonnet-4-6"
12
14
 
@@ -21,7 +23,7 @@ module Llmemory
21
23
  result = nil
22
24
  payload = { provider: :anthropic, model: @model, prompt_chars: prompt.to_s.length }
23
25
  Llmemory::Instrumentation.instrument(:llm_invoke, payload) do
24
- response = connection.post("v1/messages") do |req|
26
+ response = post_with_resilience(connection, "v1/messages") do |req|
25
27
  req.body = {
26
28
  model: @model,
27
29
  max_tokens: 1024,
@@ -47,11 +49,7 @@ module Llmemory
47
49
  private
48
50
 
49
51
  def connection
50
- @connection ||= Faraday.new(url: @base_url) do |f|
51
- f.request :json
52
- f.response :json
53
- f.adapter Faraday.default_adapter
54
- end
52
+ @connection ||= build_faraday_connection(@base_url)
55
53
  end
56
54
  end
57
55
  end
@@ -0,0 +1,55 @@
1
+ # frozen_string_literal: true
2
+
3
+ require "json"
4
+
5
+ module Llmemory
6
+ module LLM
7
+ module HttpClient
8
+ RETRYABLE_STATUSES = [429, 500, 502, 503, 504].freeze
9
+
10
+ private
11
+
12
+ def http_timeout
13
+ (Llmemory.configuration.llm_timeout_seconds || 60).to_i
14
+ end
15
+
16
+ def http_open_timeout
17
+ (Llmemory.configuration.llm_open_timeout_seconds || 10).to_i
18
+ end
19
+
20
+ def build_faraday_connection(url)
21
+ Faraday.new(url: url) do |f|
22
+ f.request :json
23
+ f.response :json
24
+ f.options.timeout = http_timeout
25
+ f.options.open_timeout = http_open_timeout
26
+ f.adapter Faraday.default_adapter
27
+ end
28
+ end
29
+
30
+ def post_with_resilience(connection, path, &block)
31
+ attempts = 0
32
+ max_attempts = (Llmemory.configuration.llm_http_retries || 2).to_i + 1
33
+
34
+ loop do
35
+ begin
36
+ response = connection.post(path, &block)
37
+ if RETRYABLE_STATUSES.include?(response.status) && attempts < max_attempts - 1
38
+ attempts += 1
39
+ sleep(0.5 * (2**(attempts - 1)))
40
+ next
41
+ end
42
+ return response
43
+ rescue Faraday::TimeoutError, Faraday::ConnectionFailed => e
44
+ attempts += 1
45
+ raise Llmemory::LLMError, "HTTP request failed: #{e.message}" if attempts >= max_attempts
46
+
47
+ sleep(0.5 * (2**(attempts - 1)))
48
+ rescue Faraday::Error => e
49
+ raise Llmemory::LLMError, "HTTP request failed: #{e.message}"
50
+ end
51
+ end
52
+ end
53
+ end
54
+ end
55
+ end
@@ -3,10 +3,12 @@
3
3
  require "faraday"
4
4
  require "json"
5
5
  require_relative "base"
6
+ require_relative "http_client"
6
7
 
7
8
  module Llmemory
8
9
  module LLM
9
10
  class OpenAI < Base
11
+ include HttpClient
10
12
  DEFAULT_BASE_URL = "https://api.openai.com/v1"
11
13
  DEFAULT_MODEL = "gpt-4"
12
14
 
@@ -21,7 +23,7 @@ module Llmemory
21
23
  result = nil
22
24
  payload = { provider: :openai, model: @model, prompt_chars: prompt.to_s.length }
23
25
  Llmemory::Instrumentation.instrument(:llm_invoke, payload) do
24
- response = connection.post("chat/completions") do |req|
26
+ response = post_with_resilience(connection, "chat/completions") do |req|
25
27
  req.body = {
26
28
  model: @model,
27
29
  messages: [{ role: "user", content: prompt }],
@@ -63,7 +65,7 @@ module Llmemory
63
65
  parsed = nil
64
66
  instrument_payload = { provider: :openai, model: @model, prompt_chars: prompt.to_s.length }
65
67
  Llmemory::Instrumentation.instrument(:llm_invoke, instrument_payload) do
66
- response = connection.post("chat/completions") do |req|
68
+ response = post_with_resilience(connection, "chat/completions") do |req|
67
69
  req.body = payload.to_json
68
70
  req.headers["Content-Type"] = "application/json"
69
71
  req.headers["Authorization"] = "Bearer #{@api_key}"
@@ -88,11 +90,7 @@ module Llmemory
88
90
  private
89
91
 
90
92
  def connection
91
- @connection ||= Faraday.new(url: @base_url) do |f|
92
- f.request :json
93
- f.response :json
94
- f.adapter Faraday.default_adapter
95
- end
93
+ @connection ||= build_faraday_connection(@base_url)
96
94
  end
97
95
  end
98
96
  end
@@ -15,22 +15,20 @@ module Llmemory
15
15
 
16
16
  def invoke(prompt)
17
17
  response = inner_client.invoke(prompt)
18
- usage = if response.respond_to?(:usage)
19
- response.usage
20
- elsif inner_client.respond_to?(:last_usage)
21
- inner_client.last_usage
22
- else
23
- Usage.zero
24
- end
25
- UsageRecorder.record(user_id: @user_id, usage: usage, operation: :invoke, store: @store)
18
+ record_usage_from_client(:invoke)
26
19
  response
27
20
  end
28
21
 
29
22
  def invoke_with_json_schema(prompt, json_schema)
30
- result = inner_client.invoke_with_json_schema(prompt, json_schema)
31
- usage = inner_client.respond_to?(:last_usage) ? inner_client.last_usage : Usage.zero
32
- UsageRecorder.record(user_id: @user_id, usage: usage, operation: :invoke, store: @store)
33
- result
23
+ return nil unless structured_output_supported?
24
+
25
+ result = nil
26
+ begin
27
+ result = inner_client.invoke_with_json_schema(prompt, json_schema)
28
+ result
29
+ ensure
30
+ record_usage_from_client(:invoke) if structured_output_supported?
31
+ end
34
32
  end
35
33
 
36
34
  def last_usage
@@ -56,6 +54,18 @@ module Llmemory
56
54
  def inner_client
57
55
  @inner_client ||= @inner || Llmemory::LLM.client(api_key: @api_key)
58
56
  end
57
+
58
+ def structured_output_supported?
59
+ inner_client.class != Llmemory::LLM::Base &&
60
+ inner_client.method(:invoke_with_json_schema).owner != Llmemory::LLM::Base
61
+ end
62
+
63
+ def record_usage_from_client(operation)
64
+ usage = inner_client.respond_to?(:last_usage) ? inner_client.last_usage : Usage.zero
65
+ return if usage.zero?
66
+
67
+ UsageRecorder.record(user_id: @user_id, usage: usage, operation: operation, store: @store)
68
+ end
59
69
  end
60
70
  end
61
71
  end
@@ -26,6 +26,10 @@ module Llmemory
26
26
  def to_h
27
27
  { input_tokens: @input_tokens, output_tokens: @output_tokens, total_tokens: @total_tokens }
28
28
  end
29
+
30
+ def zero?
31
+ @input_tokens.zero? && @output_tokens.zero? && @total_tokens.zero?
32
+ end
29
33
  end
30
34
  end
31
35
  end
@@ -15,31 +15,32 @@ module Llmemory
15
15
  end
16
16
 
17
17
  def record(user_id, usage, operation:)
18
- state = load_raw(user_id)
19
- case operation.to_sym
20
- when :invoke
21
- bucket = symbolize_bucket(state[:invoke] || state["invoke"])
22
- state = state.merge(
23
- invoke: {
24
- input_tokens: bucket[:input_tokens] + usage.input_tokens,
25
- output_tokens: bucket[:output_tokens] + usage.output_tokens,
26
- total_tokens: bucket[:total_tokens] + usage.total_tokens,
27
- calls: bucket[:calls] + 1
28
- }
29
- )
30
- when :embed
31
- bucket = symbolize_bucket(state[:embed] || state["embed"], embed: true)
32
- state = state.merge(
33
- embed: {
34
- total_tokens: bucket[:total_tokens] + usage.total_tokens,
35
- calls: bucket[:calls] + 1
36
- }
37
- )
38
- else
39
- return totals(user_id)
18
+ @store.update(user_id, SESSION_KEY) do |state|
19
+ state = normalize(state || default_state)
20
+ case operation.to_sym
21
+ when :invoke
22
+ bucket = state[:invoke]
23
+ state = state.merge(
24
+ invoke: {
25
+ input_tokens: bucket[:input_tokens] + usage.input_tokens,
26
+ output_tokens: bucket[:output_tokens] + usage.output_tokens,
27
+ total_tokens: bucket[:total_tokens] + usage.total_tokens,
28
+ calls: bucket[:calls] + 1
29
+ }
30
+ )
31
+ when :embed
32
+ bucket = state[:embed]
33
+ state = state.merge(
34
+ embed: {
35
+ total_tokens: bucket[:total_tokens] + usage.total_tokens,
36
+ calls: bucket[:calls] + 1
37
+ }
38
+ )
39
+ else
40
+ next state
41
+ end
42
+ state.merge(updated_at: Time.now.iso8601)
40
43
  end
41
- state[:updated_at] = Time.now.iso8601
42
- @store.save(user_id, SESSION_KEY, stringify(state))
43
44
  totals(user_id)
44
45
  end
45
46