Skip to content

Commit 774537e

Browse files
authored
fix: Update MemReader configuration with backup support
Enhanced MemReader configuration to support backup client and general model.
1 parent 8db7895 commit 774537e

File tree

1 file changed

+32
-15
lines changed

1 file changed

+32
-15
lines changed

src/memos/api/config.py

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -335,23 +335,40 @@ def get_activation_config() -> dict[str, Any]:
335335

336336
@staticmethod
337337
def get_memreader_config() -> dict[str, Any]:
338-
"""Get MemReader configuration for chat/doc extraction (fine-tuned 0.6B model)."""
339-
return {
340-
"backend": "openai",
341-
"config": {
342-
"model_name_or_path": os.getenv("MEMRADER_MODEL", "gpt-4o-mini"),
343-
"temperature": 0.6,
344-
"max_tokens": int(os.getenv("MEMRADER_MAX_TOKENS", "8000")),
345-
"top_p": 0.95,
346-
"top_k": 20,
347-
"api_key": os.getenv("MEMRADER_API_KEY", "EMPTY"),
348-
# Default to OpenAI base URL when env var is not provided to satisfy pydantic
349-
# validation requirements during tests/import.
350-
"api_base": os.getenv("MEMRADER_API_BASE", "https://api.openai.com/v1"),
351-
"remove_think_prefix": True,
352-
},
338+
"""Get MemReader configuration for chat/doc extraction (fine-tuned 0.6B model).
339+
340+
When MEMREADER_GENERAL_MODEL is configured (i.e. a separate stable LLM exists),
341+
the backup client is automatically enabled so that primary failures (self-deployed
342+
model) fall back to the general LLM.
343+
"""
344+
config = {
345+
"model_name_or_path": os.getenv("MEMRADER_MODEL", "gpt-4o-mini"),
346+
"temperature": 0.6,
347+
"max_tokens": int(os.getenv("MEMRADER_MAX_TOKENS", "8000")),
348+
"top_p": 0.95,
349+
"top_k": 20,
350+
"api_key": os.getenv("MEMRADER_API_KEY", "EMPTY"),
351+
# Default to OpenAI base URL when env var is not provided to satisfy pydantic
352+
# validation requirements during tests/import.
353+
"api_base": os.getenv("MEMRADER_API_BASE", "https://api.openai.com/v1"),
354+
"remove_think_prefix": True,
353355
}
354356

357+
general_model = os.getenv("MEMREADER_GENERAL_MODEL")
358+
enable_backup = os.getenv("MEMREADER_ENABLE_BACKUP", "false").lower() == "true"
359+
if general_model and enable_backup:
360+
config["backup_client"] = True
361+
config["backup_model_name_or_path"] = general_model
362+
config["backup_api_key"] = os.getenv(
363+
"MEMREADER_GENERAL_API_KEY", os.getenv("OPENAI_API_KEY", "EMPTY")
364+
)
365+
config["backup_api_base"] = os.getenv(
366+
"MEMREADER_GENERAL_API_BASE",
367+
os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
368+
)
369+
370+
return {"backend": "openai", "config": config}
371+
355372
@staticmethod
356373
def get_memreader_general_llm_config() -> dict[str, Any]:
357374
"""Get general LLM configuration for non-chat/doc tasks.

0 commit comments

Comments
 (0)