Fix provider settings model persistence

Author: macayu17

api/db/schema.sql

@@ -16,8 +16,9 @@ ALTER TABLE users ADD COLUMN IF NOT EXISTS max_memories_injected INT NOT NULL DE
 ALTER TABLE users ADD COLUMN IF NOT EXISTS retrieval_threshold DOUBLE PRECISION NOT NULL DEFAULT 0.5;
 ALTER TABLE users ADD COLUMN IF NOT EXISTS dedup_threshold DOUBLE PRECISION NOT NULL DEFAULT 0.95;
 
-ALTER TABLE users ADD COLUMN IF NOT EXISTS extraction_provider TEXT NOT NULL DEFAULT 'openai';
-ALTER TABLE users ADD COLUMN IF NOT EXISTS openai_api_key_encrypted BYTEA;
+ALTER TABLE users ADD COLUMN IF NOT EXISTS extraction_provider TEXT NOT NULL DEFAULT 'openai';
+ALTER TABLE users ADD COLUMN IF NOT EXISTS extraction_model TEXT NOT NULL DEFAULT 'gpt-4o-mini';
+ALTER TABLE users ADD COLUMN IF NOT EXISTS openai_api_key_encrypted BYTEA;
 ALTER TABLE users ADD COLUMN IF NOT EXISTS gemini_api_key_encrypted BYTEA;
 ALTER TABLE users ADD COLUMN IF NOT EXISTS anthropic_api_key_encrypted BYTEA;
 

api/models/user.py

@@ -38,9 +38,10 @@ class UserConfigResponse(BaseModel):
     dedup_threshold: float
 
 
-class UserProviderConfigUpdate(BaseModel):
-    extraction_provider: str | None = Field(default=None, pattern="^(openai|gemini|ollama|anthropic)$")
-    openai_api_key: str | None = Field(default=None, max_length=512)
+class UserProviderConfigUpdate(BaseModel):
+    extraction_provider: str | None = Field(default=None, pattern="^(openai|gemini|ollama|anthropic)$")
+    extraction_model: str | None = Field(default=None, min_length=1, max_length=120)
+    openai_api_key: str | None = Field(default=None, max_length=512)
     gemini_api_key: str | None = Field(default=None, max_length=512)
     anthropic_api_key: str | None = Field(default=None, max_length=512)
     clear_openai_key: bool = Field(default=False)

api/routes/users.py

@@ -142,9 +142,10 @@ async def update_current_user_provider_route(
 ) -> dict[str, object]:
     try:
         return await update_user_provider_config(
-            user["id"],
-            payload.extraction_provider,
-            payload.openai_api_key,
+            user["id"],
+            payload.extraction_provider,
+            payload.extraction_model,
+            payload.openai_api_key,
             payload.gemini_api_key,
             payload.anthropic_api_key,
             payload.clear_openai_key,

api/services/embedding.py

@@ -1,12 +1,19 @@
-from sentence_transformers import SentenceTransformer
+from typing import Protocol
 
 from api.config import settings
 
 
-_model: SentenceTransformer | None = None
+class EmbeddingModel(Protocol):
+    def encode(self, sentences: str | list[str], normalize_embeddings: bool) -> object:
+        ...
+
+
+_model: EmbeddingModel | None = None
 
 
 def load_model() -> None:
+    from sentence_transformers import SentenceTransformer
+
     global _model
     _model = SentenceTransformer(settings.embedding_model)
 
@@ -15,18 +22,31 @@ def is_model_loaded() -> bool:
     return _model is not None
 
 
-def get_model() -> SentenceTransformer:
+def get_model() -> EmbeddingModel:
     if _model is None:
         raise RuntimeError("Embedding model not loaded")
     return _model
 
 
 def embed(text: str) -> list[float]:
-    return get_model().encode(text, normalize_embeddings=True).tolist()
+    return to_float_vector(get_model().encode(text, normalize_embeddings=True))
 
 
 def embed_batch(texts: list[str]) -> list[list[float]]:
-    return get_model().encode(texts, normalize_embeddings=True).tolist()
+    encoded = get_model().encode(texts, normalize_embeddings=True)
+    if hasattr(encoded, "tolist"):
+        encoded = encoded.tolist()
+    if not isinstance(encoded, list):
+        raise RuntimeError("Embedding batch output is not a list")
+    return [to_float_vector(vector) for vector in encoded]
+
+
+def to_float_vector(value: object) -> list[float]:
+    if hasattr(value, "tolist"):
+        value = value.tolist()
+    if not isinstance(value, list):
+        raise RuntimeError("Embedding output is not a list")
+    return [float(item) for item in value]
 
 
 def format_embedding_for_pgvector(embedding: list[float]) -> str:

api/services/proxy.py

@@ -61,10 +61,11 @@ async def build_proxy_result(
         except Exception as error:
             injected_count = 0
             logger.warning("Retrieval failed, proceeding without memories: %s", error)
-    user_row = await db.fetchrow(
-        """SELECT id, external_id, extraction_provider,
-                  openai_api_key_encrypted, gemini_api_key_encrypted, anthropic_api_key_encrypted
-           FROM users WHERE id = $1""",
+    user_row = await db.fetchrow(
+        """SELECT id, external_id, extraction_provider,
+                  extraction_model,
+                  openai_api_key_encrypted, gemini_api_key_encrypted, anthropic_api_key_encrypted
+           FROM users WHERE id = $1""",
         user_id,
     )
     if user_row is None:
@@ -102,9 +103,10 @@ async def build_proxy_passthrough_result(
     from api.config import settings as _settings
     fallback_row = {
         "id": None,
-        "external_id": external_id,
-        "extraction_provider": _settings.extraction_provider,
-        "openai_api_key_encrypted": None,
+        "external_id": external_id,
+        "extraction_provider": _settings.extraction_provider,
+        "extraction_model": _settings.extraction_model,
+        "openai_api_key_encrypted": None,
         "gemini_api_key_encrypted": None,
         "anthropic_api_key_encrypted": None,
     }

api/services/users.py

@@ -3,40 +3,29 @@
 
 import asyncpg
 
-from api.config import settings
-from api.services.security import api_key_hashes_match, generate_api_key, hash_api_key
+from api.config import settings
+from api.services.provider_keys import normalize_provider_name
+from api.services.security import api_key_hashes_match, encrypt_provider_key, generate_api_key, hash_api_key
 
-from api.services.provider_keys import (
-    ProviderConfigError,
-    normalize_provider_name,
-)
-from api.services.security import encrypt_provider_key
 
+_USER_COLUMNS: str = (
+    "id, external_id, api_key_hash, created_at, max_memories_injected, "
+    "retrieval_threshold, dedup_threshold, extraction_provider, "
+    "extraction_model, openai_api_key_encrypted, gemini_api_key_encrypted, anthropic_api_key_encrypted"
+)
 
-_USER_COLUMNS: str = (
-    "id, external_id, api_key_hash, created_at, max_memories_injected, "
-    "retrieval_threshold, dedup_threshold, extraction_provider, "
-    "openai_api_key_encrypted, gemini_api_key_encrypted, anthropic_api_key_encrypted"
-)
 
-
-_PROVIDER_KEY_ENCRYPTED_COLUMNS: dict[str, str] = {
-    "openai": "openai_api_key_encrypted",
-    "gemini": "gemini_api_key_encrypted",
-    "anthropic": "anthropic_api_key_encrypted",
-}
-
-
-class CachedUser(TypedDict):
+class CachedUser(TypedDict):
     id: object
     external_id: str
     api_key_hash: str
     created_at: object
     max_memories_injected: int
     retrieval_threshold: float
-    dedup_threshold: float
-    extraction_provider: str
-    cached_at: float
+    dedup_threshold: float
+    extraction_provider: str
+    extraction_model: str
+    cached_at: float
 
 
 _user_auth_cache: dict[str, CachedUser] = {}
@@ -111,9 +100,9 @@ async def get_user_by_api_key(api_key: str, db: asyncpg.Connection) -> asyncpg.R
         row = await db.fetchrow(
             f"""
         SELECT users.id, users.external_id, user_api_keys.api_key_hash, users.created_at,
-               users.max_memories_injected, users.retrieval_threshold, users.dedup_threshold,
-               users.extraction_provider,
-               users.openai_api_key_encrypted, users.gemini_api_key_encrypted, users.anthropic_api_key_encrypted
+               users.max_memories_injected, users.retrieval_threshold, users.dedup_threshold,
+               users.extraction_provider, users.extraction_model,
+               users.openai_api_key_encrypted, users.gemini_api_key_encrypted, users.anthropic_api_key_encrypted
         FROM user_api_keys
         JOIN users ON users.id = user_api_keys.user_id
         WHERE user_api_keys.api_key_hash = $1
@@ -248,19 +237,25 @@ async def get_user_provider_config(user_id: object, db: asyncpg.Connection) -> d
 
 
 async def update_user_provider_config(
-    user_id: object,
-    extraction_provider: str | None,
-    openai_api_key: str | None,
+    user_id: object,
+    extraction_provider: str | None,
+    extraction_model: str | None,
+    openai_api_key: str | None,
     gemini_api_key: str | None,
     anthropic_api_key: str | None,
     clear_openai_key: bool,
     clear_gemini_key: bool,
     clear_anthropic_key: bool,
     db: asyncpg.Connection,
 ) -> dict[str, object]:
-    chosen_provider: str | None = None
-    if extraction_provider is not None:
-        chosen_provider = normalize_provider_name(extraction_provider)
+    chosen_provider: str | None = None
+    if extraction_provider is not None:
+        chosen_provider = normalize_provider_name(extraction_provider)
+    clean_extraction_model: str | None = None
+    if extraction_model is not None:
+        clean_extraction_model = extraction_model.strip()
+        if not clean_extraction_model:
+            raise ValueError("Extraction model is required")
 
     openai_blob: object = _SENTINEL_NO_UPDATE
     gemini_blob: object = _SENTINEL_NO_UPDATE
@@ -280,11 +275,14 @@ async def update_user_provider_config(
 
     assignments: list[str] = []
     params: list[object] = []
-    if chosen_provider is not None:
-        assignments.append("extraction_provider = $" + str(len(params) + 2))
-        params.append(chosen_provider)
-    if openai_blob is not _SENTINEL_NO_UPDATE:
-        assignments.append("openai_api_key_encrypted = $" + str(len(params) + 2))
+    if chosen_provider is not None:
+        assignments.append("extraction_provider = $" + str(len(params) + 2))
+        params.append(chosen_provider)
+    if clean_extraction_model is not None:
+        assignments.append("extraction_model = $" + str(len(params) + 2))
+        params.append(clean_extraction_model)
+    if openai_blob is not _SENTINEL_NO_UPDATE:
+        assignments.append("openai_api_key_encrypted = $" + str(len(params) + 2))
         params.append(openai_blob)
     if gemini_blob is not _SENTINEL_NO_UPDATE:
         assignments.append("gemini_api_key_encrypted = $" + str(len(params) + 2))
@@ -318,19 +316,21 @@ def cache_user_auth(api_key_hash: str, row: asyncpg.Record | dict[str, object])
     prune_user_auth_cache()
     max_memories_injected = get_row_value(row, "max_memories_injected", settings.max_memories_injected)
     retrieval_threshold = get_row_value(row, "retrieval_threshold", settings.retrieval_threshold)
-    dedup_threshold = get_row_value(row, "dedup_threshold", settings.dedup_threshold)
-    extraction_provider = get_row_value(row, "extraction_provider", settings.extraction_provider)
-    _user_auth_cache[api_key_hash] = {
+    dedup_threshold = get_row_value(row, "dedup_threshold", settings.dedup_threshold)
+    extraction_provider = get_row_value(row, "extraction_provider", settings.extraction_provider)
+    extraction_model = get_row_value(row, "extraction_model", settings.extraction_model)
+    _user_auth_cache[api_key_hash] = {
         "id": row["id"],
         "external_id": row["external_id"],
         "api_key_hash": api_key_hash,
         "created_at": row["created_at"],
         "max_memories_injected": int(max_memories_injected),
         "retrieval_threshold": float(retrieval_threshold),
-        "dedup_threshold": float(dedup_threshold),
-        "extraction_provider": str(extraction_provider),
-        "cached_at": monotonic(),
-    }
+        "dedup_threshold": float(dedup_threshold),
+        "extraction_provider": str(extraction_provider),
+        "extraction_model": str(extraction_model),
+        "cached_at": monotonic(),
+    }
     trim_user_auth_cache()
 
 

api/test_provider_keys.py

@@ -59,20 +59,22 @@ def test_mask_provider_key_redacts_middle() -> None:
     assert mask_provider_key(None) is None
 
 
-def test_resolve_user_provider_uses_user_key(monkeypatch) -> None:
-    monkeypatch.setattr(security.settings, "provider_key_encryption_key", Fernet.generate_key().decode())
-    user = {
-        "id": "u1",
-        "external_id": "alice",
-        "extraction_provider": "gemini",
-        "openai_api_key_encrypted": None,
-        "gemini_api_key_encrypted": encrypt_provider_key("user-gemini"),
-        "anthropic_api_key_encrypted": None,
-    }
-    resolved = provider_keys.resolve_user_provider(user)
-    assert resolved.name == "gemini"
-    assert resolved.api_key == "user-gemini"
-    assert resolved.source == "user"
+def test_resolve_user_provider_uses_user_key(monkeypatch) -> None:
+    monkeypatch.setattr(security.settings, "provider_key_encryption_key", Fernet.generate_key().decode())
+    user = {
+        "id": "u1",
+        "external_id": "alice",
+        "extraction_provider": "gemini",
+        "extraction_model": "gemini-1.5-flash",
+        "openai_api_key_encrypted": None,
+        "gemini_api_key_encrypted": encrypt_provider_key("user-gemini"),
+        "anthropic_api_key_encrypted": None,
+    }
+    resolved = provider_keys.resolve_user_provider(user)
+    assert resolved.name == "gemini"
+    assert resolved.model == "gemini-1.5-flash"
+    assert resolved.api_key == "user-gemini"
+    assert resolved.source == "user"
 
 
 def test_resolve_user_provider_override_key_wins(monkeypatch) -> None:
@@ -129,14 +131,16 @@ def test_settings_accepts_engram_prefixed_env_var(monkeypatch) -> None:
 
 def test_summarize_provider_for_response_masks_key(monkeypatch) -> None:
     monkeypatch.setattr(security.settings, "provider_key_encryption_key", Fernet.generate_key().decode())
-    user = {
-        "extraction_provider": "openai",
-        "openai_api_key_encrypted": encrypt_provider_key("sk-supersecretvalue12345"),
-        "gemini_api_key_encrypted": None,
-        "anthropic_api_key_encrypted": None,
-    }
+    user = {
+        "extraction_provider": "openai",
+        "extraction_model": "gpt-4.1-mini",
+        "openai_api_key_encrypted": encrypt_provider_key("sk-supersecretvalue12345"),
+        "gemini_api_key_encrypted": None,
+        "anthropic_api_key_encrypted": None,
+    }
     summary = provider_keys.summarize_provider_for_response(user)
     assert summary["has_user_api_key"] is True
-    assert summary["user_api_key_preview"] is not None
-    assert "supersecretvalue" not in summary["user_api_key_preview"]
-    assert summary["extraction_provider"] == "openai"
+    assert summary["user_api_key_preview"] is not None
+    assert "supersecretvalue" not in summary["user_api_key_preview"]
+    assert summary["extraction_provider"] == "openai"
+    assert summary["extraction_model"] == "gpt-4.1-mini"