Dumitru Picalau
657dc07831
Phases 2 to 8 of the AI assistant overhaul plus fixes from real-model chats: one tool registry with native tool calls, server-kept chat logs with prompt caching, an FTS5 library index with opt-in meaning search, keys bound to their address and a guard on every AI route, page tools and opt-in web search, saved chats with streaming, Try again, a context meter and compaction, chat images, delete/category/tag permissions with Undo, constraint and test cards, card fixes and previews, questions with answers to click, the derivation NOT toggle (schema v10, community catalog 0021), request fixes for the newest Claude models, and OpenAI's Responses API. Review fixes: Try again clears the old turn's fixes and answers, category Undo removes only what it made, a failed delete rolls back in order, not_matches patterns are checked, busy chats refuse a second request, and chat search ignores stored JSON.
84 lines
3.3 KiB
Python
84 lines
3.3 KiB
Python
"""Embedding requests per wire format, and Test connection's swap check."""
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
from aiohttp import web
|
|
|
|
from wp_api import _ai_client, _ai_embed
|
|
|
|
|
|
@pytest.fixture
|
|
async def embed_server(aiohttp_server):
|
|
seen: list[dict] = []
|
|
state = {"loaded": ["qwen:9b", "nomic-embed-text:latest"]}
|
|
|
|
async def ollama_embed(request):
|
|
body = await request.json()
|
|
seen.append({"path": request.path, "body": body})
|
|
return web.json_response({"embeddings": [[0.1, 0.2] for _ in body["input"]]})
|
|
|
|
async def ollama_ps(request):
|
|
return web.json_response({"models": [{"name": n} for n in state["loaded"]]})
|
|
|
|
async def openai_embed(request):
|
|
body = await request.json()
|
|
seen.append({"path": request.path, "body": body,
|
|
"auth": request.headers.get("authorization")})
|
|
if body["model"] == "chatty":
|
|
return web.json_response({"choices": []})
|
|
data = [{"index": i, "embedding": [float(i), 1.0, 0.0]} for i in range(len(body["input"]))]
|
|
return web.json_response({"data": list(reversed(data))})
|
|
|
|
app = web.Application()
|
|
app.router.add_post("/api/embed", ollama_embed)
|
|
app.router.add_get("/api/ps", ollama_ps)
|
|
app.router.add_post("/v1/embeddings", openai_embed)
|
|
server = await aiohttp_server(app)
|
|
server.seen = seen
|
|
server.state = state
|
|
return server
|
|
|
|
|
|
def _cfg(server, api, **extra):
|
|
base = str(server.make_url("/v1" if api == "openai" else ""))
|
|
return {"api": api, "base_url": base.rstrip("/"), "model": "qwen:9b", "api_key": "",
|
|
"embed_model": "nomic-embed-text", "embed_base_url": "", "unload_after_s": 60,
|
|
**extra}
|
|
|
|
|
|
async def test_ollama_batch_embed(embed_server):
|
|
vecs = await _ai_embed.embed(_cfg(embed_server, "ollama"), ["a", "b"])
|
|
assert vecs == [[0.1, 0.2], [0.1, 0.2]]
|
|
sent = embed_server.seen[-1]
|
|
assert sent["path"] == "/api/embed"
|
|
assert sent["body"] == {"model": "nomic-embed-text", "input": ["a", "b"], "keep_alive": "60s"}
|
|
|
|
|
|
async def test_openai_embed_keeps_order_and_sends_key(embed_server):
|
|
cfg = _cfg(embed_server, "openai", embed_api_key="sk-x")
|
|
vecs = await _ai_embed.embed(cfg, ["a", "b", "c"])
|
|
assert [v[0] for v in vecs] == [0.0, 1.0, 2.0]
|
|
assert embed_server.seen[-1]["auth"] == "Bearer sk-x"
|
|
|
|
|
|
async def test_separate_embedding_address(embed_server):
|
|
cfg = _cfg(embed_server, "openai", base_url="http://127.0.0.1:9/v1",
|
|
embed_base_url=str(embed_server.make_url("/v1")).rstrip("/"))
|
|
assert len(await _ai_embed.embed(cfg, ["a"])) == 1
|
|
|
|
|
|
async def test_not_an_embedding_model(embed_server):
|
|
with pytest.raises(_ai_client.AIError, match="embedding model"):
|
|
await _ai_embed.embed(_cfg(embed_server, "openai", embed_model="chatty"), ["a"])
|
|
assert await _ai_embed.query_vector(_cfg(embed_server, "openai", embed_model="chatty"),
|
|
"x") is None
|
|
|
|
|
|
async def test_check_spots_ollama_unloading_the_chat_model(embed_server):
|
|
cfg = _cfg(embed_server, "ollama")
|
|
assert await _ai_embed.check(cfg) == {"ok": True, "dims": 2, "error": None,
|
|
"swapped": False}
|
|
embed_server.state["loaded"] = ["nomic-embed-text:latest"]
|
|
assert (await _ai_embed.check(cfg))["swapped"] is True
|
|
assert await _ai_embed.check({**cfg, "embed_model": ""}) is None
|