Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions changelog.d/1921-llm-credential-cache.changed.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
- **Memoized the per-agent-build LLM credential read (#1921).**
`opencontractserver/llms/model_factory.py::_get_db_credentials` ran on every
agent build — every chat message, structured-output call, and
memory-curation task — and resolved the provider's secret through
`PipelineSettings.get_full_component_settings`, which reads the encrypted
secret store twice and so derived the Fernet key with the deliberately
expensive PBKDF2 KDF (480k HMAC iterations) twice per build whenever any
provider key was configured. The resolved per-provider credentials are now
cached in-process keyed on `(class_path, PipelineSettings.modified)` — the
same key the reranker/embedder instance caches use
(`opencontractserver/pipeline/utils.py`). Repeat builds skip the decryption,
while live rotation is preserved: a superuser key change calls
`PipelineSettings.save()`, which bumps `modified` (auto_now) and clears the
singleton cache, so the next build misses the memo and re-decrypts — no
redeploy, and no staleness beyond the existing 5-minute `PipelineSettings`
cache TTL (the live-configurability guarantee of #1897 is intact, as is the
DB-wins / env-fallback precedence). New `invalidate_credential_cache()`
mirrors the existing `invalidate_embedder_cache` / `invalidate_reranker_cache`
helpers (wired into `conftest.py` for test isolation; call it after any
out-of-band singleton write that bypasses `save()`). Tests:
`opencontractserver/tests/test_llm_model_factory.py::TestCredentialCache`.
2 changes: 2 additions & 0 deletions docs/architecture/llms/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -2868,6 +2868,8 @@ Resolution is **DB-wins / env-fallback**, applied by [`opencontractserver/llms/m

Any failure to build a credentialed model degrades to the env-fallback string, so a misconfiguration can never take the chat path down. The factory is invoked at every `make_pydantic_ai_agent` call site (document, corpus, and structured-output agents, plus the memory-curation tasks); it performs ORM access, so async call sites use the `abuild_agent_model()` wrapper.

Because the factory runs on every agent build, the resolved per-provider credentials are memoized in-process keyed on `(class_path, PipelineSettings.modified)` — the same cache key the reranker/embedder instance caches use. This skips the Fernet/PBKDF2 secret decryption on repeat builds while keeping rotation live: a superuser key change calls `PipelineSettings.save()`, which bumps `modified` and clears the singleton cache, so the next build misses the memo and re-decrypts (no redeploy, no staleness beyond the existing 5-minute `PipelineSettings` cache TTL). An out-of-band write that bypasses `save()` (e.g. `QuerySet.update`) should call `invalidate_credential_cache()`.

### Validation

`Corpus.save()` and `AgentConfiguration.save()` both run the resolver's `validate_model_spec()` — a malformed string or a provider with no registered `BaseLLMProvider` subclass raises `ValidationError({"preferred_llm": ...})`. The validator does not gate against `supported_models` so users aren't blocked from passing newly-released model names. Specs are normalised to canonical `"{provider}:{model}"` form on the way into the database.
19 changes: 11 additions & 8 deletions opencontractserver/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,16 +43,18 @@ def _invalidate_pipeline_settings_singleton_cache():
migration-seeded state on cache miss.

We also drop the process-local embedder/reranker *instance* caches
(``opencontractserver/pipeline/utils.py``). Those are keyed by
``(class_path, PipelineSettings.modified)``; because ``modified`` is the
stable migration-seeded value across every ``TestCase`` test that doesn't
write the singleton, two tests that mock ``get_component_by_name`` to
return *different* classes for the *same* path would otherwise collide on
the cached instance from whichever test ran first. Clearing before each
test restores per-test isolation without every test needing its own
teardown.
(``opencontractserver/pipeline/utils.py``) and the LLM-provider
*credential* cache (``opencontractserver/llms/model_factory.py``). All are
keyed by ``(class_path, PipelineSettings.modified)``; because ``modified``
is the stable migration-seeded value across every ``TestCase`` test that
doesn't write the singleton, two tests that mock ``get_component_by_name``
to return *different* classes for the *same* path (or configure *different*
provider credentials) would otherwise collide on the cached value from
whichever test ran first. Clearing before each test restores per-test
isolation without every test needing its own teardown.
"""
from opencontractserver.documents.models import PipelineSettings
from opencontractserver.llms.model_factory import invalidate_credential_cache
from opencontractserver.pipeline.utils import (
invalidate_embedder_cache,
invalidate_reranker_cache,
Expand All @@ -61,6 +63,7 @@ def _invalidate_pipeline_settings_singleton_cache():
PipelineSettings._invalidate_cache()
invalidate_embedder_cache()
invalidate_reranker_cache()
invalidate_credential_cache()
yield


Expand Down
79 changes: 71 additions & 8 deletions opencontractserver/llms/model_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
from __future__ import annotations

import logging
import threading
from typing import Any

from asgiref.sync import sync_to_async
Expand Down Expand Up @@ -76,13 +77,60 @@ def _provider_class_path(provider_key: str) -> str | None:
return definition.class_name if definition else None


# Process-local cache of resolved per-provider credentials keyed on
# ``(class_path, PipelineSettings.modified)`` — the same key the reranker and
# embedder instance caches use (``opencontractserver/pipeline/utils.py``).
#
# ``_get_db_credentials`` runs on *every* agent build — every chat message,
# every structured-output call, every memory-curation task. ``PipelineSettings``
# already caches the singleton row (Django cache, 5-min TTL), so this is not a
# DB round-trip per call, but resolving a provider's secret still decrypts the
# Fernet store, and that decryption derives its key with a deliberately
# expensive PBKDF2 KDF (hundreds of thousands of HMAC iterations — see
# ``PipelineSettings._derive_key``). ``get_full_component_settings`` reads the
# secret store twice (merged settings + secrets overlay), so two KDF passes ran
# on every build once any provider key was configured (issue #1921).
#
# Caching the *resolved* creds amortises that across calls. Live-rotation is
# preserved: every credential write in production flows through
# ``PipelineSettings.save()``, which bumps ``modified`` (auto_now) and clears
# the singleton cache, so the next lookup misses on the new key and re-decrypts
# — a rotated or cleared key takes effect within the existing cache-TTL window
# with no redeploy (the live-configurability guarantee of issue #1897), and the
# cache adds no staleness beyond ``get_instance``'s own TTL. The same issue
# #1410 caveat applies: code paths that bypass ``save()`` (``QuerySet.update``,
# ``bulk_update``, raw migrations) won't bump ``modified``; call
# :func:`invalidate_credential_cache` explicitly in those cases.
_CREDENTIAL_CACHE: dict[tuple[str, Any], dict[str, str]] = {}
# Guards against two concurrent builds paying the decryption cost twice.
# Lookups after warm-up are read-only, so no lock is held on the hot path.
_CREDENTIAL_CACHE_LOCK = threading.Lock()


def invalidate_credential_cache() -> None:
"""Drop cached per-provider credentials.

In normal operation this is unnecessary — the cache key includes
``PipelineSettings.modified`` so any settings write naturally invalidates
the cache on the next lookup across all workers. Kept for test isolation
and for callers that mutate the singleton out-of-band (a path that bypasses
``save()`` and so leaves ``modified`` unchanged).
"""
with _CREDENTIAL_CACHE_LOCK:
_CREDENTIAL_CACHE.clear()


def _get_db_credentials(provider_key: str) -> dict[str, str]:
"""Read DB-configured credentials for a provider from ``PipelineSettings``.

Returns a dict with optional ``api_key`` / ``base_url`` keys (only
present when non-empty). Empty when the provider is unknown or has no
credentials configured. Performs ORM access — invoke from a sync
context or via :func:`abuild_agent_model`.

The resolved creds are memoized per ``(class_path, PipelineSettings.modified)``
so repeat builds skip the Fernet/PBKDF2 decryption; see the
``_CREDENTIAL_CACHE`` comment for the invalidation contract.
"""
class_path = _provider_class_path(provider_key)
if not class_path:
Expand All @@ -91,7 +139,29 @@ def _get_db_credentials(provider_key: str) -> dict[str, str]:
try:
from opencontractserver.documents.models import PipelineSettings

stored = PipelineSettings.get_instance().get_full_component_settings(class_path)
instance = PipelineSettings.get_instance()
cache_key = (class_path, instance.modified)

# Fast path: a copy so callers can never mutate the cached dict.
cached = _CREDENTIAL_CACHE.get(cache_key)
if cached is not None:
return dict(cached)

with _CREDENTIAL_CACHE_LOCK:
# Double-check — another thread may have populated it while we waited.
cached = _CREDENTIAL_CACHE.get(cache_key)
if cached is not None:
return dict(cached)

stored = instance.get_full_component_settings(class_path)
creds: dict[str, str] = {}
for key in ("api_key", "base_url"):
value = (stored or {}).get(key)
if isinstance(value, str) and value.strip():
creds[key] = value.strip()

_CREDENTIAL_CACHE[cache_key] = creds
return dict(creds)
except _DB_READ_RECOVERABLE_ERRORS:
# DB unavailable (migrations / early startup) — fall back to env.
logger.debug(
Expand All @@ -101,13 +171,6 @@ def _get_db_credentials(provider_key: str) -> dict[str, str]:
)
return {}

creds: dict[str, str] = {}
for key in ("api_key", "base_url"):
value = (stored or {}).get(key)
if isinstance(value, str) and value.strip():
creds[key] = value.strip()
return creds


async def aget_provider_credentials(provider_key: str) -> dict[str, str]:
"""DB-configured ``api_key`` / ``base_url`` for a provider (empty when unset).
Expand Down
83 changes: 83 additions & 0 deletions opencontractserver/tests/test_llm_model_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
from opencontractserver.llms.model_factory import (
abuild_agent_model,
build_agent_model,
invalidate_credential_cache,
)
from opencontractserver.pipeline.base.settings_schema import (
get_secret_settings,
Expand Down Expand Up @@ -209,3 +210,85 @@ def test_has_value_flips_after_setting_secret(self):
self.assertTrue(schema["api_key"]["has_value"])
# The value itself is never exposed.
self.assertIsNone(schema["api_key"]["current_value"])


class TestCredentialCache(TestCase):
"""Resolved provider creds are memoized per ``PipelineSettings.modified``.

Regression coverage for issue #1921: the per-build credential read must
skip repeated Fernet/PBKDF2 decryption, yet a live key rotation (which
bumps ``modified`` via ``save()``) must still take effect without a
redeploy — preserving the live-configurability guarantee of issue #1897.
"""

def setUp(self):
reset_registry()
self.addCleanup(reset_registry)
PipelineSettings.clear_cache()
self.addCleanup(PipelineSettings.clear_cache)
# The memo is process-local and survives TestCase rollback; clear it
# explicitly so the suite is isolated under the unittest runner too
# (the conftest autouse fixture only applies under pytest).
invalidate_credential_cache()
self.addCleanup(invalidate_credential_cache)
openai_defn = get_llm_provider_by_key_cached("openai")
assert openai_defn is not None
self.openai_path = openai_defn.class_name

def _configure_openai_creds(self, *, api_key):
instance = PipelineSettings.get_instance()
instance.set_secrets({self.openai_path: {"api_key": api_key}})
# save() bumps ``modified`` (auto_now) and clears the singleton cache.
instance.save()

def test_resolution_memoized_across_builds(self):
"""The second build for a provider reuses creds without re-decrypting."""
self._configure_openai_creds(api_key="sk-db-key")
with mock.patch.object(
PipelineSettings,
"get_secrets",
autospec=True,
side_effect=PipelineSettings.get_secrets,
) as get_secrets_spy:
first = build_agent_model("openai:gpt-4o")
second = build_agent_model("openai:gpt-4o")
# The first build decrypts once — get_full_component_settings reads the
# secret store twice (merged settings + secrets overlay) — and the
# second build is served from the memo without touching the store.
self.assertEqual(get_secrets_spy.call_count, 2)
self.assertIsInstance(first, Model)
self.assertIsInstance(second, Model)

def test_rotated_key_takes_effect_without_redeploy(self):
"""A key rotation via save() busts the memo (preserves #1897)."""
self._configure_openai_creds(api_key="sk-old")
seen_keys: list[str | None] = []

def _capture(provider_key, model_name, creds):
seen_keys.append(creds.get("api_key"))
return object()

with mock.patch(
"opencontractserver.llms.model_factory._construct_model",
side_effect=_capture,
):
build_agent_model("openai:gpt-4o") # warms the memo with sk-old
# Rotate exactly as a superuser would: System Settings → save().
self._configure_openai_creds(api_key="sk-new")
build_agent_model("openai:gpt-4o")

self.assertEqual(seen_keys, ["sk-old", "sk-new"])

def test_invalidate_credential_cache_forces_redecrypt(self):
"""The explicit purge re-decrypts on the next build (out-of-band path)."""
self._configure_openai_creds(api_key="sk-db-key")
with mock.patch.object(
PipelineSettings,
"get_secrets",
autospec=True,
side_effect=PipelineSettings.get_secrets,
) as get_secrets_spy:
build_agent_model("openai:gpt-4o") # 2 decrypt reads, then cached
invalidate_credential_cache()
build_agent_model("openai:gpt-4o") # memo purged → 2 more reads
self.assertEqual(get_secrets_spy.call_count, 4)