diff --git a/README.md b/README.md index 9f44b5b..05c04d5 100644 --- a/README.md +++ b/README.md @@ -358,6 +358,21 @@ response = client.chat.completions.create( ) ``` +The server also exposes the legacy text completions endpoint `POST /v1/completions` (returning `object: "text_completion"` responses), so OpenAI-legacy clients work against the inbuilt inference server too: + +```bash +curl http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $OPTILLM_API_KEY" \ + -d '{ + "model": "meta-llama/Llama-3.2-1B-Instruct", + "prompt": "Write a haiku about recursion.", + "max_tokens": 64 + }' +``` + +Streaming (`stream: true`) is not supported on `/v1/completions` yet; use `/v1/chat/completions` for streaming responses. + ### Starting the optillm proxy with an external server (e.g. llama.cpp or ollama) - Set the `OPENAI_API_KEY` env variable to a placeholder value diff --git a/optillm/server.py b/optillm/server.py index 8ed597f..707d41a 100644 --- a/optillm/server.py +++ b/optillm/server.py @@ -4,6 +4,7 @@ import secrets import time import traceback +import uuid from pathlib import Path from flask import Flask, request, jsonify from cerebras.cloud.sdk import Cerebras @@ -756,21 +757,64 @@ def check_api_key(): if not secrets.compare_digest(client_key, server_config['optillm_api_key']): return jsonify({"error": "Invalid API key"}), 401 -@app.route('/v1/chat/completions', methods=['POST']) -def proxy(): - logger.info('Received request to /v1/chat/completions') - data = request.get_json() - auth_header = request.headers.get("Authorization") - bearer_token = "" +class CompletionError(Exception): + """Raised by the shared completion pipeline on processing failure. + Endpoints map it to an HTTP 500 JSON response.""" + + +class CompletionOutcome: + """Result of the shared completion pipeline (see _execute_completion). + + response: str, list[str], or a full OpenAI-shaped completion dict passed + through from the 'none' approach / a full-response plugin + (e.g. proxy) / batch mode + completion_tokens: usage count reported by the approach + model: dispatch model name (approach-prefixed) + passthrough: None on the normal path, 'dict' when response is a + passthrough completion dict, 'batch' for batch mode + request_id: conversation logging id (or None) + """ + __slots__ = ('response', 'completion_tokens', 'model', 'passthrough', 'request_id') + + def __init__(self, response, completion_tokens, model, passthrough, request_id): + self.response = response + self.completion_tokens = completion_tokens + self.model = model + self.passthrough = passthrough + self.request_id = request_id + + +def _normalize_response_texts(response): + """Convert tagged conversation output to plain text(s).""" + if isinstance(response, list): + processed_response = tagged_conversation_to_messages(response) + # If processed_response is a list of message lists, extract last message content + if processed_response != response: # Only process if format changed + response = [msg[-1]['content'] if isinstance(msg, list) and msg else msg + for msg in processed_response] + # Otherwise keep original response + else: + messages = tagged_conversation_to_messages(response) + if isinstance(messages, list) and messages: # Only process if format changed + response = messages[-1]['content'] + return response - if auth_header and auth_header.startswith("Bearer "): - bearer_token = auth_header.split("Bearer ")[1].strip() - logger.debug(f"Intercepted Bearer Token: {bearer_token}") +def _execute_completion(data: dict, messages: list, bearer_token: str) -> CompletionOutcome: + """ + Shared completion pipeline behind /v1/chat/completions and /v1/completions. + + Parses request_config, resolves the approach/model pair, builds the client + and dispatches to the selected approach(es). Request parsing (auth header, + prompt/messages extraction) and response shaping stay in the endpoints so + each can keep its own protocol shape. + + Raises CompletionError on processing failure; endpoints map that to an + HTTP 500 JSON response. + """ logger.debug(f'Request data: {data}') stream = data.get('stream', False) - messages = data.get('messages', []) model = data.get('model', server_config['model']) n = data.get('n', server_config['n']) # Get n value from request or config # Extract response_format if present @@ -781,10 +825,6 @@ def proxy(): max_completion_tokens = data.get('max_completion_tokens') max_tokens = data.get('max_tokens') - # stream_options only applies to our own SSE output; upstream calls are - # never streamed and providers reject stream_options without stream=true - include_usage = bool((data.get('stream_options') or {}).get('include_usage')) - # Explicit keys that we are already handling explicit_keys = {'stream', 'stream_options', 'messages', 'model', 'n', 'response_format', 'max_completion_tokens', 'max_tokens'} @@ -883,11 +923,11 @@ def proxy(): logger.debug("Routing request to batch processor") result = request_batcher.add_request(batch_request_data) - return jsonify(result), 200 + return CompletionOutcome(result, 0, model, 'batch', request_id) except BatchingError as e: logger.error(f"Batch processing failed: {e}") - return jsonify({"error": str(e)}), 500 + raise CompletionError(str(e)) # Check if any of the approaches is 'none' contains_none = any(approach == 'none' for approach in approaches) @@ -903,14 +943,7 @@ def proxy(): conversation_logger.log_final_response(request_id, result) conversation_logger.finalize_conversation(request_id) - if stream: - if request_id: - logger.info(f'Request {request_id}: Completed (streaming response)') - return Response(generate_streaming_completion(result, model, include_usage), content_type='text/event-stream') - else : - if request_id: - logger.info(f'Request {request_id}: Completed') - return jsonify(result), 200 + return CompletionOutcome(result, completion_tokens, model, 'dict', request_id) elif operation == 'AND' or operation == 'OR': if contains_none: @@ -926,15 +959,10 @@ def proxy(): conversation_logger.log_final_response(request_id, response) conversation_logger.finalize_conversation(request_id) - if stream: - if request_id: - logger.info(f'Request {request_id}: Completed (streaming response)') - return Response(generate_streaming_completion(response, model, include_usage), content_type='text/event-stream') - else: - if request_id: - logger.info(f'Request {request_id}: Completed') - return jsonify(response), 200 + return CompletionOutcome(response, completion_tokens, model, 'dict', request_id) + except CompletionError: + raise except Exception as e: # Log error to conversation logger if enabled if conversation_logger and request_id: @@ -945,72 +973,208 @@ def proxy(): logger.error(f"Error processing request{request_id_str}: {str(e)}") if logger.getEffectiveLevel() == logging.DEBUG: logger.exception("Debug request exception") + raise CompletionError(str(e)) + + return CompletionOutcome(response, completion_tokens, model, None, request_id) + + +@app.route('/v1/chat/completions', methods=['POST']) +def proxy(): + logger.info('Received request to /v1/chat/completions') + data = request.get_json() + auth_header = request.headers.get("Authorization") + bearer_token = "" + + if auth_header and auth_header.startswith("Bearer "): + bearer_token = auth_header.split("Bearer ")[1].strip() + logger.debug(f"Intercepted Bearer Token: {bearer_token}") + + messages = data.get('messages', []) + stream = data.get('stream', False) + + try: + outcome = _execute_completion(data, messages, bearer_token) + except CompletionError as e: return jsonify({"error": str(e)}), 500 - # Convert tagged conversation to messages format if needed - if isinstance(response, list): - processed_response = tagged_conversation_to_messages(response) - # If processed_response is a list of message lists, extract last message content - if processed_response != response: # Only process if format changed - response = [msg[-1]['content'] if isinstance(msg, list) and msg else msg - for msg in processed_response] - # Otherwise keep original response - else: - messages = tagged_conversation_to_messages(response) - if isinstance(messages, list) and messages: # Only process if format changed - response = messages[-1]['content'] + # Batch mode returns the aggregated response as-is (no streaming wrapper, + # matching the historical behavior of the batch path) + if outcome.passthrough == 'batch': + return jsonify(outcome.response), 200 + + # stream_options only applies to our own SSE output; upstream calls are + # never streamed and providers reject stream_options without stream=true + include_usage = bool((data.get('stream_options') or {}).get('include_usage')) if stream: - return Response(generate_streaming_response(response, model), content_type='text/event-stream') - else: - # Calculate reasoning tokens from the response - reasoning_tokens = 0 - if isinstance(response, str): - reasoning_tokens = count_reasoning_tokens(response) - elif isinstance(response, list) and response: - # For multiple responses, sum up reasoning tokens from all - reasoning_tokens = sum(count_reasoning_tokens(resp) for resp in response if isinstance(resp, str)) - - response_data = { - 'model': model, - 'choices': [], - 'usage': { - 'completion_tokens': completion_tokens, - 'completion_tokens_details': { - 'reasoning_tokens': reasoning_tokens - } + if isinstance(outcome.response, dict): + # Full completion dict from the 'none' approach or a full-response plugin + if outcome.request_id: + logger.info(f'Request {outcome.request_id}: Completed (streaming response)') + return Response(generate_streaming_completion(outcome.response, outcome.model, include_usage), content_type='text/event-stream') + response = _normalize_response_texts(outcome.response) + return Response(generate_streaming_response(response, outcome.model), content_type='text/event-stream') + + if isinstance(outcome.response, dict): + # Full completion dict: return it directly + if outcome.request_id: + logger.info(f'Request {outcome.request_id}: Completed') + return jsonify(outcome.response), 200 + + response = _normalize_response_texts(outcome.response) + + # Calculate reasoning tokens from the response + reasoning_tokens = 0 + if isinstance(response, str): + reasoning_tokens = count_reasoning_tokens(response) + elif isinstance(response, list) and response: + # For multiple responses, sum up reasoning tokens from all + reasoning_tokens = sum(count_reasoning_tokens(resp) for resp in response if isinstance(resp, str)) + + response_data = { + 'model': outcome.model, + 'choices': [], + 'usage': { + 'completion_tokens': outcome.completion_tokens, + 'completion_tokens_details': { + 'reasoning_tokens': reasoning_tokens } } + } - if isinstance(response, list): - for index, resp in enumerate(response): - response_data['choices'].append({ - 'index': index, - 'message': { - 'role': 'assistant', - 'content': resp, - }, - 'finish_reason': 'stop' - }) - else: + if isinstance(response, list): + for index, resp in enumerate(response): response_data['choices'].append({ - 'index': 0, + 'index': index, 'message': { 'role': 'assistant', - 'content': response, + 'content': resp, }, 'finish_reason': 'stop' }) + else: + response_data['choices'].append({ + 'index': 0, + 'message': { + 'role': 'assistant', + 'content': response, + }, + 'finish_reason': 'stop' + }) - # Log the final response and finalize conversation logging - if conversation_logger and request_id: - conversation_logger.log_final_response(request_id, response_data) - conversation_logger.finalize_conversation(request_id) + # Log the final response and finalize conversation logging + if conversation_logger and outcome.request_id: + conversation_logger.log_final_response(outcome.request_id, response_data) + conversation_logger.finalize_conversation(outcome.request_id) + + logger.debug(f'API response: {response_data}') + if outcome.request_id: + logger.info(f'Request {outcome.request_id}: Completed') + return jsonify(response_data), 200 + +@app.route('/v1/completions', methods=['POST']) +def completions(): + """ + Legacy text completions endpoint, complementing /v1/chat/completions so + the built-in inference server also serves OpenAI-legacy clients + (issue #170). Reuses the shared completion pipeline; streaming is not + supported in this first version. + """ + logger.info('Received request to /v1/completions') + data = request.get_json() + if data is None: + return jsonify({"error": "Request body must be JSON"}), 400 + + auth_header = request.headers.get("Authorization") + bearer_token = "" + + if auth_header and auth_header.startswith("Bearer "): + bearer_token = auth_header.split("Bearer ")[1].strip() + logger.debug(f"Intercepted Bearer Token: {bearer_token}") - logger.debug(f'API response: {response_data}') - if request_id: - logger.info(f'Request {request_id}: Completed') - return jsonify(response_data), 200 + if data.get('stream', False): + return jsonify({"error": "stream=true is not supported on /v1/completions yet; use /v1/chat/completions for streaming"}), 400 + + prompt = data.get('prompt') + if prompt is None: + return jsonify({"error": "'prompt' is required for /v1/completions"}), 400 + + prompts = prompt if isinstance(prompt, list) else [prompt] + if not prompts or not all(isinstance(p, str) for p in prompts): + return jsonify({"error": "'prompt' must be a string or a list of strings"}), 400 + + # Map each prompt to a single user message and run the shared pipeline + # once per prompt (no HTTP self-call). 'prompt' itself is not a chat API + # parameter, so it is dropped before dispatch. + outcomes = [] + try: + for single_prompt in prompts: + request_data = {k: v for k, v in data.items() if k != 'prompt'} + request_data['messages'] = [{"role": "user", "content": single_prompt}] + outcomes.append(_execute_completion(request_data, request_data['messages'], bearer_token)) + except CompletionError as e: + return jsonify({"error": str(e)}), 500 + + choices = [] + completion_tokens_total = 0 + reasoning_tokens_total = 0 + passthrough_usage = None + model = outcomes[-1].model + + for outcome in outcomes: + if outcome.passthrough is not None: + # Full completion dict (batch mode / 'none' approach / full-response + # plugin): map chat choices to legacy text choices + for choice in outcome.response.get('choices', []): + message = choice.get('message') or {} + choices.append({ + 'index': len(choices), + 'text': message.get('content', ''), + 'logprobs': choice.get('logprobs'), + 'finish_reason': choice.get('finish_reason', 'stop'), + }) + if outcome.response.get('usage'): + passthrough_usage = outcome.response['usage'] + continue + + response = _normalize_response_texts(outcome.response) + texts = response if isinstance(response, list) else [response] + completion_tokens_total += outcome.completion_tokens + reasoning_tokens_total += sum(count_reasoning_tokens(text) for text in texts if isinstance(text, str)) + for text in texts: + choices.append({ + 'index': len(choices), + 'text': text, + 'logprobs': None, + 'finish_reason': 'stop', + }) + + response_data = { + 'id': f"cmpl-{uuid.uuid4().hex}", + 'object': 'text_completion', + 'created': int(time.time()), + 'model': model, + 'choices': choices, + 'usage': passthrough_usage if passthrough_usage is not None else { + 'completion_tokens': completion_tokens_total, + 'completion_tokens_details': { + 'reasoning_tokens': reasoning_tokens_total + } + } + } + + # Log the final response and finalize conversation logging for the normal + # path (passthrough paths are logged and finalized inside the pipeline) + for outcome in outcomes: + if conversation_logger and outcome.request_id and outcome.passthrough is None: + conversation_logger.log_final_response(outcome.request_id, response_data) + conversation_logger.finalize_conversation(outcome.request_id) + + logger.debug(f'API response: {response_data}') + for outcome in outcomes: + if outcome.request_id: + logger.info(f'Request {outcome.request_id}: Completed') + return jsonify(response_data), 200 @app.route('/v1/models', methods=['GET']) def proxy_models(): diff --git a/tests/test_completions_endpoint.py b/tests/test_completions_endpoint.py new file mode 100644 index 0000000..d072d18 --- /dev/null +++ b/tests/test_completions_endpoint.py @@ -0,0 +1,205 @@ +""" +Offline unit tests for the legacy /v1/completions endpoint (issue #170). + +Uses the Flask test client (no real server port, no network, no model) and +mocks the approach function (chat_with_mcts) following the pattern in +tests/test_approaches.py. Also asserts that /v1/chat/completions keeps its +exact historical response shape. +""" + +import sys +import os +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import unittest +from unittest.mock import patch + +from optillm import server + + +MOCK_TEXT = "legacy completion text" +CAPTURED = {} + + +def fake_chat_with_mcts(system_prompt, initial_query, client, model, + num_simulations, exploration_weight, depth, + request_config=None, request_id=None): + CAPTURED['model'] = model + CAPTURED['request_config'] = request_config + return MOCK_TEXT, 5 + + +def fake_chat_with_mcts_boom(*args, **kwargs): + raise RuntimeError("approach exploded") + + +class _FakeResp: + def model_dump(self): + return {"choices": [{"index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} + + +class _FakeClient: + def __init__(self): + self.chat = type("_Chat", (), { + "completions": type("_Completions", (), {"create": lambda self, **k: _FakeResp()})() + })() + + +class CompletionsEndpointTestBase(unittest.TestCase): + """Shared fixture: mock client + approach, Flask test client.""" + + def setUp(self): + CAPTURED.clear() + self._originals = { + 'get_config': server.get_config, + 'chat_with_mcts': server.chat_with_mcts, + 'config': dict(server.server_config), + } + server.get_config = lambda: (_FakeClient(), 'optillm') + server.chat_with_mcts = fake_chat_with_mcts + server.server_config['optillm_api_key'] = '' + server.app.config['TESTING'] = True + self.client = server.app.test_client() + + def tearDown(self): + server.get_config = self._originals['get_config'] + server.chat_with_mcts = self._originals['chat_with_mcts'] + server.server_config.clear() + server.server_config.update(self._originals['config']) + + +class TestLegacyCompletionsEndpoint(CompletionsEndpointTestBase): + + def test_basic_text_completion_shape(self): + resp = self.client.post( + '/v1/completions', + json={"model": "gpt-4o-mini", + "prompt": "Say hello", + "max_tokens": 32, + "optillm_approach": "mcts"}, + headers={"Authorization": "Bearer optillm"}) + + self.assertEqual(resp.status_code, 200, resp.get_data(as_text=True)) + body = resp.get_json() + + # OpenAI legacy completion shape + self.assertEqual(body["object"], "text_completion") + self.assertTrue(body["id"].startswith("cmpl-"), body["id"]) + self.assertIsInstance(body["created"], int) + self.assertEqual(body["model"], "gpt-4o-mini") + self.assertEqual(len(body["choices"]), 1) + choice = body["choices"][0] + self.assertEqual(choice["text"], MOCK_TEXT) + self.assertEqual(choice["index"], 0) + self.assertEqual(choice["finish_reason"], "stop") + self.assertIsNone(choice["logprobs"]) + + # usage reports the approach's completion token count + self.assertEqual(body["usage"]["completion_tokens"], 5) + self.assertIn("reasoning_tokens", body["usage"]["completion_tokens_details"]) + + # optillm_approach routes to the mocked mcts approach; the slug prefix + # is consumed by parse_combined_approach, so the approach itself sees + # the bare model name (same dispatch behavior as /v1/chat/completions) + self.assertEqual(CAPTURED.get("model"), "gpt-4o-mini") + + # request_config carries the token limit through to the approach + self.assertEqual(CAPTURED["request_config"].get("max_tokens"), 32) + + def test_multiple_prompts_expand_to_multiple_choices(self): + resp = self.client.post( + '/v1/completions', + json={"model": "gpt-4o-mini", + "prompt": ["first prompt", "second prompt"], + "optillm_approach": "mcts"}) + + self.assertEqual(resp.status_code, 200) + body = resp.get_json() + self.assertEqual(len(body["choices"]), 2) + self.assertEqual([c["text"] for c in body["choices"]], [MOCK_TEXT, MOCK_TEXT]) + self.assertEqual([c["index"] for c in body["choices"]], [0, 1]) + + def test_n_expansion_multiple_choices(self): + resp = self.client.post( + '/v1/completions', + json={"model": "gpt-4o-mini", + "prompt": "one prompt", + "n": 2, + "optillm_approach": "mcts"}) + + self.assertEqual(resp.status_code, 200) + body = resp.get_json() + self.assertEqual(len(body["choices"]), 2) + # completion tokens accumulate across the n attempts (2 x 5) + self.assertEqual(body["usage"]["completion_tokens"], 10) + + def test_stream_not_supported(self): + resp = self.client.post( + '/v1/completions', + json={"model": "gpt-4o-mini", + "prompt": "hi", + "stream": True, + "optillm_approach": "mcts"}) + + self.assertEqual(resp.status_code, 400) + self.assertIn("error", resp.get_json()) + + def test_prompt_required(self): + resp = self.client.post( + '/v1/completions', + json={"model": "gpt-4o-mini", "optillm_approach": "mcts"}) + self.assertEqual(resp.status_code, 400) + self.assertIn("prompt", resp.get_json()["error"]) + + def test_prompt_type_validation(self): + resp = self.client.post( + '/v1/completions', + json={"model": "gpt-4o-mini", "prompt": [], "optillm_approach": "mcts"}) + self.assertEqual(resp.status_code, 400) + + resp = self.client.post( + '/v1/completions', + json={"model": "gpt-4o-mini", "prompt": [1, 2], "optillm_approach": "mcts"}) + self.assertEqual(resp.status_code, 400) + + def test_approach_error_maps_to_500(self): + server.chat_with_mcts = fake_chat_with_mcts_boom + resp = self.client.post( + '/v1/completions', + json={"model": "gpt-4o-mini", + "prompt": "hi", + "optillm_approach": "mcts"}) + self.assertEqual(resp.status_code, 500) + self.assertIn("approach exploded", resp.get_json()["error"]) + + +class TestChatEndpointUnchanged(CompletionsEndpointTestBase): + """The /v1/chat/completions shape must not change (regression guard).""" + + def test_chat_shape_unchanged(self): + resp = self.client.post( + '/v1/chat/completions', + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "optillm_approach": "mcts"}) + + self.assertEqual(resp.status_code, 200, resp.get_data(as_text=True)) + body = resp.get_json() + self.assertNotIn("object", body) + self.assertNotIn("id", body) + self.assertEqual(body["model"], "gpt-4o-mini") + self.assertEqual(len(body["choices"]), 1) + choice = body["choices"][0] + self.assertNotIn("text", choice) + self.assertEqual(choice["message"]["role"], "assistant") + self.assertEqual(choice["message"]["content"], MOCK_TEXT) + self.assertEqual(choice["finish_reason"], "stop") + self.assertEqual(body["usage"]["completion_tokens"], 5) + self.assertEqual(CAPTURED.get("model"), "gpt-4o-mini") + + +if __name__ == '__main__': + unittest.main()