betterdb-agent-cache 0.4.0__tar.gz

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 (47) hide show
  1. betterdb_agent_cache-0.4.0/.gitignore +6 -0
  2. betterdb_agent_cache-0.4.0/PKG-INFO +34 -0
  3. betterdb_agent_cache-0.4.0/RELEASE_NOTES.md +103 -0
  4. betterdb_agent_cache-0.4.0/betterdb_agent_cache/__init__.py +89 -0
  5. betterdb_agent_cache-0.4.0/betterdb_agent_cache/adapters/__init__.py +0 -0
  6. betterdb_agent_cache-0.4.0/betterdb_agent_cache/adapters/anthropic.py +248 -0
  7. betterdb_agent_cache-0.4.0/betterdb_agent_cache/adapters/langchain.py +189 -0
  8. betterdb_agent_cache-0.4.0/betterdb_agent_cache/adapters/langgraph.py +361 -0
  9. betterdb_agent_cache-0.4.0/betterdb_agent_cache/adapters/llamaindex.py +166 -0
  10. betterdb_agent_cache-0.4.0/betterdb_agent_cache/adapters/openai.py +205 -0
  11. betterdb_agent_cache-0.4.0/betterdb_agent_cache/adapters/openai_responses.py +212 -0
  12. betterdb_agent_cache-0.4.0/betterdb_agent_cache/agent_cache.py +262 -0
  13. betterdb_agent_cache-0.4.0/betterdb_agent_cache/analytics.py +113 -0
  14. betterdb_agent_cache-0.4.0/betterdb_agent_cache/cluster.py +63 -0
  15. betterdb_agent_cache-0.4.0/betterdb_agent_cache/default_cost_table.py +1986 -0
  16. betterdb_agent_cache-0.4.0/betterdb_agent_cache/errors.py +15 -0
  17. betterdb_agent_cache-0.4.0/betterdb_agent_cache/normalizer.py +157 -0
  18. betterdb_agent_cache-0.4.0/betterdb_agent_cache/telemetry.py +136 -0
  19. betterdb_agent_cache-0.4.0/betterdb_agent_cache/tiers/__init__.py +0 -0
  20. betterdb_agent_cache-0.4.0/betterdb_agent_cache/tiers/llm_cache.py +311 -0
  21. betterdb_agent_cache-0.4.0/betterdb_agent_cache/tiers/session_store.py +312 -0
  22. betterdb_agent_cache-0.4.0/betterdb_agent_cache/tiers/tool_cache.py +274 -0
  23. betterdb_agent_cache-0.4.0/betterdb_agent_cache/types.py +217 -0
  24. betterdb_agent_cache-0.4.0/betterdb_agent_cache/utils.py +88 -0
  25. betterdb_agent_cache-0.4.0/examples/anthropic/main.py +112 -0
  26. betterdb_agent_cache-0.4.0/examples/langchain/main.py +88 -0
  27. betterdb_agent_cache-0.4.0/examples/langgraph/main.py +136 -0
  28. betterdb_agent_cache-0.4.0/examples/llamaindex/main.py +109 -0
  29. betterdb_agent_cache-0.4.0/examples/openai/main.py +116 -0
  30. betterdb_agent_cache-0.4.0/examples/openai_responses/main.py +119 -0
  31. betterdb_agent_cache-0.4.0/hatch_build.py +48 -0
  32. betterdb_agent_cache-0.4.0/pyproject.toml +53 -0
  33. betterdb_agent_cache-0.4.0/scripts/__init__.py +0 -0
  34. betterdb_agent_cache-0.4.0/scripts/update_model_prices.py +119 -0
  35. betterdb_agent_cache-0.4.0/tests/__init__.py +0 -0
  36. betterdb_agent_cache-0.4.0/tests/adapters/__init__.py +0 -0
  37. betterdb_agent_cache-0.4.0/tests/adapters/test_anthropic.py +173 -0
  38. betterdb_agent_cache-0.4.0/tests/adapters/test_langchain.py +131 -0
  39. betterdb_agent_cache-0.4.0/tests/adapters/test_langgraph.py +228 -0
  40. betterdb_agent_cache-0.4.0/tests/adapters/test_llamaindex.py +142 -0
  41. betterdb_agent_cache-0.4.0/tests/adapters/test_openai.py +194 -0
  42. betterdb_agent_cache-0.4.0/tests/conftest.py +84 -0
  43. betterdb_agent_cache-0.4.0/tests/test_agent_cache.py +90 -0
  44. betterdb_agent_cache-0.4.0/tests/test_llm_cache.py +245 -0
  45. betterdb_agent_cache-0.4.0/tests/test_session_store.py +219 -0
  46. betterdb_agent_cache-0.4.0/tests/test_tool_cache.py +206 -0
  47. betterdb_agent_cache-0.4.0/tests/test_utils.py +154 -0
@@ -0,0 +1,6 @@
1
+ __pycache__/
2
+ *.pyc
3
+ *.pyo
4
+ .venv/
5
+ dist/
6
+ *.egg-info/
@@ -0,0 +1,34 @@
1
+ Metadata-Version: 2.4
2
+ Name: betterdb-agent-cache
3
+ Version: 0.4.0
4
+ Summary: Multi-tier exact-match cache for AI agent workloads backed by Valkey. LLM responses, tool results, and session state with built-in OpenTelemetry and Prometheus instrumentation.
5
+ Project-URL: Repository, https://github.com/BetterDB-inc/monitor
6
+ License: MIT
7
+ Keywords: agent,cache,langchain,langgraph,llm,opentelemetry,prometheus,redis,valkey
8
+ Requires-Python: >=3.11
9
+ Requires-Dist: opentelemetry-api>=1.20.0
10
+ Requires-Dist: prometheus-client>=0.19.0
11
+ Requires-Dist: valkey>=6.0.0
12
+ Provides-Extra: all
13
+ Requires-Dist: anthropic>=0.20.0; extra == 'all'
14
+ Requires-Dist: langchain-core>=0.1.0; extra == 'all'
15
+ Requires-Dist: langgraph>=0.1.0; extra == 'all'
16
+ Requires-Dist: llama-index-core>=0.10.0; extra == 'all'
17
+ Requires-Dist: openai>=1.0.0; extra == 'all'
18
+ Requires-Dist: posthog>=3.0.0; extra == 'all'
19
+ Provides-Extra: analytics
20
+ Requires-Dist: posthog>=3.0.0; extra == 'analytics'
21
+ Provides-Extra: anthropic
22
+ Requires-Dist: anthropic>=0.20.0; extra == 'anthropic'
23
+ Provides-Extra: dev
24
+ Requires-Dist: fakeredis[aioredis]>=2.20.0; extra == 'dev'
25
+ Requires-Dist: pytest-asyncio>=0.23.0; extra == 'dev'
26
+ Requires-Dist: pytest>=8.0.0; extra == 'dev'
27
+ Provides-Extra: langchain
28
+ Requires-Dist: langchain-core>=0.1.0; extra == 'langchain'
29
+ Provides-Extra: langgraph
30
+ Requires-Dist: langgraph>=0.1.0; extra == 'langgraph'
31
+ Provides-Extra: llamaindex
32
+ Requires-Dist: llama-index-core>=0.10.0; extra == 'llamaindex'
33
+ Provides-Extra: openai
34
+ Requires-Dist: openai>=1.0.0; extra == 'openai'
@@ -0,0 +1,103 @@
1
+ # betterdb-agent-cache v0.3.0
2
+
3
+ Python port of `@betterdb/agent-cache`. Multi-tier exact-match cache for AI agent
4
+ workloads backed by Valkey — LLM responses, tool results, and session state, with
5
+ built-in OpenTelemetry and Prometheus instrumentation.
6
+
7
+ Runs on vanilla Valkey 7+. No modules, no RedisJSON, no RediSearch. Works on
8
+ ElastiCache for Valkey, Memorystore for Valkey, and MemoryDB.
9
+
10
+ ---
11
+
12
+ ## Installation
13
+
14
+ ```sh
15
+ pip install betterdb-agent-cache
16
+ ```
17
+
18
+ Optional extras install the provider SDKs alongside the library:
19
+
20
+ ```sh
21
+ pip install "betterdb-agent-cache[openai]"
22
+ pip install "betterdb-agent-cache[anthropic]"
23
+ pip install "betterdb-agent-cache[langchain]"
24
+ pip install "betterdb-agent-cache[langgraph]"
25
+ pip install "betterdb-agent-cache[llamaindex]"
26
+ ```
27
+
28
+ ---
29
+
30
+ ## What's included
31
+
32
+ ### Three cache tiers
33
+
34
+ | Tier | Use for |
35
+ |---|---|
36
+ | `cache.llm` | LLM API responses — `check` / `store` / `store_multipart` / `invalidate_by_model` |
37
+ | `cache.tool` | Tool call results — `check` / `store` / `set_policy` / `invalidate_by_tool` |
38
+ | `cache.session` | Agent session state — `get` / `set` / `get_all` / `destroy_thread` / `touch` |
39
+
40
+ ### Provider adapters
41
+
42
+ | Import | Provider |
43
+ |---|---|
44
+ | `betterdb_agent_cache.adapters.openai` | OpenAI Chat Completions |
45
+ | `betterdb_agent_cache.adapters.openai_responses` | OpenAI Responses API |
46
+ | `betterdb_agent_cache.adapters.anthropic` | Anthropic Messages |
47
+ | `betterdb_agent_cache.adapters.llamaindex` | LlamaIndex |
48
+ | `betterdb_agent_cache.adapters.langchain` | LangChain `BaseCache` |
49
+ | `betterdb_agent_cache.adapters.langgraph` | LangGraph `BaseCheckpointSaver` |
50
+
51
+ ### Pluggable binary normalizer
52
+
53
+ ```python
54
+ from betterdb_agent_cache import compose_normalizer, hash_base64
55
+
56
+ normalizer = compose_normalizer({"base64": hash_base64})
57
+ ```
58
+
59
+ ### Observability
60
+
61
+ - OpenTelemetry spans on every cache operation
62
+ - Prometheus counters, histograms, and gauges (`requests_total`, `operation_duration_seconds`,
63
+ `cost_saved_total`, `stored_bytes_total`, `active_sessions`)
64
+
65
+ ---
66
+
67
+ ## Quick start
68
+
69
+ ```python
70
+ import asyncio
71
+ import valkey.asyncio as valkey_client
72
+ from betterdb_agent_cache import AgentCache, ModelCost, TierDefaults
73
+ from betterdb_agent_cache.adapters.openai import prepare_params
74
+ from betterdb_agent_cache.types import AgentCacheOptions
75
+
76
+ client = valkey_client.Valkey(host="localhost", port=6379)
77
+ cache = AgentCache(AgentCacheOptions(
78
+ client=client,
79
+ tier_defaults={"llm": TierDefaults(ttl=3600)},
80
+ cost_table={"gpt-4o-mini": ModelCost(input_per_1k=0.00015, output_per_1k=0.0006)},
81
+ ))
82
+
83
+ async def main():
84
+ params = await prepare_params({
85
+ "model": "gpt-4o-mini",
86
+ "messages": [{"role": "user", "content": "What is 2+2?"}],
87
+ })
88
+
89
+ result = await cache.llm.check(params)
90
+ if result.hit:
91
+ print("Cache hit:", result.response)
92
+ else:
93
+ # ... call OpenAI ...
94
+ await cache.llm.store(params, "Four")
95
+
96
+ asyncio.run(main())
97
+ ```
98
+
99
+ ---
100
+
101
+ ## Full changelog
102
+
103
+ See [CHANGELOG.md](./CHANGELOG.md) for detailed history.
@@ -0,0 +1,89 @@
1
+ from .agent_cache import AgentCache
2
+ from .default_cost_table import DEFAULT_COST_TABLE
3
+ from .analytics import Analytics
4
+ from .errors import AgentCacheError, AgentCacheUsageError, ValkeyCommandError
5
+ from .normalizer import (
6
+ BinaryNormalizer,
7
+ BinaryRef,
8
+ NormalizerConfig,
9
+ compose_normalizer,
10
+ default_normalizer,
11
+ fetch_and_hash,
12
+ hash_base64,
13
+ hash_bytes,
14
+ hash_url,
15
+ passthrough,
16
+ )
17
+ from .types import (
18
+ AgentCacheOptions,
19
+ AgentCacheStats,
20
+ BinaryBlock,
21
+ BlockHints,
22
+ ContentBlock,
23
+ LlmCacheMessage,
24
+ LlmCacheParams,
25
+ LlmCacheResult,
26
+ LlmStoreOptions,
27
+ ModelCost,
28
+ ReasoningBlock,
29
+ SessionStats,
30
+ TextBlock,
31
+ TierDefaults,
32
+ TierStats,
33
+ ToolCallBlock,
34
+ ToolCacheResult,
35
+ ToolDefinition,
36
+ ToolEffectivenessEntry,
37
+ ToolPolicy,
38
+ ToolResultBlock,
39
+ ToolStats,
40
+ ToolStoreOptions,
41
+ )
42
+
43
+ __all__ = [
44
+ # Main class
45
+ "AgentCache",
46
+ "DEFAULT_COST_TABLE",
47
+ # Types
48
+ "AgentCacheOptions",
49
+ "AgentCacheStats",
50
+ "LlmCacheParams",
51
+ "LlmCacheMessage",
52
+ "LlmStoreOptions",
53
+ "LlmCacheResult",
54
+ "ToolStoreOptions",
55
+ "ToolPolicy",
56
+ "ToolCacheResult",
57
+ "ToolDefinition",
58
+ "ModelCost",
59
+ "TierDefaults",
60
+ "TierStats",
61
+ "SessionStats",
62
+ "ToolStats",
63
+ "ToolEffectivenessEntry",
64
+ # Content blocks
65
+ "ContentBlock",
66
+ "TextBlock",
67
+ "BinaryBlock",
68
+ "ToolCallBlock",
69
+ "ToolResultBlock",
70
+ "ReasoningBlock",
71
+ "BlockHints",
72
+ # Normalizer
73
+ "BinaryRef",
74
+ "BinaryNormalizer",
75
+ "NormalizerConfig",
76
+ "hash_base64",
77
+ "hash_bytes",
78
+ "hash_url",
79
+ "fetch_and_hash",
80
+ "passthrough",
81
+ "compose_normalizer",
82
+ "default_normalizer",
83
+ # Errors
84
+ "AgentCacheError",
85
+ "AgentCacheUsageError",
86
+ "ValkeyCommandError",
87
+ # Analytics
88
+ "Analytics",
89
+ ]
@@ -0,0 +1,248 @@
1
+ """Anthropic Messages adapter.
2
+
3
+ Converts params for ``anthropic.messages.create()`` into ``LlmCacheParams``.
4
+
5
+ Usage::
6
+
7
+ from betterdb_agent_cache.adapters.anthropic import prepare_params
8
+
9
+ params = await prepare_params({"model": "claude-opus-4-5", "messages": [...], "max_tokens": 1024})
10
+ result = await cache.llm.check(params)
11
+ """
12
+ from __future__ import annotations
13
+
14
+ import base64
15
+ import hashlib
16
+ import json
17
+ from dataclasses import dataclass, field
18
+ from typing import Any
19
+
20
+ from ..normalizer import BinaryNormalizer, default_normalizer
21
+ from ..types import (
22
+ BinaryBlock,
23
+ BlockHints,
24
+ ContentBlock,
25
+ LlmCacheParams,
26
+ ReasoningBlock,
27
+ TextBlock,
28
+ ToolCallBlock,
29
+ ToolResultBlock,
30
+ )
31
+
32
+
33
+ @dataclass
34
+ class AnthropicPrepareOptions:
35
+ normalizer: BinaryNormalizer = field(default_factory=lambda: default_normalizer)
36
+
37
+
38
+ def _build_cache_hints(cache_control: dict[str, Any] | None) -> BlockHints | None:
39
+ if not cache_control:
40
+ return None
41
+ hint: BlockHints = {"anthropicCacheControl": {"type": "ephemeral"}}
42
+ if cache_control.get("ttl"):
43
+ hint["anthropicCacheControl"]["ttl"] = cache_control["ttl"] # type: ignore[typeddict-item]
44
+ return hint
45
+
46
+
47
+ async def _normalize_block(
48
+ block: dict[str, Any],
49
+ normalizer: BinaryNormalizer,
50
+ ) -> ContentBlock | None:
51
+ t = block.get("type")
52
+
53
+ if t == "text":
54
+ result: TextBlock = {"type": "text", "text": block["text"]}
55
+ hints = _build_cache_hints(block.get("cache_control"))
56
+ if hints:
57
+ result["hints"] = hints
58
+ return result
59
+
60
+ if t == "image":
61
+ src = block.get("source", {})
62
+ src_type = src.get("type")
63
+ media_type = "image/*"
64
+
65
+ if src_type == "base64":
66
+ source: dict[str, Any] = {"type": "base64", "data": src["data"]}
67
+ media_type = src.get("media_type", "image/*")
68
+ elif src_type == "url":
69
+ source = {"type": "url", "url": src["url"]}
70
+ elif src_type == "file":
71
+ source = {"type": "file_id", "file_id": src["file_id"], "provider": "anthropic"}
72
+ else:
73
+ return None
74
+
75
+ ref = await normalizer({"kind": "image", "source": source})
76
+ img_block: BinaryBlock = {"type": "binary", "kind": "image", "mediaType": media_type, "ref": ref}
77
+ hints = _build_cache_hints(block.get("cache_control"))
78
+ if hints:
79
+ img_block["hints"] = hints
80
+ return img_block
81
+
82
+ if t == "document":
83
+ src = block.get("source", {})
84
+ src_type = src.get("type")
85
+
86
+ if src_type == "content":
87
+ full_json = json.dumps(src.get("content"), sort_keys=True)
88
+ digest = hashlib.sha256(full_json.encode()).hexdigest()
89
+ ref = f"nested:sha256:{digest}"
90
+ doc_block: BinaryBlock = {
91
+ "type": "binary", "kind": "document",
92
+ "mediaType": "application/x-nested-content", "ref": ref,
93
+ }
94
+ hints = _build_cache_hints(block.get("cache_control"))
95
+ if hints:
96
+ doc_block["hints"] = hints
97
+ return doc_block
98
+
99
+ media_type = "application/octet-stream"
100
+ if src_type == "base64":
101
+ source = {"type": "base64", "data": src["data"]}
102
+ media_type = src.get("media_type", "application/pdf")
103
+ elif src_type == "text":
104
+ encoded = base64.b64encode(src["text"].encode()).decode()
105
+ source = {"type": "base64", "data": encoded}
106
+ media_type = "text/plain"
107
+ elif src_type == "url":
108
+ source = {"type": "url", "url": src["url"]}
109
+ media_type = "application/pdf"
110
+ elif src_type == "file":
111
+ source = {"type": "file_id", "file_id": src["file_id"], "provider": "anthropic"}
112
+ else:
113
+ return None
114
+
115
+ ref = await normalizer({"kind": "document", "source": source})
116
+ doc_b: BinaryBlock = {"type": "binary", "kind": "document", "mediaType": media_type, "ref": ref}
117
+ hints = _build_cache_hints(block.get("cache_control"))
118
+ if hints:
119
+ doc_b["hints"] = hints
120
+ return doc_b
121
+
122
+ if t == "tool_use":
123
+ return {
124
+ "type": "tool_call",
125
+ "id": block["id"],
126
+ "name": block["name"],
127
+ "args": block.get("input", {}),
128
+ }
129
+
130
+ if t == "thinking":
131
+ result_r: ReasoningBlock = {"type": "reasoning", "text": block["thinking"]}
132
+ if block.get("signature"):
133
+ result_r["opaqueSignature"] = block["signature"]
134
+ return result_r
135
+
136
+ if t == "redacted_thinking":
137
+ return {
138
+ "type": "reasoning",
139
+ "text": "",
140
+ "redacted": True,
141
+ "opaqueSignature": block.get("data", ""),
142
+ }
143
+
144
+ return None
145
+
146
+
147
+ async def _normalize_tool_result_content(
148
+ content: str | list[dict[str, Any]],
149
+ normalizer: BinaryNormalizer,
150
+ ) -> list[TextBlock | BinaryBlock]:
151
+ if isinstance(content, str):
152
+ return [{"type": "text", "text": content}]
153
+ blocks: list[TextBlock | BinaryBlock] = []
154
+ for item in content:
155
+ if item.get("type") == "text":
156
+ blocks.append({"type": "text", "text": item["text"]})
157
+ elif item.get("type") == "image":
158
+ b = await _normalize_block(item, normalizer)
159
+ if b and b.get("type") == "binary":
160
+ blocks.append(b) # type: ignore[arg-type]
161
+ return blocks
162
+
163
+
164
+ async def prepare_params(
165
+ params: dict[str, Any],
166
+ opts: AnthropicPrepareOptions | None = None,
167
+ ) -> LlmCacheParams:
168
+ """Normalise Anthropic Messages params to ``LlmCacheParams``."""
169
+ normalizer = opts.normalizer if opts else default_normalizer
170
+ messages: list[Any] = []
171
+
172
+ # System message
173
+ system = params.get("system")
174
+ if system:
175
+ if isinstance(system, str):
176
+ messages.append({"role": "system", "content": system})
177
+ else:
178
+ sys_blocks: list[TextBlock] = []
179
+ for b in system:
180
+ tb: TextBlock = {"type": "text", "text": b["text"]}
181
+ hints = _build_cache_hints(b.get("cache_control"))
182
+ if hints:
183
+ tb["hints"] = hints
184
+ sys_blocks.append(tb)
185
+ messages.append({"role": "system", "content": sys_blocks})
186
+
187
+ for msg in params.get("messages", []):
188
+ role = msg["role"]
189
+ content = msg["content"]
190
+
191
+ if role == "assistant":
192
+ if isinstance(content, str):
193
+ messages.append({
194
+ "role": "assistant",
195
+ "content": [{"type": "text", "text": content}],
196
+ })
197
+ else:
198
+ blocks: list[ContentBlock] = []
199
+ for blk in content:
200
+ b = await _normalize_block(blk, normalizer)
201
+ if b is not None:
202
+ blocks.append(b)
203
+ messages.append({"role": "assistant", "content": blocks})
204
+
205
+ else: # user
206
+ if isinstance(content, str):
207
+ messages.append({
208
+ "role": "user",
209
+ "content": [{"type": "text", "text": content}],
210
+ })
211
+ continue
212
+
213
+ tool_results = [p for p in content if p.get("type") == "tool_result"]
214
+ others = [p for p in content if p.get("type") != "tool_result"]
215
+
216
+ for tr in tool_results:
217
+ tr_content = await _normalize_tool_result_content(
218
+ tr.get("content") or "", normalizer
219
+ )
220
+ messages.append({
221
+ "role": "tool",
222
+ "toolCallId": tr["tool_use_id"],
223
+ "content": tr_content,
224
+ })
225
+
226
+ if others:
227
+ user_blocks: list[ContentBlock] = []
228
+ for blk in others:
229
+ b = await _normalize_block(blk, normalizer)
230
+ if b is not None:
231
+ user_blocks.append(b)
232
+ messages.append({"role": "user", "content": user_blocks})
233
+
234
+ result: LlmCacheParams = {"model": params["model"], "messages": messages}
235
+ if params.get("temperature") is not None:
236
+ result["temperature"] = params["temperature"]
237
+ if params.get("top_p") is not None:
238
+ result["top_p"] = params["top_p"]
239
+ if params.get("max_tokens") is not None:
240
+ result["max_tokens"] = params["max_tokens"]
241
+ if params.get("tools") is not None:
242
+ result["tools"] = params["tools"]
243
+ if params.get("tool_choice") is not None:
244
+ result["tool_choice"] = params["tool_choice"]
245
+ if params.get("stop_sequences") is not None:
246
+ result["stop"] = params["stop_sequences"]
247
+
248
+ return result
@@ -0,0 +1,189 @@
1
+ """LangChain cache adapter.
2
+
3
+ Implements LangChain's ``BaseCache`` interface backed by the AgentCache LLM tier.
4
+
5
+ Usage::
6
+
7
+ from langchain_openai import ChatOpenAI
8
+ from betterdb_agent_cache.adapters.langchain import BetterDBLlmCache
9
+
10
+ lc_cache = BetterDBLlmCache(cache=agent_cache)
11
+ llm = ChatOpenAI(model="gpt-4o", cache=lc_cache)
12
+ """
13
+ from __future__ import annotations
14
+
15
+ import json
16
+ import re
17
+ from typing import TYPE_CHECKING, Any, Optional, Sequence
18
+
19
+ try:
20
+ from langchain_core.caches import BaseCache
21
+ from langchain_core.messages import AIMessage
22
+ from langchain_core.outputs import ChatGeneration, Generation
23
+ _LANGCHAIN_AVAILABLE = True
24
+ except ImportError:
25
+ _LANGCHAIN_AVAILABLE = False
26
+ BaseCache = object # type: ignore[assignment,misc]
27
+ Generation = Any # type: ignore[misc,assignment]
28
+ ChatGeneration = Any # type: ignore[misc,assignment]
29
+
30
+ if TYPE_CHECKING:
31
+ from ..agent_cache import AgentCache
32
+
33
+
34
+ _KV_RE = re.compile(r'\b(\w+):"([^"]*)"')
35
+
36
+ _NUMERIC_PARAMS = {"temperature", "top_p"}
37
+ _INT_PARAMS = {"max_tokens"}
38
+
39
+
40
+ def _parse_llm_params(llm_string: str) -> dict[str, Any]:
41
+ """Extract model and sampling parameters from LangChain's serialised llm_string.
42
+
43
+ Handles two formats:
44
+ - Comma-separated key:value pairs: ``model_name:"gpt-4o-mini",temperature:"0.7",...``
45
+ - LangChain JSON serialisation: ``{"kwargs": {"model_name": "gpt-4o-mini", ...}, ...}---[...]``
46
+
47
+ Returns a dict suitable for use as ``LlmCacheParams`` (without ``messages``).
48
+ Falls back to using the raw ``llm_string`` as the model name if extraction fails.
49
+ """
50
+ # Format 2: JSON serialisation with optional ---[...] suffix
51
+ try:
52
+ json_part = llm_string.split("---")[0].strip()
53
+ data = json.loads(json_part)
54
+ kwargs = data.get("kwargs", data)
55
+ params: dict[str, Any] = {}
56
+ model = kwargs.get("model_name") or kwargs.get("model")
57
+ if model:
58
+ params["model"] = str(model)
59
+ for field in ("temperature", "top_p", "max_tokens"):
60
+ if (v := kwargs.get(field)) is not None:
61
+ params[field] = v
62
+ if params.get("model"):
63
+ return params
64
+ except (json.JSONDecodeError, TypeError, AttributeError):
65
+ pass
66
+
67
+ # Format 1: comma-separated key:"value" pairs
68
+ params = {}
69
+ for m in _KV_RE.finditer(llm_string):
70
+ k, v = m.group(1), m.group(2)
71
+ if k in ("model_name", "model"):
72
+ params["model"] = v
73
+ elif k in _NUMERIC_PARAMS:
74
+ try:
75
+ params[k] = float(v)
76
+ except ValueError:
77
+ pass
78
+ elif k in _INT_PARAMS:
79
+ try:
80
+ params[k] = int(v)
81
+ except ValueError:
82
+ pass
83
+
84
+ if not params.get("model"):
85
+ params["model"] = llm_string
86
+ return params
87
+
88
+
89
+ class BetterDBLlmCache(BaseCache):
90
+ """LangChain ``BaseCache`` implementation backed by AgentCache.
91
+
92
+ This cache is async-only. LangChain calls ``alookup`` / ``aupdate`` in
93
+ async pipelines. The synchronous ``lookup`` / ``update`` methods raise
94
+ ``RuntimeError`` — use an async LangChain invocation (``ainvoke`` /
95
+ ``astream``) to avoid hitting them.
96
+ """
97
+
98
+ def __init__(self, cache: "AgentCache") -> None:
99
+ if not _LANGCHAIN_AVAILABLE:
100
+ raise ImportError(
101
+ "langchain-core is required for BetterDBLlmCache. "
102
+ "Install it with: pip install betterdb-agent-cache[langchain]"
103
+ )
104
+ super().__init__()
105
+ self._cache = cache
106
+
107
+ # ── Async interface (primary) ──────────────────────────────────────────
108
+
109
+ async def alookup(
110
+ self, prompt: str, llm_string: str
111
+ ) -> Optional[list[Generation]]:
112
+ result = await self._cache.llm.check({
113
+ **_parse_llm_params(llm_string),
114
+ "messages": [{"role": "user", "content": prompt}],
115
+ })
116
+ if not result.hit or not result.response:
117
+ return None
118
+ try:
119
+ parsed: list[dict[str, Any]] = json.loads(result.response)
120
+ return [
121
+ ChatGeneration(text=g.get("text", ""), message=AIMessage(g.get("text", "")))
122
+ for g in parsed
123
+ ]
124
+ except (json.JSONDecodeError, TypeError):
125
+ text = result.response
126
+ return [ChatGeneration(text=text, message=AIMessage(text))]
127
+
128
+ async def aupdate(
129
+ self, prompt: str, llm_string: str, return_val: Sequence[Any]
130
+ ) -> None:
131
+ if not return_val:
132
+ return
133
+ stripped = [{"text": g.text if hasattr(g, "text") else g.get("text", "")} for g in return_val]
134
+ text = json.dumps(stripped)
135
+
136
+ # Extract token counts — try usage_metadata first, fall back to response_metadata
137
+ first = return_val[0] if return_val else None
138
+ tokens = None
139
+ if first is not None:
140
+ msg = getattr(first, "message", None) or (
141
+ first.get("message") if isinstance(first, dict) else None
142
+ )
143
+ if msg is not None:
144
+ # Newer LangChain: AIMessage.usage_metadata = {"input_tokens": N, "output_tokens": M}
145
+ usage = getattr(msg, "usage_metadata", None)
146
+ if isinstance(usage, dict) and usage:
147
+ inp = usage.get("input_tokens")
148
+ out = usage.get("output_tokens")
149
+ if inp is not None and out is not None:
150
+ tokens = {"input": inp, "output": out}
151
+
152
+ # Fallback: AIMessage.response_metadata["token_usage"] = {"prompt_tokens": N, "completion_tokens": M}
153
+ if tokens is None:
154
+ rm = getattr(msg, "response_metadata", None) or {}
155
+ tu = rm.get("token_usage") if isinstance(rm, dict) else None
156
+ if isinstance(tu, dict) and tu:
157
+ inp = tu.get("prompt_tokens")
158
+ out = tu.get("completion_tokens")
159
+ if inp is not None and out is not None:
160
+ tokens = {"input": inp, "output": out}
161
+
162
+ from ..types import LlmStoreOptions
163
+ await self._cache.llm.store(
164
+ {**_parse_llm_params(llm_string), "messages": [{"role": "user", "content": prompt}]},
165
+ text,
166
+ LlmStoreOptions(tokens=tokens) if tokens else None,
167
+ )
168
+
169
+ # ── Sync interface (not supported) ────────────────────────────────────
170
+
171
+ def lookup(self, prompt: str, llm_string: str) -> Optional[list[Generation]]:
172
+ raise RuntimeError(
173
+ "BetterDBLlmCache is async-only. "
174
+ "Use an async LangChain invocation (ainvoke / astream)."
175
+ )
176
+
177
+ def update(self, prompt: str, llm_string: str, return_val: Sequence[Any]) -> None:
178
+ raise RuntimeError(
179
+ "BetterDBLlmCache is async-only. "
180
+ "Use an async LangChain invocation (ainvoke / astream)."
181
+ )
182
+
183
+ def clear(self, **kwargs: Any) -> None:
184
+ raise RuntimeError(
185
+ "BetterDBLlmCache is async-only. Use aclear() instead."
186
+ )
187
+
188
+ async def aclear(self, **kwargs: Any) -> None:
189
+ await self._cache.llm.clear()