| """Method registry: get_method(name, **params) returns a fresh agent. |
| |
| Methods with heavy optional dependencies are imported lazily. |
| """ |
| from __future__ import annotations |
|
|
| from .base import BaseAgent, MemoryBuildResult, MemoryContext |
| from .long_context import LongContextAgent |
| from .bm25_rag import BM25RAGAgent |
| from .dense_rag import DenseRAGAgent |
|
|
| __all__ = ["BaseAgent", "MemoryBuildResult", "MemoryContext", "LongContextAgent", "BM25RAGAgent", |
| "DenseRAGAgent", "SUPPORTED_METHODS", "DEFAULT_METHODS", "INTERNAL_LLM_METHODS", |
| "EMBEDDING_DIM_METHODS", "EMBEDDING_MODEL_METHODS", "EMBEDDING_BATCH_METHODS", |
| "EMBEDDING_TOKENIZER_METHODS", "TOP_K_PARAMS", "embedding_params", "get_method"] |
|
|
| SUPPORTED_METHODS = ("long_context", "bm25_rag", "dense_rag", |
| "hipporag", "remem", "memos", "letta", "mirix", |
| "amem", "lightmem", "mem0") |
| DEFAULT_METHODS = ("long_context", "bm25_rag", "dense_rag") |
| INTERNAL_LLM_METHODS = ("memos", "letta", "mirix", "amem", "lightmem", |
| "hipporag", "remem", "mem0") |
| EMBEDDING_DIM_METHODS = ("hipporag", "remem", "letta", "mirix", "lightmem", "mem0", "memos") |
| EMBEDDING_MODEL_METHODS = EMBEDDING_DIM_METHODS + ("dense_rag", "amem") |
| |
| EMBEDDING_BATCH_METHODS = ("dense_rag", "hipporag", "remem", "mirix", "amem", "mem0", "memos") |
| |
| EMBEDDING_TOKENIZER_METHODS = ("amem", "mem0", "mirix") |
| |
| TOP_K_PARAMS = {"bm25_rag": "top_k", "dense_rag": "top_k", "hipporag": "qa_top_k", "remem": "qa_top_k", |
| "letta": "retrieve_num", "mirix": "retrieve_num", "amem": "retrieve_num", |
| "lightmem": "retrieve_num", "mem0": "retrieve_num", "memos": "retrieve_num"} |
|
|
|
|
| def embedding_params(name: str, embedding: dict) -> dict: |
| """The embedding settings of the backbone config that a method takes as parameters.""" |
| out = {} |
| if name in EMBEDDING_MODEL_METHODS: |
| out["embedding_model"] = embedding["model"] |
| if name in EMBEDDING_DIM_METHODS: |
| out["embedding_dim"] = embedding["dim"] |
| if name in EMBEDDING_BATCH_METHODS: |
| out["embedding_batch_size"] = embedding.get("batch_size") |
| return out |
|
|
|
|
| def get_method(name: str, **params) -> BaseAgent: |
| if name == "long_context": |
| return LongContextAgent(**params) |
| if name == "bm25_rag": |
| return BM25RAGAgent(**params) |
| if name == "dense_rag": |
| from ..llm_client import make_embedder |
| embed_model = params.pop("embedding_model", None) |
| batch_size = params.pop("embedding_batch_size", None) |
| return DenseRAGAgent(embedder=make_embedder(model=embed_model, batch_size=batch_size), **params) |
| if name == "hipporag": |
| from .hipporag_agent import HippoRAGAgent |
| return HippoRAGAgent(**params) |
| if name == "remem": |
| from .remem_agent import ReMemAgent |
| return ReMemAgent(**params) |
| if name == "memos": |
| from .memos_agent import MemOSAgent |
| return MemOSAgent(**params) |
| if name == "letta": |
| from .letta_agent import LettaAgent |
| return LettaAgent(**params) |
| if name == "mirix": |
| from .mirix_agent import MIRIXAgent |
| return MIRIXAgent(**params) |
| if name == "amem": |
| from .amem_agent import AMemAgent |
| return AMemAgent(**params) |
| if name == "lightmem": |
| from .lightmem_agent import LightMemAgent |
| return LightMemAgent(**params) |
| if name == "mem0": |
| from .mem0_agent import Mem0Agent |
| return Mem0Agent(**params) |
| raise ValueError(f"unknown method '{name}'; supported: {', '.join(SUPPORTED_METHODS)}") |
|
|
