flowent 0.1.5 → 0.2.1

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 (67) hide show
  1. package/backend/pyproject.toml +31 -5
  2. package/backend/src/flowent/agent.py +107 -37
  3. package/backend/src/flowent/compact.py +35 -14
  4. package/backend/src/flowent/llm.py +198 -12
  5. package/backend/src/flowent/main.py +260 -59
  6. package/backend/src/flowent/static/assets/index-CRSV2xu1.css +2 -0
  7. package/backend/src/flowent/static/assets/index-DUYj6rgD.js +82 -0
  8. package/backend/src/flowent/static/index.html +2 -2
  9. package/backend/src/flowent/storage.py +135 -3
  10. package/backend/src/flowent/usage.py +315 -0
  11. package/backend/uv.lock +971 -3
  12. package/dist/frontend/assets/index-CRSV2xu1.css +2 -0
  13. package/dist/frontend/assets/index-DUYj6rgD.js +82 -0
  14. package/dist/frontend/index.html +2 -2
  15. package/package.json +24 -3
  16. package/backend/src/flowent/__pycache__/__init__.cpython-313.pyc +0 -0
  17. package/backend/src/flowent/__pycache__/_version.cpython-313.pyc +0 -0
  18. package/backend/src/flowent/__pycache__/agent.cpython-313.pyc +0 -0
  19. package/backend/src/flowent/__pycache__/approval.cpython-313.pyc +0 -0
  20. package/backend/src/flowent/__pycache__/channels.cpython-313.pyc +0 -0
  21. package/backend/src/flowent/__pycache__/cli.cpython-313.pyc +0 -0
  22. package/backend/src/flowent/__pycache__/compact.cpython-313.pyc +0 -0
  23. package/backend/src/flowent/__pycache__/context.cpython-313.pyc +0 -0
  24. package/backend/src/flowent/__pycache__/llm.cpython-313.pyc +0 -0
  25. package/backend/src/flowent/__pycache__/logging.cpython-313.pyc +0 -0
  26. package/backend/src/flowent/__pycache__/main.cpython-313.pyc +0 -0
  27. package/backend/src/flowent/__pycache__/mcp.cpython-313.pyc +0 -0
  28. package/backend/src/flowent/__pycache__/mcp_import.cpython-313.pyc +0 -0
  29. package/backend/src/flowent/__pycache__/patch.cpython-313.pyc +0 -0
  30. package/backend/src/flowent/__pycache__/paths.cpython-313.pyc +0 -0
  31. package/backend/src/flowent/__pycache__/permissions.cpython-313.pyc +0 -0
  32. package/backend/src/flowent/__pycache__/sandbox.cpython-313.pyc +0 -0
  33. package/backend/src/flowent/__pycache__/skills.cpython-313.pyc +0 -0
  34. package/backend/src/flowent/__pycache__/storage.cpython-313.pyc +0 -0
  35. package/backend/src/flowent/__pycache__/tools.cpython-313.pyc +0 -0
  36. package/backend/src/flowent/static/assets/index-Cl20cARb.css +0 -2
  37. package/backend/src/flowent/static/assets/index-dsDDsEym.js +0 -81
  38. package/backend/tests/__pycache__/conftest.cpython-313-pytest-9.0.3.pyc +0 -0
  39. package/backend/tests/__pycache__/test_agent_tools.cpython-313-pytest-9.0.3.pyc +0 -0
  40. package/backend/tests/__pycache__/test_approval.cpython-313-pytest-9.0.3.pyc +0 -0
  41. package/backend/tests/__pycache__/test_channels.cpython-313-pytest-9.0.3.pyc +0 -0
  42. package/backend/tests/__pycache__/test_health.cpython-313-pytest-9.0.3.pyc +0 -0
  43. package/backend/tests/__pycache__/test_llm_providers.cpython-313-pytest-9.0.3.pyc +0 -0
  44. package/backend/tests/__pycache__/test_logging.cpython-313-pytest-9.0.3.pyc +0 -0
  45. package/backend/tests/__pycache__/test_mcp.cpython-313-pytest-9.0.3.pyc +0 -0
  46. package/backend/tests/__pycache__/test_patch.cpython-313-pytest-9.0.3.pyc +0 -0
  47. package/backend/tests/__pycache__/test_permissions.cpython-313-pytest-9.0.3.pyc +0 -0
  48. package/backend/tests/__pycache__/test_persistence.cpython-313-pytest-9.0.3.pyc +0 -0
  49. package/backend/tests/__pycache__/test_skills.cpython-313-pytest-9.0.3.pyc +0 -0
  50. package/backend/tests/__pycache__/test_startup_requirements.cpython-313-pytest-9.0.3.pyc +0 -0
  51. package/backend/tests/__pycache__/test_workspace_chat.cpython-313-pytest-9.0.3.pyc +0 -0
  52. package/backend/tests/conftest.py +0 -21
  53. package/backend/tests/test_agent_tools.py +0 -988
  54. package/backend/tests/test_approval.py +0 -283
  55. package/backend/tests/test_channels.py +0 -360
  56. package/backend/tests/test_health.py +0 -12
  57. package/backend/tests/test_llm_providers.py +0 -387
  58. package/backend/tests/test_logging.py +0 -212
  59. package/backend/tests/test_mcp.py +0 -788
  60. package/backend/tests/test_patch.py +0 -112
  61. package/backend/tests/test_permissions.py +0 -588
  62. package/backend/tests/test_persistence.py +0 -249
  63. package/backend/tests/test_skills.py +0 -462
  64. package/backend/tests/test_startup_requirements.py +0 -144
  65. package/backend/tests/test_workspace_chat.py +0 -2122
  66. package/dist/frontend/assets/index-Cl20cARb.css +0 -2
  67. package/dist/frontend/assets/index-dsDDsEym.js +0 -81
@@ -6,8 +6,8 @@
6
6
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
7
7
  <title>Flowent</title>
8
8
  <meta name="description" content="Flowent application" />
9
- <script type="module" crossorigin src="/assets/index-dsDDsEym.js"></script>
10
- <link rel="stylesheet" crossorigin href="/assets/index-Cl20cARb.css">
9
+ <script type="module" crossorigin src="/assets/index-DUYj6rgD.js"></script>
10
+ <link rel="stylesheet" crossorigin href="/assets/index-CRSV2xu1.css">
11
11
  </head>
12
12
  <body>
13
13
  <div id="root"></div>
@@ -7,6 +7,7 @@ from pydantic import BaseModel, ConfigDict, Field
7
7
 
8
8
  from flowent.llm import ChatMessage, ProviderFormat, ReasoningEffort
9
9
  from flowent.paths import data_directory
10
+ from flowent.usage import TokenUsageInfo
10
11
 
11
12
 
12
13
  class StoredTelegramSession(BaseModel):
@@ -172,6 +173,9 @@ class StoredMessage(BaseModel):
172
173
  )
173
174
  thinking: str = Field(default="", exclude_if=lambda value: value == "")
174
175
  tools: list[StoredToolItem] = Field(default_factory=list)
176
+ usage_info: TokenUsageInfo | None = Field(
177
+ default=None, exclude_if=lambda value: value is None
178
+ )
175
179
 
176
180
 
177
181
  class StoredCompactionCheckpoint(BaseModel):
@@ -193,12 +197,16 @@ class StoredState(BaseModel):
193
197
 
194
198
  active_run_event_index: int = 0
195
199
  active_run_id: str | None = None
200
+ is_compacting: bool = False
196
201
  mcp_servers: list[StoredMcpServer]
197
202
  messages: list[StoredMessage]
198
203
  providers: list[StoredProvider]
199
204
  settings: StoredSettings
200
205
  skills: list[StoredSkill]
201
206
  telegram_bot: StoredTelegramBot
207
+ usage_info: TokenUsageInfo | None = Field(
208
+ default=None, exclude_if=lambda value: value is None
209
+ )
202
210
  writable_paths: list[StoredWritablePath] = Field(default_factory=list)
203
211
 
204
212
 
@@ -261,18 +269,34 @@ class StateStore:
261
269
  StoredToolItem.model_validate(tool)
262
270
  for tool in json.loads(row["tools"] or "[]")
263
271
  ],
272
+ usage_info=TokenUsageInfo.model_validate_json(row["usage_info"])
273
+ if row["usage_info"]
274
+ else None,
264
275
  )
265
276
  for row in connection.execute(
266
277
  """
267
- SELECT id, author, content, tools, thinking, groups, status
278
+ SELECT id, author, content, tools, thinking, groups, status, usage_info
268
279
  FROM messages
269
280
  ORDER BY position, id
270
281
  """
271
282
  )
272
283
  ]
284
+ usage_row = connection.execute(
285
+ """
286
+ SELECT is_compacting, usage_info
287
+ FROM workspace_context
288
+ WHERE id = 1
289
+ """
290
+ ).fetchone()
291
+ usage_info = (
292
+ TokenUsageInfo.model_validate_json(usage_row["usage_info"])
293
+ if usage_row and usage_row["usage_info"]
294
+ else None
295
+ )
273
296
 
274
297
  return StoredState(
275
298
  mcp_servers=mcp_servers,
299
+ is_compacting=bool(usage_row["is_compacting"]) if usage_row else False,
276
300
  messages=messages,
277
301
  providers=providers,
278
302
  settings=StoredSettings(
@@ -287,6 +311,7 @@ class StateStore:
287
311
  ),
288
312
  skills=[],
289
313
  telegram_bot=telegram_bot,
314
+ usage_info=usage_info,
290
315
  writable_paths=writable_paths,
291
316
  )
292
317
 
@@ -600,6 +625,26 @@ class StateStore:
600
625
  def save_messages(self, messages: list[StoredMessage]) -> list[StoredMessage]:
601
626
  with self.connect() as connection:
602
627
  connection.execute("DELETE FROM messages")
628
+ if messages:
629
+ latest_usage_info = next(
630
+ (
631
+ message.usage_info
632
+ for message in reversed(messages)
633
+ if message.usage_info is not None
634
+ ),
635
+ None,
636
+ )
637
+ if latest_usage_info is not None:
638
+ connection.execute(
639
+ """
640
+ INSERT INTO workspace_context (id, usage_info)
641
+ VALUES (1, ?)
642
+ ON CONFLICT(id) DO UPDATE SET
643
+ usage_info = excluded.usage_info,
644
+ updated_at = unixepoch()
645
+ """,
646
+ (latest_usage_info.model_dump_json(),),
647
+ )
603
648
  connection.executemany(
604
649
  """
605
650
  INSERT INTO messages (
@@ -610,9 +655,10 @@ class StateStore:
610
655
  thinking,
611
656
  groups,
612
657
  status,
658
+ usage_info,
613
659
  position
614
660
  )
615
- VALUES (?, ?, ?, ?, ?, ?, ?, ?)
661
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
616
662
  """,
617
663
  [
618
664
  (
@@ -634,6 +680,9 @@ class StateStore:
634
680
  ensure_ascii=False,
635
681
  ),
636
682
  message.status,
683
+ message.usage_info.model_dump_json()
684
+ if message.usage_info
685
+ else None,
637
686
  position,
638
687
  )
639
688
  for position, message in enumerate(messages)
@@ -665,9 +714,10 @@ class StateStore:
665
714
  thinking,
666
715
  groups,
667
716
  status,
717
+ usage_info,
668
718
  position
669
719
  )
670
- VALUES (?, ?, ?, ?, ?, ?, ?, ?)
720
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
671
721
  ON CONFLICT(id) DO UPDATE SET
672
722
  author = excluded.author,
673
723
  content = excluded.content,
@@ -675,6 +725,7 @@ class StateStore:
675
725
  thinking = excluded.thinking,
676
726
  groups = excluded.groups,
677
727
  status = excluded.status,
728
+ usage_info = excluded.usage_info,
678
729
  position = excluded.position
679
730
  """,
680
731
  (
@@ -693,11 +744,52 @@ class StateStore:
693
744
  ensure_ascii=False,
694
745
  ),
695
746
  message.status,
747
+ message.usage_info.model_dump_json()
748
+ if message.usage_info
749
+ else None,
696
750
  position,
697
751
  ),
698
752
  )
753
+ if message.usage_info is not None:
754
+ connection.execute(
755
+ """
756
+ INSERT INTO workspace_context (id, usage_info)
757
+ VALUES (1, ?)
758
+ ON CONFLICT(id) DO UPDATE SET
759
+ usage_info = excluded.usage_info,
760
+ updated_at = unixepoch()
761
+ """,
762
+ (message.usage_info.model_dump_json(),),
763
+ )
699
764
  return message
700
765
 
766
+ def read_usage_info(self) -> TokenUsageInfo | None:
767
+ with self.connect() as connection:
768
+ row = connection.execute(
769
+ """
770
+ SELECT usage_info
771
+ FROM workspace_context
772
+ WHERE id = 1
773
+ """
774
+ ).fetchone()
775
+ if row is None or not row["usage_info"]:
776
+ return None
777
+ return TokenUsageInfo.model_validate_json(row["usage_info"])
778
+
779
+ def save_usage_info(self, usage_info: TokenUsageInfo) -> TokenUsageInfo:
780
+ with self.connect() as connection:
781
+ connection.execute(
782
+ """
783
+ INSERT INTO workspace_context (id, usage_info)
784
+ VALUES (1, ?)
785
+ ON CONFLICT(id) DO UPDATE SET
786
+ usage_info = excluded.usage_info,
787
+ updated_at = unixepoch()
788
+ """,
789
+ (usage_info.model_dump_json(),),
790
+ )
791
+ return usage_info
792
+
701
793
  def read_compacted_context(self) -> str:
702
794
  with self.connect() as connection:
703
795
  row = connection.execute(
@@ -718,12 +810,38 @@ class StateStore:
718
810
  ON CONFLICT(id) DO UPDATE SET
719
811
  compacted_summary = excluded.compacted_summary,
720
812
  active_compaction_id = NULL,
813
+ usage_info = NULL,
721
814
  updated_at = unixepoch()
722
815
  """,
723
816
  (summary,),
724
817
  )
725
818
  return summary
726
819
 
820
+ def read_is_compacting(self) -> bool:
821
+ with self.connect() as connection:
822
+ row = connection.execute(
823
+ """
824
+ SELECT is_compacting
825
+ FROM workspace_context
826
+ WHERE id = 1
827
+ """
828
+ ).fetchone()
829
+ return bool(row["is_compacting"]) if row else False
830
+
831
+ def save_is_compacting(self, is_compacting: bool) -> bool:
832
+ with self.connect() as connection:
833
+ connection.execute(
834
+ """
835
+ INSERT INTO workspace_context (id, is_compacting)
836
+ VALUES (1, ?)
837
+ ON CONFLICT(id) DO UPDATE SET
838
+ is_compacting = excluded.is_compacting,
839
+ updated_at = unixepoch()
840
+ """,
841
+ (int(is_compacting),),
842
+ )
843
+ return is_compacting
844
+
727
845
  def read_active_compaction_checkpoint(
728
846
  self,
729
847
  ) -> StoredCompactionCheckpoint | None:
@@ -1020,6 +1138,7 @@ class StateStore:
1020
1138
  author TEXT NOT NULL,
1021
1139
  content TEXT NOT NULL,
1022
1140
  status TEXT NOT NULL DEFAULT 'completed',
1141
+ usage_info TEXT,
1023
1142
  position INTEGER NOT NULL
1024
1143
  );
1025
1144
 
@@ -1027,6 +1146,8 @@ class StateStore:
1027
1146
  id INTEGER PRIMARY KEY CHECK (id = 1),
1028
1147
  compacted_summary TEXT NOT NULL DEFAULT '',
1029
1148
  active_compaction_id TEXT,
1149
+ is_compacting INTEGER NOT NULL DEFAULT 0,
1150
+ usage_info TEXT,
1030
1151
  updated_at INTEGER NOT NULL DEFAULT (unixepoch())
1031
1152
  );
1032
1153
 
@@ -1089,6 +1210,8 @@ class StateStore:
1089
1210
  connection.execute(
1090
1211
  "ALTER TABLE messages ADD COLUMN groups TEXT NOT NULL DEFAULT '[]'"
1091
1212
  )
1213
+ if "usage_info" not in columns:
1214
+ connection.execute("ALTER TABLE messages ADD COLUMN usage_info TEXT")
1092
1215
  settings_columns = {
1093
1216
  row["name"] for row in connection.execute("PRAGMA table_info(settings)")
1094
1217
  }
@@ -1109,3 +1232,12 @@ class StateStore:
1109
1232
  connection.execute(
1110
1233
  "ALTER TABLE workspace_context ADD COLUMN active_compaction_id TEXT"
1111
1234
  )
1235
+ if "usage_info" not in workspace_context_columns:
1236
+ connection.execute(
1237
+ "ALTER TABLE workspace_context ADD COLUMN usage_info TEXT"
1238
+ )
1239
+ if "is_compacting" not in workspace_context_columns:
1240
+ connection.execute(
1241
+ "ALTER TABLE workspace_context "
1242
+ "ADD COLUMN is_compacting INTEGER NOT NULL DEFAULT 0"
1243
+ )
@@ -0,0 +1,315 @@
1
+ import json
2
+ from collections.abc import Mapping, Sequence
3
+ from typing import Any
4
+
5
+ from pydantic import BaseModel, ConfigDict, Field
6
+
7
+ DEFAULT_MODEL_CONTEXT_WINDOW = 120_000
8
+
9
+ MODEL_CONTEXT_WINDOWS: dict[str, int] = {
10
+ "claude-3-7-sonnet-20250219": 200_000,
11
+ "claude-3-haiku-20240307": 200_000,
12
+ "claude-3-opus-20240229": 200_000,
13
+ "claude-4-opus-20250514": 200_000,
14
+ "claude-4-sonnet-20250514": 1_000_000,
15
+ "claude-haiku-4-5": 200_000,
16
+ "claude-haiku-4-5-20251001": 200_000,
17
+ "claude-opus-4-1": 200_000,
18
+ "claude-opus-4-1-20250805": 200_000,
19
+ "claude-opus-4-20250514": 200_000,
20
+ "claude-opus-4-5": 200_000,
21
+ "claude-opus-4-5-20251101": 200_000,
22
+ "claude-opus-4-6": 1_000_000,
23
+ "claude-opus-4-6-20260205": 1_000_000,
24
+ "claude-opus-4-7": 1_000_000,
25
+ "claude-opus-4-7-20260416": 1_000_000,
26
+ "claude-opus-4-8": 1_000_000,
27
+ "claude-sonnet-4-20250514": 1_000_000,
28
+ "claude-sonnet-4-5": 200_000,
29
+ "claude-sonnet-4-5-20250929": 200_000,
30
+ "claude-sonnet-4-5-20250929-v1:0": 200_000,
31
+ "claude-sonnet-4-6": 1_000_000,
32
+ "gemini-2.5-computer-use-preview-10-2025": 128_000,
33
+ "gemini-2.5-flash": 1_048_576,
34
+ "gemini-2.5-flash-image": 32_768,
35
+ "gemini-2.5-flash-lite": 1_048_576,
36
+ "gemini-2.5-flash-lite-preview-06-17": 1_048_576,
37
+ "gemini-2.5-flash-lite-preview-09-2025": 1_048_576,
38
+ "gemini-2.5-flash-native-audio-latest": 1_048_576,
39
+ "gemini-2.5-flash-native-audio-preview-09-2025": 1_048_576,
40
+ "gemini-2.5-flash-native-audio-preview-12-2025": 1_048_576,
41
+ "gemini-2.5-flash-preview-09-2025": 1_048_576,
42
+ "gemini-2.5-pro": 1_048_576,
43
+ "gemini-2.5-pro-preview-tts": 1_048_576,
44
+ "gemini-3-flash-preview": 1_048_576,
45
+ "gemini-3-pro-image-preview": 65_536,
46
+ "gemini-3-pro-preview": 1_048_576,
47
+ "gemini-3.1-flash-image-preview": 65_536,
48
+ "gemini-3.1-flash-lite": 1_048_576,
49
+ "gemini-3.1-flash-lite-preview": 1_048_576,
50
+ "gemini-3.1-flash-live-preview": 131_072,
51
+ "gemini-3.1-pro-preview": 1_048_576,
52
+ "gemini-3.1-pro-preview-customtools": 1_048_576,
53
+ "gemini-3.5-flash": 1_048_576,
54
+ "gpt-4.1": 1_047_576,
55
+ "gpt-4.1-2025-04-14": 1_047_576,
56
+ "gpt-4.1-mini": 1_047_576,
57
+ "gpt-4.1-mini-2025-04-14": 1_047_576,
58
+ "gpt-4.1-nano": 1_047_576,
59
+ "gpt-4.1-nano-2025-04-14": 1_047_576,
60
+ "gpt-5": 272_000,
61
+ "gpt-5-2025-08-07": 272_000,
62
+ "gpt-5-chat": 128_000,
63
+ "gpt-5-chat-latest": 128_000,
64
+ "gpt-5-codex": 272_000,
65
+ "gpt-5-mini": 272_000,
66
+ "gpt-5-mini-2025-08-07": 272_000,
67
+ "gpt-5-nano": 272_000,
68
+ "gpt-5-nano-2025-08-07": 272_000,
69
+ "gpt-5-pro": 128_000,
70
+ "gpt-5-pro-2025-10-06": 128_000,
71
+ "gpt-5-search-api": 272_000,
72
+ "gpt-5-search-api-2025-10-14": 272_000,
73
+ "gpt-5.1": 272_000,
74
+ "gpt-5.1-2025-11-13": 272_000,
75
+ "gpt-5.1-chat-latest": 128_000,
76
+ "gpt-5.1-codex": 272_000,
77
+ "gpt-5.1-codex-max": 272_000,
78
+ "gpt-5.1-codex-mini": 272_000,
79
+ "gpt-5.2": 272_000,
80
+ "gpt-5.2-2025-12-11": 272_000,
81
+ "gpt-5.2-chat-latest": 128_000,
82
+ "gpt-5.2-codex": 272_000,
83
+ "gpt-5.2-pro": 272_000,
84
+ "gpt-5.2-pro-2025-12-11": 272_000,
85
+ "gpt-5.3-chat-latest": 128_000,
86
+ "gpt-5.3-codex": 272_000,
87
+ "gpt-5.4": 1_050_000,
88
+ "gpt-5.4-2026-03-05": 1_050_000,
89
+ "gpt-5.4-mini": 272_000,
90
+ "gpt-5.4-mini-2026-03-17": 272_000,
91
+ "gpt-5.4-nano": 272_000,
92
+ "gpt-5.4-nano-2026-03-17": 272_000,
93
+ "gpt-5.4-pro": 1_050_000,
94
+ "gpt-5.4-pro-2026-03-05": 1_050_000,
95
+ "gpt-5.5": 1_050_000,
96
+ "gpt-5.5-2026-04-23": 1_050_000,
97
+ "gpt-5.5-pro": 1_050_000,
98
+ "gpt-5.5-pro-2026-04-23": 1_050_000,
99
+ "o3": 200_000,
100
+ "o3-2025-04-16": 200_000,
101
+ "o3-deep-research": 200_000,
102
+ "o3-deep-research-2025-06-26": 200_000,
103
+ "o3-mini": 200_000,
104
+ "o3-mini-2025-01-31": 200_000,
105
+ "o3-pro": 200_000,
106
+ "o3-pro-2025-06-10": 200_000,
107
+ "o4-mini": 200_000,
108
+ "o4-mini-2025-04-16": 200_000,
109
+ "o4-mini-deep-research": 200_000,
110
+ "o4-mini-deep-research-2025-06-26": 200_000,
111
+ }
112
+
113
+ MODEL_CONTEXT_WINDOW_NAMES = tuple(sorted(MODEL_CONTEXT_WINDOWS, key=len, reverse=True))
114
+
115
+
116
+ class TokenUsage(BaseModel):
117
+ model_config = ConfigDict(extra="forbid")
118
+
119
+ input_tokens: int = 0
120
+ cached_input_tokens: int = 0
121
+ output_tokens: int = 0
122
+ reasoning_output_tokens: int = 0
123
+ total_tokens: int = 0
124
+
125
+ def add(self, other: "TokenUsage") -> "TokenUsage":
126
+ return TokenUsage(
127
+ input_tokens=self.input_tokens + other.input_tokens,
128
+ cached_input_tokens=self.cached_input_tokens + other.cached_input_tokens,
129
+ output_tokens=self.output_tokens + other.output_tokens,
130
+ reasoning_output_tokens=self.reasoning_output_tokens
131
+ + other.reasoning_output_tokens,
132
+ total_tokens=self.total_tokens + other.total_tokens,
133
+ )
134
+
135
+
136
+ class TokenUsageInfo(BaseModel):
137
+ model_config = ConfigDict(extra="forbid")
138
+
139
+ total_token_usage: TokenUsage = Field(default_factory=TokenUsage)
140
+ last_token_usage: TokenUsage = Field(default_factory=TokenUsage)
141
+ model_context_window: int | None = None
142
+
143
+
144
+ def current_model_context_window(model_name: str | None = None) -> int:
145
+ return model_context_window_for(model_name)
146
+
147
+
148
+ def model_context_window_for(model_name: str | None = None) -> int:
149
+ candidates = normalized_model_name_candidates(model_name)
150
+ for candidate in candidates:
151
+ context_window = MODEL_CONTEXT_WINDOWS.get(candidate)
152
+ if context_window is not None:
153
+ return context_window
154
+ for candidate in candidates:
155
+ for known_model in MODEL_CONTEXT_WINDOW_NAMES:
156
+ if is_model_context_window_prefix_match(candidate, known_model):
157
+ return MODEL_CONTEXT_WINDOWS[known_model]
158
+ return DEFAULT_MODEL_CONTEXT_WINDOW
159
+
160
+
161
+ def normalized_model_name_candidates(model_name: str | None) -> tuple[str, ...]:
162
+ if model_name is None:
163
+ return ()
164
+ normalized = model_name.strip().lower()
165
+ if not normalized:
166
+ return ()
167
+ candidates = [normalized]
168
+ if "/" in normalized:
169
+ candidates.append(normalized.rsplit("/", 1)[-1])
170
+ return tuple(dict.fromkeys(candidates))
171
+
172
+
173
+ def is_model_context_window_prefix_match(candidate: str, known_model: str) -> bool:
174
+ if candidate == known_model:
175
+ return True
176
+ if not candidate.startswith(known_model):
177
+ return False
178
+ return candidate[len(known_model)] in {"-", ".", ":", "/"}
179
+
180
+
181
+ def append_token_usage(
182
+ usage_info: TokenUsageInfo | None,
183
+ usage: TokenUsage,
184
+ *,
185
+ model_context_window: int | None = None,
186
+ ) -> TokenUsageInfo:
187
+ info = usage_info or TokenUsageInfo(model_context_window=model_context_window)
188
+ return TokenUsageInfo(
189
+ total_token_usage=info.total_token_usage.add(usage),
190
+ last_token_usage=usage,
191
+ model_context_window=model_context_window or info.model_context_window,
192
+ )
193
+
194
+
195
+ def recompute_context_usage(
196
+ usage_info: TokenUsageInfo | None,
197
+ active_context_tokens: int,
198
+ *,
199
+ model_context_window: int | None = None,
200
+ ) -> TokenUsageInfo:
201
+ info = usage_info or TokenUsageInfo(model_context_window=model_context_window)
202
+ return TokenUsageInfo(
203
+ total_token_usage=info.total_token_usage,
204
+ last_token_usage=TokenUsage(total_tokens=max(0, active_context_tokens)),
205
+ model_context_window=model_context_window or info.model_context_window,
206
+ )
207
+
208
+
209
+ def token_usage_from_response(response: Any) -> TokenUsage | None:
210
+ usage = value_at(response, "usage")
211
+ if usage is None:
212
+ return None
213
+
214
+ input_tokens = first_int_value(
215
+ value_at(usage, "input_tokens"),
216
+ value_at(usage, "prompt_tokens"),
217
+ )
218
+ output_tokens = first_int_value(
219
+ value_at(usage, "output_tokens"),
220
+ value_at(usage, "completion_tokens"),
221
+ )
222
+ total_tokens = first_int_value(value_at(usage, "total_tokens"))
223
+ cached_input_tokens = first_int_value(
224
+ value_at(usage, "cached_input_tokens"),
225
+ value_at(usage, "cache_read_input_tokens"),
226
+ value_at(usage, "cached_tokens"),
227
+ nested_value_at(usage, "prompt_tokens_details", "cached_tokens"),
228
+ nested_value_at(usage, "input_tokens_details", "cached_tokens"),
229
+ nested_value_at(usage, "cache_read", "input_tokens"),
230
+ )
231
+ reasoning_output_tokens = first_int_value(
232
+ value_at(usage, "reasoning_output_tokens"),
233
+ nested_value_at(usage, "completion_tokens_details", "reasoning_tokens"),
234
+ nested_value_at(usage, "output_tokens_details", "reasoning_tokens"),
235
+ )
236
+
237
+ if total_tokens is None:
238
+ total_tokens = (input_tokens or 0) + (output_tokens or 0)
239
+
240
+ return TokenUsage(
241
+ input_tokens=input_tokens or 0,
242
+ cached_input_tokens=cached_input_tokens or 0,
243
+ output_tokens=output_tokens or 0,
244
+ reasoning_output_tokens=reasoning_output_tokens or 0,
245
+ total_tokens=total_tokens,
246
+ )
247
+
248
+
249
+ def estimated_token_usage_for_messages(
250
+ messages: Sequence[Mapping[str, object]],
251
+ *,
252
+ output_content: str = "",
253
+ ) -> TokenUsage:
254
+ total_tokens = sum(estimate_mapping_message_tokens(message) for message in messages)
255
+ output_tokens = approximate_token_count(output_content)
256
+ return TokenUsage(
257
+ input_tokens=max(total_tokens - output_tokens, 0),
258
+ output_tokens=output_tokens,
259
+ total_tokens=total_tokens,
260
+ )
261
+
262
+
263
+ def estimate_mapping_message_tokens(message: Mapping[str, object]) -> int:
264
+ total = approximate_token_count(string_content(message.get("content")))
265
+ tool_calls = message.get("tool_calls")
266
+ if tool_calls:
267
+ total += approximate_token_count(json.dumps(tool_calls, ensure_ascii=False))
268
+ if message.get("role") == "tool":
269
+ total += approximate_token_count(string_content(message.get("tool_call_id")))
270
+ return total
271
+
272
+
273
+ def approximate_token_count(content: str) -> int:
274
+ if not content:
275
+ return 0
276
+ return max(1, (len(content) + 3) // 4)
277
+
278
+
279
+ def string_content(value: object) -> str:
280
+ if value is None:
281
+ return ""
282
+ if isinstance(value, str):
283
+ return value
284
+ return json.dumps(value, ensure_ascii=False)
285
+
286
+
287
+ def value_at(value: Any, key: str, default: Any = None) -> Any:
288
+ if isinstance(value, Mapping):
289
+ return value.get(key, default)
290
+ return getattr(value, key, default)
291
+
292
+
293
+ def nested_value_at(value: Any, *keys: str) -> Any:
294
+ current = value
295
+ for key in keys:
296
+ current = value_at(current, key)
297
+ if current is None:
298
+ return None
299
+ return current
300
+
301
+
302
+ def first_int_value(*values: Any) -> int | None:
303
+ for value in values:
304
+ if isinstance(value, bool) or value is None:
305
+ continue
306
+ if isinstance(value, int):
307
+ return max(0, value)
308
+ if isinstance(value, float):
309
+ return max(0, int(value))
310
+ if isinstance(value, str):
311
+ try:
312
+ return max(0, int(value))
313
+ except ValueError:
314
+ continue
315
+ return None