"""How full the context is, what a reply cost, and how fast it arrived.""" from __future__ import annotations from fastapi.testclient import TestClient from sqlalchemy import select from lembas.db.models import Connection, Model from lembas.services import tokens from lembas.services.crypto import encrypt from lembas.services.llm.openai_client import chunk_usage, context_from def _model(db, **kwargs) -> Model: connection = Connection( name="Test", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("") ) db.add(connection) db.commit() model = Model(connection_id=connection.id, model_id="test-model", **kwargs) db.add(model) db.commit() return model # --- Reading a context length off /v1/models --------------------------------- def test_context_length_is_read_from_any_of_the_spellings(): assert context_from({"id": "m", "context_length": 8192}) == 8192 assert context_from({"id": "m", "max_model_len": 32768}) == 32768 assert context_from({"id": "m", "context_window": 4096}) == 4096 assert context_from({"id": "m", "meta": {"n_ctx": 2048}}) == 2048 def test_a_quoted_number_is_accepted_but_a_label_is_not(): """Some servers quote it. "8192 tokens" is a label, not a measurement.""" assert context_from({"id": "m", "context_length": "8192"}) == 8192 assert context_from({"id": "m", "context_length": "8192 tokens"}) == 0 def test_an_absent_or_implausible_context_length_is_zero(): assert context_from({"id": "m"}) == 0 assert context_from({"id": "m", "context_length": 0}) == 0 assert context_from({"id": "m", "context_length": 64}) == 0 assert context_from({"id": "m", "context_length": 10**9}) == 0 # True is an int in Python, and it is not a context length. assert context_from({"id": "m", "context_length": True}) == 0 # --- Discovery --------------------------------------------------------------- async def test_discovery_fills_in_a_context_length(client: TestClient, db, registered, mock_http): import httpx mock_http( lambda _r: httpx.Response( 200, json={"data": [{"id": "big-model", "context_length": 16384}]} ) ) client.post( "/admin/connections", data={"name": "Local", "base_url": "http://x.test", "api_key": ""}, follow_redirects=False, ) model = db.scalar(select(Model).where(Model.model_id == "big-model")) assert model.context_length == 16384 async def test_discovery_never_overwrites_a_number_an_admin_typed( client: TestClient, db, registered, mock_http ): """A refresh must not undo a correction. Administrators set this precisely because the endpoint was wrong or silent.""" import httpx mock_http( lambda _r: httpx.Response(200, json={"data": [{"id": "m", "context_length": 4096}]}) ) client.post( "/admin/connections", data={"name": "Local", "base_url": "http://x.test", "api_key": ""}, follow_redirects=False, ) model = db.scalar(select(Model).where(Model.model_id == "m")) model.context_length = 131072 db.commit() connection = db.scalar(select(Connection)) client.post(f"/admin/connections/{connection.id}/refresh", follow_redirects=False) db.refresh(model) assert model.context_length == 131072 # --- The admin field --------------------------------------------------------- def test_an_admin_can_set_and_clear_the_context_length(client: TestClient, db, registered): model = _model(db) client.post( f"/admin/models/{model.id}", data={"context_length": "8192", "position": ""}, follow_redirects=False, ) db.refresh(model) assert model.context_length == 8192 client.post( f"/admin/models/{model.id}", data={"context_length": "", "position": ""}, follow_redirects=False, ) db.refresh(model) assert model.context_length == 0 def test_junk_in_the_context_length_field_is_ignored_not_a_500( client: TestClient, db, registered ): model = _model(db, context_length=4096) response = client.post( f"/admin/models/{model.id}", data={"context_length": "eight thousand", "position": ""}, follow_redirects=False, ) assert response.status_code == 303 db.refresh(model) assert model.context_length == 4096 # --- Usage off the wire ------------------------------------------------------- def test_usage_is_read_from_a_usage_chunk(): chunk = { "choices": [], "usage": {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120}, } assert chunk_usage(chunk) == { "prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120, } def test_a_missing_total_is_worked_out(): chunk = {"choices": [], "usage": {"prompt_tokens": 100, "completion_tokens": 20}} assert chunk_usage(chunk)["total_tokens"] == 120 def test_an_ordinary_chunk_carries_no_usage(): assert chunk_usage({"choices": [{"delta": {"content": "hi"}}]}) is None assert chunk_usage({}) is None assert chunk_usage({"usage": "lots"}) is None def test_an_all_zero_usage_object_is_not_an_answer(): """Some servers attach zeros to every chunk and the real numbers only at the end. Believing the zeros freezes the count at nothing.""" chunk = {"choices": [], "usage": {"prompt_tokens": 0, "completion_tokens": 0}} assert chunk_usage(chunk) is None def test_the_other_accessors_still_ignore_a_usage_chunk(): """They return early on `choices: []`, which is exactly the shape of one. That is what lets a usage chunk through the loop untouched.""" from lembas.services.llm.openai_client import ( delta_reasoning, delta_text, delta_tool_calls, finish_reason, ) chunk = {"choices": [], "usage": {"prompt_tokens": 1, "completion_tokens": 1}} assert delta_text(chunk) == "" assert delta_reasoning(chunk) == "" assert delta_tool_calls(chunk) == [] assert finish_reason(chunk) == "" async def test_stream_options_is_asked_for(mock_http): import json as json_module import httpx from lembas.services.llm.openai_client import Endpoint, stream_chat seen: list[dict] = [] def handler(request: httpx.Request) -> httpx.Response: seen.append(json_module.loads(request.content)) return httpx.Response(200, text="data: [DONE]\n\n") mock_http(handler) async for _ in stream_chat(Endpoint("http://ask.test", "", {}), {"model": "m"}): pass assert seen[0]["stream_options"] == {"include_usage": True} async def test_an_endpoint_that_rejects_stream_options_is_asked_once(mock_http): """A 400 for an unknown key is the same hazard as sending `tools` to an endpoint without support. Retry without it, then stop asking.""" import json as json_module import httpx from lembas.services.llm.openai_client import ( _NO_STREAM_OPTIONS, Endpoint, stream_chat, ) _NO_STREAM_OPTIONS.discard("http://fussy.test") seen: list[dict] = [] def handler(request: httpx.Request) -> httpx.Response: body = json_module.loads(request.content) seen.append(body) if "stream_options" in body: return httpx.Response(400, json={"error": {"message": "unknown field"}}) reply = 'data: {"choices":[{"delta":{"content":"hi"}}]}\n\ndata: [DONE]\n\n' return httpx.Response(200, text=reply) mock_http(handler) endpoint = Endpoint("http://fussy.test", "", {}) text = [c async for c in stream_chat(endpoint, {"model": "m"})] assert text, "the retry should have produced the reply" assert len(seen) == 2 # Second reply: it already knows, so one request and no stream_options. async for _ in stream_chat(endpoint, {"model": "m"}): pass assert len(seen) == 3 assert "stream_options" not in seen[2] _NO_STREAM_OPTIONS.discard("http://fussy.test") # --- The estimate ------------------------------------------------------------- def test_the_estimate_is_about_four_characters_a_token(): assert tokens.estimate("") == 0 assert tokens.estimate("x" * 400) == 100 def test_typed_content_parts_are_counted_and_images_are_not(): """An image's cost depends on the model's tiling. A number invented here would be worse than the omission.""" content = [ {"type": "text", "text": "x" * 40}, {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA" * 500}}, ] assert tokens.estimate_content(content) == 10 def test_a_request_estimate_includes_the_tools_array(): """Thirteen schemas is a meaningful slice of a short window; leaving them out would read low exactly when it matters.""" payload = { "messages": [{"role": "user", "content": "x" * 40}], "tools": [ {"function": {"name": "web_search", "description": "y" * 400, "parameters": {}}} ], } assert tokens.estimate_request(payload) > tokens.estimate_request( {"messages": payload["messages"]} )