File size: 3,820 Bytes
750b88d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """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")
# letta and lightmem embed one text per request inside their engines
EMBEDDING_BATCH_METHODS = ("dense_rag", "hipporag", "remem", "mirix", "amem", "mem0", "memos")
# methods that count embedding inputs with embedding.tokenizer (truncation, usage estimate)
EMBEDDING_TOKENIZER_METHODS = ("amem", "mem0", "mirix")
# the parameter that holds the number of retrieved items (top_k / --top-k)
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)}")
|