diff --git a/Autotests/import_knowledge/README.md b/Autotests/import_knowledge/README.md index fa943687..6280b439 100644 --- a/Autotests/import_knowledge/README.md +++ b/Autotests/import_knowledge/README.md @@ -37,7 +37,7 @@ The integration tests start and remove the container themselves through `scripts ## Unit: scripts/import_knowledge.sh -The script reads `EMBEDDING_PROVIDER`, runs `import-knowledge` (or `import-knowledge --local`), and writes a per-provider sentinel under `CHROMA_DB_PATH` so a second start skips the import. The stub on PATH records how `import-knowledge` was called. +The script reads `EMBEDDING_PROVIDER` and `EMBEDDING_MODEL`, runs `import-knowledge --provider openai|asicloud [--model ]` (or `import-knowledge --local`), and writes a per-provider sentinel under `CHROMA_DB_PATH` so a second start skips the import. The stub on PATH records how `import-knowledge` was called. ### 1. test_local_runs_import_and_writes_sentinel @@ -49,7 +49,7 @@ Local provider runs the import and leaves the local sentinel behind. OpenAI with `OPENAI_API_KEY` set runs the import and leaves the OpenAI sentinel. -- Checks: exit 0; the stub was called without `--local`; `.import-kb.openai.done` exists. +- Checks: exit 0; the stub was called with `--provider openai`; `.import-kb.openai.done` exists. ### 3. test_openai_without_key_exit1 @@ -99,23 +99,41 @@ When `import-knowledge` exits non-zero, no sentinel is written, so the next star - Checks: non-zero exit; the stub was called; `.import-kb.local.done` is absent. +### 11. test_asicloud_with_key_runs_import_and_writes_sentinel + +ASICloud with `ASI_API_KEY` set runs the import and leaves the ASICloud sentinel. + +- Checks: exit 0; the stub was called with `--provider asicloud`; `.import-kb.asicloud.done` exists. + +### 12. test_asicloud_without_key_exit1 + +ASICloud without a key stops before importing. + +- Checks: exit 1; stderr says `ASI_API_KEY is required`; the stub was never called. + +### 13. test_embedding_model_is_passed_to_the_import + +`EMBEDDING_MODEL` reaches the import as `--model`. + +- Checks: exit 0; the stub was called with `--provider asicloud --model BAAI/bge-base-en-v1.5`. + ## Integration: container startup They launch the image through `scripts/omega`, so the real entrypoint runs (nginx, env scrub, import gating). -### 11. test_entrypoint_imports_when_enabled +### 14. test_entrypoint_imports_when_enabled With `IMPORT_KB_ON_START=1` the entrypoint starts the import on boot. - Checks: `[import-kb] Running` appears in the container log within 180 s. -### 12. test_entrypoint_skips_when_disabled +### 15. test_entrypoint_skips_when_disabled With `IMPORT_KB_ON_START=0` the entrypoint never touches import-kb. - Checks: after a 25 s window no `[import-kb]` line appears in the log. -### 13. test_local_real_import_runs +### 16. test_local_real_import_runs A real local import runs and lands in the same `chroma_db` the agent reads from. diff --git a/Autotests/import_knowledge/test_import_knowledge.py b/Autotests/import_knowledge/test_import_knowledge.py index 29973524..df1e125b 100644 --- a/Autotests/import_knowledge/test_import_knowledge.py +++ b/Autotests/import_knowledge/test_import_knowledge.py @@ -34,7 +34,9 @@ def env(tmp_path): e["CHROMA_DB_PATH"] = str(chroma) e.pop("IMPORT_KB_FORCE", None) e.pop("OPENAI_API_KEY", None) + e.pop("ASI_API_KEY", None) e.pop("EMBEDDING_PROVIDER", None) + e.pop("EMBEDDING_MODEL", None) return {"env": e, "marker": marker, "chroma": chroma} @@ -62,7 +64,7 @@ def test_openai_with_key_runs_import_and_writes_sentinel(env): r = run(env["env"]) assert r.returncode == 0, r.stderr assert env["marker"].exists() - assert env["marker"].read_text().strip() == "" + assert env["marker"].read_text().strip() == "--provider openai" assert (env["chroma"] / ".import-kb.openai.done").exists() @@ -165,3 +167,31 @@ def test_failed_import_does_not_write_sentinel(env): assert r.returncode != 0 assert env["marker"].exists() assert not (env["chroma"] / ".import-kb.local.done").exists() + + +def test_asicloud_with_key_runs_import_and_writes_sentinel(env): + env["env"]["EMBEDDING_PROVIDER"] = "ASICloud" + env["env"]["ASI_API_KEY"] = "dummy-key" + r = run(env["env"]) + assert r.returncode == 0, r.stderr + assert env["marker"].read_text().strip() == "--provider asicloud" + assert (env["chroma"] / ".import-kb.asicloud.done").exists() + + +def test_asicloud_without_key_exit1(env): + env["env"]["EMBEDDING_PROVIDER"] = "ASICloud" + r = run(env["env"]) + assert r.returncode == 1 + assert "ASI_API_KEY is required" in r.stderr + assert not env["marker"].exists() + + +def test_embedding_model_is_passed_to_the_import(env): + env["env"]["EMBEDDING_PROVIDER"] = "ASICloud" + env["env"]["ASI_API_KEY"] = "dummy-key" + env["env"]["EMBEDDING_MODEL"] = "BAAI/bge-base-en-v1.5" + r = run(env["env"]) + assert r.returncode == 0, r.stderr + assert env["marker"].read_text().strip() == ( + "--provider asicloud --model BAAI/bge-base-en-v1.5" + ) diff --git a/Autotests/run_mandatory b/Autotests/run_mandatory index 8c76355a..1720c8c9 100644 --- a/Autotests/run_mandatory +++ b/Autotests/run_mandatory @@ -41,3 +41,4 @@ unit/test_fileio_verified_writes.py unit/test_fileio_verified_deletes.py unit/test_helper_parsing.py unit/test_openclaw_unit.py +import_knowledge/test_import_knowledge.py diff --git a/Autotests/test_openai_runtime_embeddings.py b/Autotests/test_openai_runtime_embeddings.py index 1bf6bf30..7fee0481 100644 --- a/Autotests/test_openai_runtime_embeddings.py +++ b/Autotests/test_openai_runtime_embeddings.py @@ -3,18 +3,24 @@ import types from pathlib import Path +import pytest + REPO_ROOT = Path(__file__).resolve().parents[1] RAG_MODULE_PATH = REPO_ROOT / "src" / "rag.py" MEMORY_METTA_PATH = REPO_ROOT / "src" / "memory.metta" -def load_rag_module(monkeypatch): +def load_rag_module(monkeypatch, config=None, expected_model="text-embedding-3-large", + error=None): created_clients = [] + settings = {"GATEWAY_URL": "http://gateway:8080", **(config or {})} class FakeEmbeddings: def create(self, *, model, input): - assert model == "text-embedding-3-large" + if error is not None: + raise error + assert model == expected_model assert input == ["runtime probe"] return types.SimpleNamespace( data=[types.SimpleNamespace(embedding=[0.1, 0.2, 0.3])] @@ -32,7 +38,7 @@ def __init__(self, *, base_url=None, api_key=None): chromadb_module = types.ModuleType("chromadb") config_module = types.ModuleType("config") config_module.config_get_by_key = ( - lambda key, default=None: "http://gateway:8080" if key == "GATEWAY_URL" else default + lambda key, default=None: settings.get(key, default) ) llm_module = types.ModuleType("lib_llm_ext") llm_module.initLocalEmbedding = lambda: None @@ -53,14 +59,64 @@ def __init__(self, *, base_url=None, api_key=None): def test_runtime_openai_embedding_uses_proxy_and_returns_single_vector(monkeypatch): rag, clients = load_rag_module(monkeypatch) - assert rag.openai_embed("runtime probe") == [0.1, 0.2, 0.3] + assert rag.cloud_embed("runtime probe") == [0.1, 0.2, 0.3] assert len(clients) == 1 assert clients[0].base_url == "http://gateway:8080/openai/" assert clients[0].api_key == "unused" +def test_runtime_embedding_uses_the_configured_provider_and_model(monkeypatch): + rag, clients = load_rag_module( + monkeypatch, + config={"embeddingprovider": "ASICloud", + "embeddingModel": "WhereIsAI/UAE-Large-V1"}, + expected_model="WhereIsAI/UAE-Large-V1", + ) + + assert rag.cloud_embed("runtime probe") == [0.1, 0.2, 0.3] + assert clients[0].base_url == "http://gateway:8080/asicloud/" + + +def test_runtime_asicloud_without_model_uses_the_asicloud_default(monkeypatch): + rag, clients = load_rag_module( + monkeypatch, + config={"embeddingprovider": "ASICloud"}, + expected_model="WhereIsAI/UAE-Large-V1", + ) + + assert rag.cloud_embed("runtime probe") == [0.1, 0.2, 0.3] + assert clients[0].base_url == "http://gateway:8080/asicloud/" + + +def test_runtime_empty_model_falls_back_to_the_provider_default(monkeypatch): + rag, _ = load_rag_module( + monkeypatch, + config={"embeddingprovider": "ASICloud", "embeddingModel": ""}, + expected_model="WhereIsAI/UAE-Large-V1", + ) + + assert rag.cloud_embed("runtime probe") == [0.1, 0.2, 0.3] + + +def test_runtime_embedding_failure_logs_the_provider_error(monkeypatch): + rag, _ = load_rag_module( + monkeypatch, + config={"embeddingprovider": "ASICloud"}, + error=Exception("Error code: 400 - {'error': 'Model not found'}"), + ) + logged = [] + rag.logger = types.SimpleNamespace( + error=lambda message, *args, **kwargs: logged.append(message) + ) + + with pytest.raises(RuntimeError): + rag.cloud_embed("runtime probe") + + assert any("Model not found" in message for message in logged) + + def test_memory_metta_routes_openai_embeddings_to_rag_wrapper(): memory_metta = MEMORY_METTA_PATH.read_text(encoding="utf-8") - assert "(py-call (rag.openai_embed (string-safe $str)))" in memory_metta + assert "(py-call (rag.cloud_embed (string-safe $str)))" in memory_metta assert "useGPTEmbedding" not in memory_metta diff --git a/README.md b/README.md index 0886643b..a56b554f 100644 --- a/README.md +++ b/README.md @@ -133,7 +133,7 @@ After start go to https://webchat.quakenet.org/ to communicate with the agent. J If you are running Omega without Docker and would like to load it with preset knowledge, follow these steps: -1. Set EMBEDDING_PROVIDER in your environment. It can be set to either OpenAI or Local. OpenAI embeddings also require OPENAI_API_KEY to be set in your environment. +1. Set EMBEDDING_PROVIDER in your environment. It can be set to OpenAI, ASICloud or Local. OpenAI embeddings also require OPENAI_API_KEY, and ASICloud embeddings require ASI_API_KEY to be set in your environment. 2. Run: ``` diff --git a/config/config.yaml b/config/config.yaml index 9d74a649..3009bb3d 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -36,8 +36,11 @@ maxEpisodeRecallLines: 20 maxHistory: 30000 # ChromaDB persistence directory used for long-term memory portability. chromaDbPath: "./chroma_db" -# `Local` (Python-side model) or `OpenAI` (requires `OPENAI_API_KEY`) +# `Local` (Python-side model) or the id of a provider serving OpenAI-compatible embeddings: `OpenAI`, `ASICloud` embeddingprovider: Local +# Model asked of a non-`Local` embeddingprovider; empty means the provider default: +# `text-embedding-3-large` for `OpenAI`, `WhereIsAI/UAE-Large-V1` for `ASICloud` +embeddingModel: "" # Enable authenticated operator-triggered /memory-export commands (disabled by default). memoryExportEnabled: false diff --git a/docs/reference-configuration.md b/docs/reference-configuration.md index 0b9c6262..fa30be30 100644 --- a/docs/reference-configuration.md +++ b/docs/reference-configuration.md @@ -33,7 +33,8 @@ This reads a command-line override via `argk` (`name=value` on the MeTTa command | `maxHistory` | 30000 (chars) | Tail of `memory/history.metta` included in the prompt. | | `memoryDirectory` | `./repos/Omega/memory` | Directory containing persistent memory files such as `history.metta`. | | `chromaDbPath` | `./chroma_db` | ChromaDB persistence directory used for memory backup and restore. | -| `embeddingprovider` | `Local` | `Local` (Python-side model) or `OpenAI`. | +| `embeddingprovider` | `Local` | `Local` (Python-side model), or the id of a provider that serves an OpenAI-compatible `/embeddings` endpoint — `OpenAI` and `ASICloud` are known to. The gateway supplies that provider's key. | +| `embeddingModel` | empty | Model asked of a non-`Local` `embeddingprovider`. Empty means the provider default: `text-embedding-3-large` for `OpenAI`, `WhereIsAI/UAE-Large-V1` for `ASICloud`. | ## Channels (`src/channels.metta`, `initChannels`) diff --git a/docs/reference-internals-extension-points.md b/docs/reference-internals-extension-points.md index 9aaab3d8..06c7817e 100644 --- a/docs/reference-internals-extension-points.md +++ b/docs/reference-internals-extension-points.md @@ -53,10 +53,14 @@ In `src/memory.metta`, the `embed` function dispatches on `embeddingprovider`: (= (embed $str) (if (== (embeddingprovider) Local) (py-call (lib_llm_ext.useLocalEmbedding (string-safe $str))) - (useGPTEmbedding (string-safe $str)))) + (py-call (rag.cloud_embed (string-safe $str))))) ``` -To add a new backend, add a branch and implement the Python function. +Any value other than `Local` is a provider id: the remote branch posts +`embeddingModel` to `//`, the +location that already injects that provider's key. Switching vendor is +configuration, not code — provided the vendor serves embeddings at all. It +changes the vector space, so reset the ChromaDB store when you do. ## Change the reasoning library diff --git a/docs/reference-skills-memory.md b/docs/reference-skills-memory.md index 2658f8fb..f55940d9 100644 --- a/docs/reference-skills-memory.md +++ b/docs/reference-skills-memory.md @@ -30,7 +30,7 @@ The result of the ChromaDB write (internally). The agent treats a successful cal ### Notes / Limits - Text is passed through `string-safe` before embedding, which escapes newlines, quotes, and apostrophes. -- Embedding provider is selected by `embeddingprovider` (`Local` or `OpenAI`). +- Embedding provider is selected by `embeddingprovider`, the model by `embeddingModel`. - Nothing deduplicates automatically — repeated `remember` calls store multiple items. --- diff --git a/docs/tutorial-01-teaching-memories.md b/docs/tutorial-01-teaching-memories.md index 330fa9cb..adaa00ff 100644 --- a/docs/tutorial-01-teaching-memories.md +++ b/docs/tutorial-01-teaching-memories.md @@ -65,7 +65,7 @@ From `src/memory.metta`: - `maxRecallItems` — how many items `query` returns (default 20). - `maxEpisodeRecallLines` — how many lines `episodes` returns (default 20). - `maxHistory` — characters of history fed back into the prompt (default 30000). -- `embeddingprovider` — `OpenAI` or `Local`. +- `embeddingprovider` — `OpenAI`, `ASICloud` or `Local`. Change any of these by editing the `configure` calls in `initMemory` or passing command-line overrides — see [reference-configuration.md](./reference-configuration.md). diff --git a/entrypoint.sh b/entrypoint.sh index a89df4dc..d5c4dda4 100644 --- a/entrypoint.sh +++ b/entrypoint.sh @@ -11,6 +11,7 @@ nginx_url() { cd /PeTTa EMBEDDING_PROVIDER="${EMBEDDING_PROVIDER:-Local}" +EMBEDDING_MODEL="${EMBEDDING_MODEL:-}" OPENAIAPI_URL="http://localhost:8080/" # dummy value MM_URL="http://localhost:8080/" # dummy value OPENCLAW_URL="http://localhost:8080/" # dummy value @@ -18,6 +19,9 @@ for arg in "$@"; do if [[ "$arg" == embeddingprovider=* ]]; then EMBEDDING_PROVIDER="${arg#*=}" fi + if [[ "$arg" == embeddingModel=* ]]; then + EMBEDDING_MODEL="${arg#*=}" + fi # URL to redirect OpenAIAPI provider requests if [[ "$arg" == openaiapi_url=* ]]; then OPENAIAPI_URL=$(nginx_url "${arg#*=}") @@ -31,7 +35,7 @@ for arg in "$@"; do OPENCLAW_URL=$(nginx_url "${arg#*=}") fi done -export EMBEDDING_PROVIDER OPENAIAPI_URL MM_URL OPENCLAW_URL +export EMBEDDING_PROVIDER EMBEDDING_MODEL OPENAIAPI_URL MM_URL OPENCLAW_URL su www-data -s /bin/sh -c "sh /opt/nginx/nginx.sh" @@ -41,11 +45,12 @@ if [[ "${IMPORT_KB_ON_START}" == "1" ]]; then fi MEMORY_PORTABILITY_PYTHON='import os +import sys from config import init_config from memory_export import create_memory_store from memory_portability import MemoryTransfer -init_config([]) +init_config(sys.argv[1:]) transfer = MemoryTransfer( transfer_dir="/memory-transfer", store=create_memory_store(), @@ -66,13 +71,13 @@ export MEMORY_PORTABILITY_PYTHON export PYTHONPATH="${OMEGA_DIR}:${OMEGA_DIR}/src${PYTHONPATH:+:${PYTHONPATH}}" export MEMORY_PORTABILITY_OPERATION=recover -su nobody -s /bin/sh -c 'exec python3 -c "$MEMORY_PORTABILITY_PYTHON"' \ +su nobody -s /bin/sh -c 'exec python3 -c "$MEMORY_PORTABILITY_PYTHON" "$@"' sh "$@" \ || { echo "Memory import recovery failed. Aborting startup." >&2; exit 1; } if [[ -n "${MEMORY_IMPORT_FILE:-}" ]]; then echo "memory_portability: importing ${MEMORY_IMPORT_FILE}" export MEMORY_PORTABILITY_OPERATION=import - su nobody -s /bin/sh -c 'exec python3 -c "$MEMORY_PORTABILITY_PYTHON"' \ + su nobody -s /bin/sh -c 'exec python3 -c "$MEMORY_PORTABILITY_PYTHON" "$@"' sh "$@" \ || { echo "Memory import failed. Aborting startup." >&2; exit 1; } echo "memory_portability: import complete" fi diff --git a/requirements.txt b/requirements.txt index 45e17674..19f8bd16 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ chromadb==1.5.9 openai==2.38.0 transformers==5.8.0 sentence-transformers==5.5.1 -import-kb==0.2.3 +import-kb==0.2.4 py-landlock==0.1.1 pyyaml==6.0.3 ddgs==9.14.4 diff --git a/scripts/import_knowledge.sh b/scripts/import_knowledge.sh index 42af4a09..ef7c5685 100644 --- a/scripts/import_knowledge.sh +++ b/scripts/import_knowledge.sh @@ -13,20 +13,29 @@ mkdir -p "${CHROMA_DB_PATH}" PROVIDER="$(normalize_provider "${EMBEDDING_PROVIDER}")" case "${PROVIDER}" in - openai) - if [[ -z "${OPENAI_API_KEY:-}" ]]; then - echo "ERROR: OPENAI_API_KEY is required when EMBEDDING_PROVIDER=OpenAI." >&2 + openai|asicloud) + case "${PROVIDER}" in + openai) API_KEY_VAR="OPENAI_API_KEY" ;; + asicloud) API_KEY_VAR="ASI_API_KEY" ;; + esac + + if [[ -z "${!API_KEY_VAR:-}" ]]; then + echo "ERROR: ${API_KEY_VAR} is required when EMBEDDING_PROVIDER=${EMBEDDING_PROVIDER}." >&2 exit 1 fi - SENTINEL="${CHROMA_DB_PATH}/.import-kb.openai.done" + SENTINEL="${CHROMA_DB_PATH}/.import-kb.${PROVIDER}.done" if [[ -f "${SENTINEL}" && "${IMPORT_KB_FORCE}" != "1" ]]; then - echo "[import-kb] Already initialized with OpenAI embeddings; skipping." + echo "[import-kb] Already initialized with ${PROVIDER} embeddings; skipping." else - echo "[import-kb] Running import-knowledge with OpenAI embeddings." + echo "[import-kb] Running import-knowledge with ${PROVIDER} embeddings." echo "[import-kb] CHROMA_DB_PATH=${CHROMA_DB_PATH}" - import-knowledge + if [[ -n "${EMBEDDING_MODEL:-}" ]]; then + import-knowledge --provider "${PROVIDER}" --model "${EMBEDDING_MODEL}" + else + import-knowledge --provider "${PROVIDER}" + fi date -Iseconds > "${SENTINEL}" echo "[import-kb] Import complete." fi @@ -48,7 +57,7 @@ case "${PROVIDER}" in *) echo "ERROR: Unsupported embeddingprovider='${EMBEDDING_PROVIDER}'." >&2 - echo "Use embeddingprovider=OpenAI or embeddingprovider=Local." >&2 + echo "Use embeddingprovider=OpenAI, embeddingprovider=ASICloud or embeddingprovider=Local." >&2 exit 1 ;; esac diff --git a/src/embedding_models.py b/src/embedding_models.py new file mode 100644 index 00000000..cb9a6650 --- /dev/null +++ b/src/embedding_models.py @@ -0,0 +1,16 @@ +DEFAULT_MODELS = { + "openai": "text-embedding-3-large", + "asicloud": "WhereIsAI/UAE-Large-V1", +} + +DIMENSIONS = { + "text-embedding-3-large": 3072, + "WhereIsAI/UAE-Large-V1": 1024, +} + + +def embedding_model(provider, configured=None): + configured = str(configured or "").strip() + if configured: + return configured + return DEFAULT_MODELS.get(str(provider).casefold(), DEFAULT_MODELS["openai"]) diff --git a/src/loop.metta b/src/loop.metta index 783b2daa..ed51f391 100644 --- a/src/loop.metta +++ b/src/loop.metta @@ -26,9 +26,8 @@ (= (initKnowledge) (progn (log INFO "loop" "Initializing knowledge base") - (if (== (embeddingprovider) OpenAI) - (log INFO "loop" (py-call (rag.init_knowledge "OpenAI"))) - (log INFO "loop" (py-call (rag.init_knowledge "Local")))))) + (log INFO "loop" (py-call (rag.init_knowledge (embeddingprovider)))) + ())) (= (getContext) (string-safe (py-str ("PROMPT: " (getPrompt (provider)) " SKILLS: " (getSkills) diff --git a/src/memory.metta b/src/memory.metta index 3f5fae36..c7d18698 100644 --- a/src/memory.metta +++ b/src/memory.metta @@ -11,7 +11,7 @@ (configure maxRecallItems 20) (configure maxEpisodeRecallLines 20) (configure maxHistory 30000) - (configure embeddingprovider Local) ;OpenAI or Local + (configure embeddingprovider Local) ;OpenAI, ASICloud or Local (if (== (embeddingprovider) Local) (py-call (lib_llm_ext.initLocalEmbedding)) _))) @@ -45,7 +45,7 @@ (= (embed $str) (if (== (embeddingprovider) Local) (py-call (lib_llm_ext.useLocalEmbedding (string-safe $str))) - (py-call (rag.openai_embed (string-safe $str))))) + (py-call (rag.cloud_embed (string-safe $str))))) (= (appendToHistory $addition) (append-file-raw (library Omega ./memory/history.metta) (swrite $addition))) diff --git a/src/memory_export.py b/src/memory_export.py index f914ee2b..65177875 100644 --- a/src/memory_export.py +++ b/src/memory_export.py @@ -4,6 +4,7 @@ from pathlib import Path from config import config_get_by_key +from embedding_models import DEFAULT_MODELS, DIMENSIONS, embedding_model from helper import omega_version, projectRootDirectory from src.logger import get_logger @@ -33,6 +34,31 @@ def _resolve_chroma_path() -> Path: return Path(str(configured)).expanduser().resolve() +def _cloud_embedder(provider, model): + def embed(texts): + from import_knowledge.import_knowledge import embed_batch, init_embeddings + + init_embeddings(mode=provider.casefold(), model_name=model) + return embed_batch(texts) + + return embed + + +def _embedding_options(): + provider = str(config_get_by_key("embeddingprovider", "Local")).strip() + if provider.casefold() == "local": + return {} + model = embedding_model(provider, config_get_by_key("embeddingModel", "")) + return { + "embed_batch": _cloud_embedder(provider, model), + "embedding_profile": { + "provider": provider, + "model": model, + "vector_dimension": DIMENSIONS.get(model), + }, + } + + def create_memory_store(): """Build an import-kb store from Omega's effective configuration.""" from memory_portability.storage import MemoryStore @@ -41,6 +67,7 @@ def create_memory_store(): memory_dir=_resolve_memory_dir(), chroma_path=_resolve_chroma_path(), collection_name="memories", + **_embedding_options(), ) @@ -50,7 +77,7 @@ def _get_transfer(): from memory_portability import MemoryTransfer embedding_provider = str(config_get_by_key("embeddingprovider", "Local")).strip() - if embedding_provider.casefold() not in {"local", "openai"}: + if embedding_provider.casefold() not in {"local", *DEFAULT_MODELS}: raise ValueError(f"Unsupported embedding provider: {embedding_provider!r}") os.environ["EMBEDDING_PROVIDER"] = embedding_provider diff --git a/src/rag.py b/src/rag.py index f9274fd6..471540c3 100644 --- a/src/rag.py +++ b/src/rag.py @@ -7,12 +7,12 @@ from lib_llm_ext import initLocalEmbedding, useLocalEmbedding from src.logger import get_logger from config import config_get_by_key +from embedding_models import embedding_model logger = get_logger(__name__) # --- Constants ----------------------------------------------------------- -EMBEDDING_MODEL = "text-embedding-3-large" COLLECTION_NAME = "memories" TOP_K = 5 MIN_CHUNK_CHARS = 100 @@ -138,23 +138,26 @@ def _chunk_markdown(text, filename): # --- Embedding ----------------------------------------------------------- -def openai_embed_batch(texts): - """Embed a list of texts via OpenAI. Returns list of float vectors.""" +def cloud_embed_batch(texts): + """Embed a list of texts via an OpenAI-compatible API. Returns list of float vectors.""" + provider = str(config_get_by_key("embeddingprovider", "OpenAI")) + model = embedding_model(provider, config_get_by_key("embeddingModel", "")) proxy_url = config_get_by_key("GATEWAY_URL") if proxy_url: - client = openai.OpenAI(base_url=f"{proxy_url.rstrip('/')}/openai/", api_key="unused") + client = openai.OpenAI(base_url=f"{proxy_url.rstrip('/')}/{provider.lower()}/", api_key="unused") else: client = openai.OpenAI() try: - resp = client.embeddings.create(model=EMBEDDING_MODEL, input=texts) + resp = client.embeddings.create(model=model, input=texts) except Exception as e: + logger.error(f"Embedding request failed: provider={provider} model={model}: {e}") raise RuntimeError(f"Embedding request failed: {e}") from e return [item.embedding for item in resp.data] -def openai_embed(text): +def cloud_embed(text): """Embed one runtime memory string via the configured OpenAI route.""" - return openai_embed_batch([text])[0] + return cloud_embed_batch([text])[0] def local_embed_batch(texts): @@ -253,13 +256,10 @@ def init_knowledge(embedding_selection): texts = [c["text"] for c in chunks] if embedding_selection == "Local": embeddings = local_embed_batch(texts) - elif embedding_selection == "OpenAI": - embeddings = openai_embed_batch(texts) + elif embedding_selection.lower() in ("asicloud", "openai"): + embeddings = cloud_embed_batch(texts) else: - raise ValueError( - f"Invalid embedding_selection={embedding_selection!r}. " - "Expected 'Local' or 'OpenAI'.") - + raise NotImplementedError(f"Unsupported embedding provider: {embedding_selection}") if not embeddings: logger.warning(f"{filename}: embedding failed, skipping") continue diff --git a/tests/test_embedding_models.py b/tests/test_embedding_models.py new file mode 100644 index 00000000..c34d4192 --- /dev/null +++ b/tests/test_embedding_models.py @@ -0,0 +1,38 @@ +import pytest + +import config +import embedding_models + + +@pytest.fixture(autouse=True) +def isolated_config(monkeypatch): + for name in ("_CONFIG", "_COMMAND_LINE", "_CONFIG_FILE"): + monkeypatch.setattr(config, name, getattr(config, name)) + + +@pytest.mark.parametrize( + "provider, model", + [ + ("ASICloud", "WhereIsAI/UAE-Large-V1"), + ("OpenAI", "text-embedding-3-large"), + ], +) +def test_shipped_config_resolves_to_the_provider_default(provider, model): + config.init_config([f"embeddingprovider={provider}"]) + + configured = config.config_get_by_key("embeddingModel", "") + + assert embedding_models.embedding_model(provider, configured) == model + + +def test_explicit_model_wins_over_the_provider_default(): + assert embedding_models.embedding_model( + "ASICloud", "BAAI/bge-base-en-v1.5" + ) == "BAAI/bge-base-en-v1.5" + + +def test_provider_defaults_match_import_kb(): + import_knowledge = pytest.importorskip("import_knowledge.import_knowledge") + + for provider, model in embedding_models.DEFAULT_MODELS.items(): + assert import_knowledge.PROVIDERS[provider]["default_model"] == model diff --git a/tests/test_memory_export.py b/tests/test_memory_export.py index feedf891..dd47d95b 100644 --- a/tests/test_memory_export.py +++ b/tests/test_memory_export.py @@ -185,6 +185,11 @@ def __init__(self, **kwargs): chroma_path = tmp_path / "custom-chroma" monkeypatch.setattr(handler, "_resolve_memory_dir", lambda: memory_dir) monkeypatch.setattr(handler, "_resolve_chroma_path", lambda: chroma_path) + monkeypatch.setattr( + handler, + "config_get_by_key", + lambda key, default=None: "Local" if key == "embeddingprovider" else default, + ) store = handler.create_memory_store() @@ -198,6 +203,183 @@ def __init__(self, **kwargs): ] +def install_fake_memory_store(monkeypatch, tmp_path): + created_stores = [] + + class FakeStore: + def __init__(self, **kwargs): + created_stores.append(kwargs) + + package = types.ModuleType("memory_portability") + package.__path__ = [] + storage = types.ModuleType("memory_portability.storage") + storage.MemoryStore = FakeStore + monkeypatch.setitem(sys.modules, "memory_portability", package) + monkeypatch.setitem(sys.modules, "memory_portability.storage", storage) + return created_stores + + +def install_fake_import_kb(monkeypatch): + calls = [] + package = types.ModuleType("import_knowledge") + package.__path__ = [] + module = types.ModuleType("import_knowledge.import_knowledge") + module.init_embeddings = lambda mode, model_name=None: calls.append( + ("init", mode, model_name) + ) + module.embed_batch = lambda texts: calls.append(("embed", list(texts))) or [ + [0.5] * 3 for _ in texts + ] + monkeypatch.setitem(sys.modules, "import_knowledge", package) + monkeypatch.setitem(sys.modules, "import_knowledge.import_knowledge", module) + return calls + + +def configure(monkeypatch, handler, tmp_path, provider, config_model="", + env_provider="Local", env_model="env-only-model"): + monkeypatch.setenv("EMBEDDING_PROVIDER", env_provider) + monkeypatch.setenv("EMBEDDING_MODEL", env_model) + settings = {"embeddingprovider": provider, "embeddingModel": config_model} + monkeypatch.setattr( + handler, + "config_get_by_key", + lambda key, default=None: settings.get(key, default), + ) + monkeypatch.setattr(handler, "_resolve_memory_dir", lambda: tmp_path / "memory") + monkeypatch.setattr(handler, "_resolve_chroma_path", lambda: tmp_path / "chroma") + + +def test_asicloud_memory_store_embeds_with_asicloud_default_model( + handler, + monkeypatch, + tmp_path, +): + stores = install_fake_memory_store(monkeypatch, tmp_path) + calls = install_fake_import_kb(monkeypatch) + configure(monkeypatch, handler, tmp_path, "ASICloud") + + handler.create_memory_store() + + assert stores[0]["embedding_profile"] == { + "provider": "ASICloud", + "model": "WhereIsAI/UAE-Large-V1", + "vector_dimension": 1024, + } + assert stores[0]["embed_batch"](["imported fact"]) == [[0.5, 0.5, 0.5]] + assert calls == [ + ("init", "asicloud", "WhereIsAI/UAE-Large-V1"), + ("embed", ["imported fact"]), + ] + + +def test_memory_store_uses_the_configured_model( + handler, + monkeypatch, + tmp_path, +): + stores = install_fake_memory_store(monkeypatch, tmp_path) + calls = install_fake_import_kb(monkeypatch) + configure( + monkeypatch, handler, tmp_path, "ASICloud", config_model="BAAI/bge-base-en-v1.5" + ) + + handler.create_memory_store() + stores[0]["embed_batch"](["imported fact"]) + + assert stores[0]["embedding_profile"] == { + "provider": "ASICloud", + "model": "BAAI/bge-base-en-v1.5", + "vector_dimension": None, + } + assert calls[0] == ("init", "asicloud", "BAAI/bge-base-en-v1.5") + + +def test_memory_store_uses_the_runtime_model_inside_the_agent( + handler, + monkeypatch, + tmp_path, +): + stores = install_fake_memory_store(monkeypatch, tmp_path) + install_fake_import_kb(monkeypatch) + configure( + monkeypatch, + handler, + tmp_path, + "OpenAI", + config_model="text-embedding-3-small", + ) + + handler.create_memory_store() + + assert stores[0]["embedding_profile"]["provider"] == "OpenAI" + assert stores[0]["embedding_profile"]["model"] == "text-embedding-3-small" + + +def test_memory_store_ignores_embedding_environment_variables( + handler, + monkeypatch, + tmp_path, +): + stores = install_fake_memory_store(monkeypatch, tmp_path) + install_fake_import_kb(monkeypatch) + configure(monkeypatch, handler, tmp_path, "Local", env_provider="ASICloud") + + handler.create_memory_store() + + assert stores == [ + { + "memory_dir": tmp_path / "memory", + "chroma_path": tmp_path / "chroma", + "collection_name": "memories", + } + ] + + +def test_entrypoint_passes_container_arguments_to_memory_portability(): + entrypoint = (REPO_ROOT / "entrypoint.sh").read_text(encoding="utf-8") + + assert "init_config(sys.argv[1:])" in entrypoint + assert "init_config([])" not in entrypoint + assert entrypoint.count( + """su nobody -s /bin/sh -c 'exec python3 -c "$MEMORY_PORTABILITY_PYTHON" "$@"' sh "$@\"""" + ) == 2 + + +def test_export_is_allowed_for_asicloud_embeddings(handler, monkeypatch): + created = [] + package = types.ModuleType("memory_portability") + package.MemoryTransfer = lambda **kwargs: created.append(kwargs) or "transfer" + monkeypatch.setitem(sys.modules, "memory_portability", package) + monkeypatch.setenv("EMBEDDING_PROVIDER", "Local") + monkeypatch.setenv("OMEGA_VERSION", "unset") + monkeypatch.setattr(handler, "omega_version", lambda: "Omega version=test") + monkeypatch.setattr(handler, "create_memory_store", lambda: "configured-store") + monkeypatch.setattr( + handler, + "config_get_by_key", + lambda key, default=None: "ASICloud" if key == "embeddingprovider" else default, + ) + handler._transfer = None + + assert handler._get_transfer() == "transfer" + assert os.environ["EMBEDDING_PROVIDER"] == "ASICloud" + + +def test_export_rejects_providers_without_embeddings(handler, monkeypatch): + package = types.ModuleType("memory_portability") + package.MemoryTransfer = lambda **kwargs: pytest.fail("transfer must not start") + monkeypatch.setitem(sys.modules, "memory_portability", package) + monkeypatch.setattr( + handler, + "config_get_by_key", + lambda key, default=None: "Anthropic" if key == "embeddingprovider" else default, + ) + handler._transfer = None + + with pytest.raises(ValueError, match="Anthropic"): + handler._get_transfer() + + def test_storage_paths_are_resolved_from_omega_config( handler, monkeypatch,