Spaces:
Sleeping
Sleeping
| """Generator backends against local HTTP servers that implement the Ollama and OpenAI wire protocols.""" | |
| import base64 | |
| import io | |
| import json | |
| import threading | |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer | |
| import pytest | |
| from PIL import Image | |
| from mrag.config import Config | |
| from mrag.generate import Generator, ImageHit, TextHit | |
| class FakeLLM: | |
| def __init__(self): | |
| self.requests = [] | |
| outer = self | |
| class H(BaseHTTPRequestHandler): | |
| def log_message(self, *a): | |
| pass | |
| def do_POST(self): | |
| body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) | |
| outer.requests.append((self.path, dict(self.headers), body)) | |
| words = ["The ", "warranty ", "is ", "36 ", "months ", "[1]."] | |
| if self.path == "/api/chat": | |
| if body["model"] == "missing": | |
| return self._send(404, {"error": "model 'missing' not found"}) | |
| if not body["stream"]: | |
| return self._send(200, {"message": {"role": "assistant", "content": "".join(words)}, | |
| "done": True}) | |
| self.send_response(200) | |
| self.send_header("Content-Type", "application/x-ndjson") | |
| self.end_headers() | |
| for w in words: | |
| self.wfile.write((json.dumps({"message": {"content": w}, "done": False}) + "\n").encode()) | |
| self.wfile.write((json.dumps({"message": {"content": ""}, "done": True}) + "\n").encode()) | |
| elif self.path == "/v1/chat/completions": | |
| if not body["stream"]: | |
| return self._send(200, {"choices": [{"message": {"content": "".join(words)}}]}) | |
| self.send_response(200) | |
| self.send_header("Content-Type", "text/event-stream") | |
| self.end_headers() | |
| for w in words: | |
| self.wfile.write(f"data: {json.dumps({'choices': [{'delta': {'content': w}}]})}\n\n".encode()) | |
| self.wfile.write(b"data: [DONE]\n\n") | |
| else: | |
| self._send(404, {"error": "no route"}) | |
| def _send(self, code, obj): | |
| b = json.dumps(obj).encode() | |
| self.send_response(code) | |
| self.send_header("Content-Type", "application/json") | |
| self.send_header("Content-Length", str(len(b))) | |
| self.end_headers() | |
| self.wfile.write(b) | |
| self.server = ThreadingHTTPServer(("127.0.0.1", 0), H) | |
| self.url = f"http://127.0.0.1:{self.server.server_address[1]}" | |
| threading.Thread(target=self.server.serve_forever, daemon=True).start() | |
| def fake(): | |
| f = FakeLLM() | |
| yield f | |
| f.server.shutdown() | |
| def hits(tmp_path): | |
| p = tmp_path / "chart.png" | |
| Image.new("RGB", (2000, 1000), "red").save(p) | |
| texts = [TextHit("t1", "Every R2 ships with a 36-month warranty.", "/d/report.pdf", 3, "ocr", 0.8)] | |
| images = [ImageHit("i1", str(p), "/d/report.pdf", 3, "warranty", 0.3)] | |
| return texts, images | |
| def test_ollama(fake, hits, stream): | |
| g = Generator(Config(llm_backend="ollama", llm_url=fake.url, llm_model="qwen2.5vl:3b", max_image_side=512)) | |
| out = g.answer("How long is the warranty?", *hits, stream=stream) | |
| text = "".join(out) if stream else out | |
| assert text == "The warranty is 36 months [1]." | |
| path, _, body = fake.requests[-1] | |
| assert path == "/api/chat" and body["model"] == "qwen2.5vl:3b" and body["stream"] is stream | |
| sys_msg, user = body["messages"] | |
| assert sys_msg["role"] == "system" and "[1] (report.pdf, p.3, OCR of an image)" in user["content"] | |
| img = Image.open(io.BytesIO(base64.b64decode(user["images"][0]))) | |
| assert img.size == (512, 256) # downscaled, aspect preserved | |
| assert body["options"]["num_ctx"] == 8192 | |
| def test_ollama_text_only_model(fake, hits): | |
| g = Generator(Config(llm_backend="ollama", llm_url=fake.url, llm_model="llama3.2:3b", llm_vision=False)) | |
| g.answer("q", *hits) | |
| assert "images" not in fake.requests[-1][2]["messages"][1] | |
| def test_ollama_errors(fake, hits): | |
| g = Generator(Config(llm_backend="ollama", llm_url=fake.url, llm_model="missing")) | |
| with pytest.raises(RuntimeError, match="404"): | |
| g.answer("q", *hits) | |
| g = Generator(Config(llm_backend="ollama", llm_url="http://127.0.0.1:9", llm_model="x")) | |
| with pytest.raises(RuntimeError, match="ollama serve"): | |
| g.answer("q", *hits) | |
| def test_openai_compatible(fake, hits, stream, suffix): | |
| g = Generator(Config(llm_backend="openai", llm_url=fake.url + suffix, llm_model="local", llm_api_key="k")) | |
| out = g.answer("How long is the warranty?", *hits, stream=stream) | |
| assert ("".join(out) if stream else out) == "The warranty is 36 months [1]." | |
| path, headers, body = fake.requests[-1] | |
| assert path == "/v1/chat/completions" and headers["Authorization"] == "Bearer k" | |
| parts = body["messages"][1]["content"] | |
| assert parts[0]["type"] == "text" and parts[1]["image_url"]["url"].startswith("data:image/png;base64,") | |