refactor: use literal type for retrieve method config (#97)

Author: ankaisen

src/memu/app/__init__.py

@@ -1,4 +1,4 @@
-from memu.app.service import MemoryUser
+from memu.app.service import MemoryService
 from memu.app.settings import BlobConfig, DatabaseConfig, LLMConfig, MemorizeConfig, RetrieveConfig
 
-__all__ = ["BlobConfig", "DatabaseConfig", "LLMConfig", "MemorizeConfig", "MemoryUser", "RetrieveConfig"]
+__all__ = ["BlobConfig", "DatabaseConfig", "LLMConfig", "MemorizeConfig", "MemoryService", "RetrieveConfig"]

src/memu/app/service.py

@@ -34,19 +34,21 @@
 TConfigModel = TypeVar("TConfigModel", bound=BaseModel)
 
 
-class MemoryUser:
+class MemoryService:
     def __init__(
         self,
         *,
-        llm_config: dict[str, Any] | LLMConfig | None = None,
-        blob_config: dict[str, Any] | BlobConfig | None = None,
-        database_config: dict[str, Any] | DatabaseConfig | None = None,
-        memorize_config: dict[str, Any] | MemorizeConfig | None = None,
+        llm_config: LLMConfig | dict[str, Any] | None = None,
+        blob_config: BlobConfig | dict[str, Any] | None = None,
+        database_config: DatabaseConfig | dict[str, Any] | None = None,
+        memorize_config: MemorizeConfig | dict[str, Any] | None = None,
+        retrieve_config: RetrieveConfig | dict[str, Any] | None = None,
     ):
         self.llm_config = self._validate_config(llm_config, LLMConfig)
         self.blob_config = self._validate_config(blob_config, BlobConfig)
         self.database_config = self._validate_config(database_config, DatabaseConfig)
         self.memorize_config = self._validate_config(memorize_config, MemorizeConfig)
+        self.retrieve_config = self._validate_config(retrieve_config, RetrieveConfig)
         self.fs = LocalFS(self.blob_config.resources_dir)
         self.store = InMemoryStore()
         backend = self.llm_config.client_backend
@@ -788,19 +790,13 @@ async def retrieve(
         self,
         query: str,
         *,
-        retrieve_config: dict[str, Any] | RetrieveConfig | None = None,
         conversation_history: list[dict[str, str]] | None = None,
     ) -> dict[str, Any]:
         """
         Retrieve relevant memories based on the query using either RAG-based or LLM-based search.
 
         Args:
             query: The search query string
-            retrieve_config: Configuration for retrieval method and parameters.
-                           Can be a dict or RetrieveConfig object with:
-                           - method: 'rag' for embedding-based vector search (default),
-                                   'llm' for LLM-based semantic ranking
-                           - top_k: Maximum number of results per category (default: 5)
             conversation_history: Optional list of last 3 conversation turns, each with 'role' and 'content'.
                                 Example: [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
 
@@ -819,14 +815,6 @@ async def retrieve(
             - Pre-retrieval decision checks if retrieval is needed based on query type
             - Query rewriting incorporates conversation history for better context
         """
-        # Validate and resolve config
-        config = self._validate_config(retrieve_config, RetrieveConfig)
-
-        # Validate method
-        if config.method not in ("rag", "llm"):
-            msg = f"Invalid retrieval method '{config.method}'. Must be 'rag' or 'llm'."
-            raise ValueError(msg)
-
         # Step 1: Decide if retrieval is needed
         needs_retrieval, rewritten_query = await self._decide_if_retrieval_needed(query, conversation_history)
 
@@ -844,13 +832,13 @@ async def retrieve(
         logger.info(f"Query rewritten: '{query}' -> '{rewritten_query}'")
 
         # Step 2: Perform retrieval with rewritten query using configured method
-        if config.method == "llm":
+        if self.retrieve_config.method == "llm":
             results = await self._llm_based_retrieve(
-                rewritten_query, top_k=config.top_k, conversation_history=conversation_history
+                rewritten_query, top_k=self.retrieve_config.top_k, conversation_history=conversation_history
             )
         else:  # rag
             results = await self._embedding_based_retrieve(
-                rewritten_query, top_k=config.top_k, conversation_history=conversation_history
+                rewritten_query, top_k=self.retrieve_config.top_k, conversation_history=conversation_history
             )
 
         # Add metadata

src/memu/app/settings.py

@@ -1,9 +1,20 @@
-from pydantic import BaseModel, Field
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, BeforeValidator, Field
 
 from memu.prompts.memory_type import DEFAULT_MEMORY_TYPES
 from memu.prompts.memory_type import PROMPTS as DEFAULT_MEMORY_TYPE_PROMPTS
 
 
+def normalize_value(v: str) -> str:
+    if isinstance(v, str):
+        return v.strip().lower()
+    return v
+
+
+Normalize = BeforeValidator(normalize_value)
+
+
 def _default_memory_types() -> list[str]:
     return list(DEFAULT_MEMORY_TYPES)
 
@@ -56,10 +67,16 @@ class DatabaseConfig(BaseModel):
 
 
 class RetrieveConfig(BaseModel):
-    method: str = Field(
-        default="rag",
-        description="Retrieval method: 'rag' for embedding-based vector search, 'llm' for LLM-based ranking.",
-    )
+    """Configure retrieval behavior for `MemoryUser.retrieve`.
+
+    Attributes:
+        method: Retrieval strategy. Use "rag" for embedding-based vector search or
+            "llm" to delegate ranking to the LLM.
+        top_k: Maximum number of results to return per category (and per stage),
+            controlling breadth of the retrieved context.
+    """
+
+    method: Annotated[Literal["rag", "llm"], Normalize] = "rag"
     top_k: int = Field(
         default=5,
         description="Maximum number of results to return per category.",