From d1cffb2f40b3aabea2a6afa002f5f3cd5883ced0 Mon Sep 17 00:00:00 2001 From: KiriAky 107 Date: Fri, 4 Sep 2026 19:33:57 +0800 Subject: [PATCH] =?UTF-8?q?fix(embedding):=20=E4=BC=A0=E9=80=92=E6=9C=AC?= =?UTF-8?q?=E5=9C=B0=E7=B4=A2=E5=BC=95=E9=99=90=E5=88=B6=E5=B9=B6=E5=86=BB?= =?UTF-8?q?=E7=BB=93=E6=8E=A8=E7=90=86=E9=85=8D=E7=BD=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/knowledge/parser.py | 2 ++ backend/app/local_models/runtime.py | 17 +++++++++-- backend/app/providers/routing.py | 12 ++++---- backend/app/retrieval/routed_vectors.py | 6 ++-- backend/app/services/media_notes.py | 3 ++ backend/app/services/note_service.py | 4 +-- backend/tests/test_media_jobs.py | 26 ++++++++++++++++ backend/tests/test_model_routing.py | 40 +++++++++++++++++++++++++ 8 files changed, 97 insertions(+), 13 deletions(-) diff --git a/backend/app/knowledge/parser.py b/backend/app/knowledge/parser.py index 268525a..d3217f5 100644 --- a/backend/app/knowledge/parser.py +++ b/backend/app/knowledge/parser.py @@ -31,6 +31,7 @@ class ParsedNote: created_at: datetime updated_at: datetime blocks: list[NoteBlock] = field(default_factory=list) + embedding_local_only: bool = False def note_id_for_path(rel_path: str) -> str: @@ -69,6 +70,7 @@ def parse_note( created_at=created_at, updated_at=updated_at, blocks=blocks, + embedding_local_only=str(frontmatter.get("embedding_local_only", "")).lower() == "true", ) diff --git a/backend/app/local_models/runtime.py b/backend/app/local_models/runtime.py index 01f93fc..f31e6ed 100644 --- a/backend/app/local_models/runtime.py +++ b/backend/app/local_models/runtime.py @@ -165,21 +165,32 @@ runtime = Runtime() class LocalEmbedding: dim = 384 + def __init__(self, config=None): + self._config = config + + def snapshot(self): + return LocalEmbedding((self._config or configuration()).model_copy(deep=True)) + @property def model_id(self): - spec = CATALOG[configuration().embedding_model] + spec = CATALOG[(self._config or configuration()).embedding_model] return f"{spec.repository}@{spec.revision}" @property def version(self): - return CATALOG[configuration().embedding_model].revision + return CATALOG[(self._config or configuration()).embedding_model].revision @property def available(self): return read_state(configuration().embedding_model)["status"] == "installed" and interpreter().is_file() async def embed_documents(self, texts): - return await runtime.infer(configuration().embedding_model, "embedding", {"texts": texts}, priority=0) + config = (self._config or configuration()).model_copy(deep=True) + token = runtime_context.set(config) + try: + return await runtime.infer(config.embedding_model, "embedding", {"texts": texts}, priority=0) + finally: + runtime_context.reset(token) async def embed_query(self, query): return (await self.embed_documents([query]))[0] diff --git a/backend/app/providers/routing.py b/backend/app/providers/routing.py index 580da01..dcd009e 100644 --- a/backend/app/providers/routing.py +++ b/backend/app/providers/routing.py @@ -206,9 +206,9 @@ class ModelRoutingService: raise invalid_response() return data, url - async def embed(self, texts: list[str]) -> EmbeddingResult: + async def embed(self, texts: list[str], *, local_only=False) -> EmbeddingResult: config = self.configuration() - binding = config.embedding + binding = None if local_only else config.embedding record_embedding(route_version=config.version, requested_route=binding.model_dump() if binding else None) reason = None @@ -255,12 +255,14 @@ class ModelRoutingService: model_id="api-" + hashlib.sha256(identity.encode()).hexdigest()) except ProviderError as exc: reason = exc.code + from app.local_models.runtime import LocalEmbedding + local_embedding = self.local_embedding.snapshot() if isinstance(self.local_embedding, LocalEmbedding) else self.local_embedding try: - vectors = await self.local_embedding.embed_documents(texts) + vectors = await local_embedding.embed_documents(texts) except ProviderError as exc: raise ApiError(503, exc.code, exc.message, {"fallback_reason": reason}) from exc - return EmbeddingResult(vectors=vectors, source="local", model_id=self.local_embedding.model_id, - dimensions=self.local_embedding.dim, fallback_reason=reason) + return EmbeddingResult(vectors=vectors, source="local", model_id=local_embedding.model_id, + dimensions=local_embedding.dim, fallback_reason=reason) @staticmethod def _media_file(path: Path): diff --git a/backend/app/retrieval/routed_vectors.py b/backend/app/retrieval/routed_vectors.py index 12286af..666526c 100644 --- a/backend/app/retrieval/routed_vectors.py +++ b/backend/app/retrieval/routed_vectors.py @@ -35,7 +35,7 @@ class EmbeddingResult(Protocol): class EmbeddingRuntime(Protocol): - async def embed(self, texts: list[str]) -> EmbeddingResult: ... + async def embed(self, texts: list[str], *, local_only=False) -> EmbeddingResult: ... @dataclass(frozen=True) @@ -69,7 +69,7 @@ def _unit_vector(vector: list[float], dimensions: int) -> list[float]: return [value / norm for value in scaled] -async def embed_remote(texts: list[str], *, accept_local=False, strict=False) -> RemoteEmbeddings | None: +async def embed_remote(texts: list[str], *, accept_local=False, strict=False, local_only=False) -> RemoteEmbeddings | None: """Return validated API vectors, or None to use the caller's local baseline. Do not use the runtime's local result: the caller may have injected its own @@ -83,7 +83,7 @@ async def embed_remote(texts: list[str], *, accept_local=False, strict=False) -> if strict: raise ApiError(503, "EMBEDDING_UNAVAILABLE", "Embedding 服务未就绪，请检查模型路由和本地运行环境。") return None - result = await runtime.embed(texts) + result = await runtime.embed(texts, local_only=True) if local_only else await runtime.embed(texts) if result.source != "api" and not accept_local: record_embedding(fallback_reason=result.fallback_reason) return None diff --git a/backend/app/services/media_notes.py b/backend/app/services/media_notes.py index cf0f938..e2ec13e 100644 --- a/backend/app/services/media_notes.py +++ b/backend/app/services/media_notes.py @@ -41,6 +41,9 @@ async def create_transcript_note(job_id, options): lines.append("") else: lines.append(job.text or "") + if job.local_only: + # Persist the indexing policy in the Vault, including later rebuilds. + lines = ["---", "embedding_local_only: true", "---", "", *lines] try: note = await note_service.create_note(title=title, markdown="\n".join(lines), folder=options.folder, tags=["转写"]) except ApiError as exc: diff --git a/backend/app/services/note_service.py b/backend/app/services/note_service.py index 58f9e40..3c7da5c 100644 --- a/backend/app/services/note_service.py +++ b/backend/app/services/note_service.py @@ -82,10 +82,10 @@ async def prepare_note_index(parsed: ParsedNote, *, strict=False) -> PreparedInd texts = [block.content for block in parsed.blocks] if isinstance(embedding, LocalEmbedding): # One routed invocation: API first, validated local fallback. No hash vectors. - remote = await routed_vectors.embed_remote(texts, accept_local=True, strict=strict) + remote = await routed_vectors.embed_remote(texts, accept_local=True, strict=strict, local_only=parsed.embedding_local_only) return [], remote vectors = await embedding.embed_documents(texts) - remote = await routed_vectors.embed_remote(texts) + remote = await routed_vectors.embed_remote(texts, local_only=parsed.embedding_local_only) return vectors, remote diff --git a/backend/tests/test_media_jobs.py b/backend/tests/test_media_jobs.py index cb085c2..20ecfa4 100644 --- a/backend/tests/test_media_jobs.py +++ b/backend/tests/test_media_jobs.py @@ -102,3 +102,29 @@ def test_terminology_export_and_privacy_cleanup(): assert cleaned['text'] is None and cleaned['original_text'] is None and cleaned['corrections'] == [] assert client.post(f'/api/media/transcriptions/{job_id}/retry').status_code == 409 assert client.get('/api/media/attachments/lecture.txt').status_code == 404 + + +def test_local_only_export_and_rebuild_keep_local_embedding_policy(monkeypatch): + from types import SimpleNamespace + from app.contracts import TranscriptNoteRequest, IndexRebuildRequest + from app.local_models.runtime import LocalEmbedding + from app.retrieval import routed_vectors + from app.services import note_service, index_service + from app.services.media_notes import create_transcript_note + calls = [] + class Routing: + async def embed(self, texts, *, local_only=False): + calls.append(local_only) + assert local_only + return SimpleNamespace(source='local', model_id='local-test', dimensions=2, + vectors=[[1.0, 0.0] for _ in texts], fallback_reason=None) + monkeypatch.setattr(routed_vectors, 'get_model_routing', lambda: Routing()) + monkeypatch.setattr(note_service, 'embedding', LocalEmbedding()) + text_attachment() + async def scenario(): + job = await jobs.create_transcription('lecture.txt', local_only=True) + note = await create_transcript_note(job.job_id, TranscriptNoteRequest(title='Private')) + assert note.markdown.startswith('---\nembedding_local_only: true\n---') + await index_service.rebuild(IndexRebuildRequest()) + assert len(calls) >= 2 and all(calls) + asyncio.run(scenario()) diff --git a/backend/tests/test_model_routing.py b/backend/tests/test_model_routing.py index 1205e5b..985017b 100644 --- a/backend/tests/test_model_routing.py +++ b/backend/tests/test_model_routing.py @@ -678,3 +678,43 @@ def test_remote_segments_are_validated_and_local_only_skips_api(rig, audio): count = len(rig.requests) result = run(rig.service.transcribe(audio[0], "zh", local_only=True)) assert result.source == "local" and len(rig.requests) == count + + +def test_embedding_local_only_does_not_change_normal_api_fallback(rig): + bind(rig) + result = run(rig.service.embed(['private'], local_only=True)) + assert result.source == 'local' and result.fallback_reason is None + assert rig.requests == [] and rig.credentials.calls == [] + rig.http.handler = lambda request: response({'data': [{'index': 0, 'embedding': [1, 0, 0]}]}) + assert run(rig.service.embed(['normal'])).source == 'api' + rig.http.handler = lambda request: response({}, status=503) + result = run(rig.service.embed(['fallback'])) + assert result.source == 'local' and result.fallback_reason + + +@pytest.mark.parametrize('api_failure', [False, True]) +def test_local_embedding_identity_and_device_are_frozen_during_inference(rig, monkeypatch, api_failure): + import app.local_models.runtime as module + config = module.RuntimeConfig(embedding_model='bekko') + monkeypatch.setattr(module, 'configuration', lambda: module.runtime_context.get() or config) + calls = [] + async def infer(key, *args, **kwargs): + calls.append(key) + config.embedding_model = 'granite' + config.device = 'cuda' + await asyncio.sleep(0) + assert module.configuration().embedding_model == key + assert module.configuration().device == ('cpu' if len(calls) == 1 else 'cuda') + return [[1.0] + [0.0] * 383] + monkeypatch.setattr(module.runtime, 'infer', infer) + rig.service.local_embedding = module.LocalEmbedding() + if api_failure: + bind(rig) + rig.http.handler = lambda request: response({}, status=503) + first = run(rig.service.embed(['first'])) + assert 'bekko' in first.model_id + assert module.runtime_context.get() is None + second = run(rig.service.embed(['second'])) + assert 'granite' in second.model_id + assert calls == ['bekko', 'granite'] + assert bool(first.fallback_reason) == api_failure