one prompt per completion

Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>

Author: sagearc

api/tokenizerpb/tokenizer.pb.go

@@ -185,23 +185,24 @@ message RenderChatCompletionResponse {
 // RenderCompletionRequest contains the fields from an OpenAI completion request
 // that are relevant to rendering (validation + tokenization).
 message RenderCompletionRequest {
-  string model_name = 1;        // Model name to use for renderer selection
-  repeated string prompts = 2;  // Text prompts to render (one item per response)
+  string model_name = 1; // Model name to use for renderer selection
+  string prompt = 2;     // Text prompt to render
 }
 
-// RenderCompletionResponse contains the rendered output for each prompt in the completion request
+// RenderCompletionResponse contains the rendered output for a completion request.
 message RenderCompletionResponse {
-  repeated RenderChatCompletionResponse items = 1; // One item per prompt in the request
-  bool success = 2;                                // Whether the request was successful
-  string error_message = 3;                        // Error message if the request failed
+  string request_id = 1;         // Request ID from the render response
+  repeated uint32 token_ids = 2; // Token IDs for the rendered prompt
+  bool success = 3;              // Whether the request was successful
+  string error_message = 4;      // Error message if the request failed
 }
 
 // TokenizationService defines the gRPC service for tokenization
 service TokenizationService {
   // Tokenize converts a text input to token IDs
   rpc Tokenize(TokenizeRequest) returns (TokenizeResponse);
 
-  // RenderChatTemplate renders a chat template with the given messages
+  // Deprecated: use RenderChatCompletion instead.
   rpc RenderChatTemplate(ChatTemplateRequest) returns (ChatTemplateResponse);
 
   // InitializeTokenizer initializes the tokenizer for a specific model

api/tokenizerpb/tokenizer.proto

@@ -193,7 +193,7 @@ func (u *UdsTokenizer) Render(prompt string) ([]uint32, []types.Offset, error) {
 
 	resp, err := u.client.RenderCompletion(ctx, &tokenizerpb.RenderCompletionRequest{
 		ModelName: u.model,
-		Prompts:   []string{prompt},
+		Prompt:    prompt,
 	})
 	if err != nil {
 		return nil, nil, fmt.Errorf("gRPC RenderCompletion request failed: %w", err)
@@ -203,16 +203,7 @@ func (u *UdsTokenizer) Render(prompt string) ([]uint32, []types.Offset, error) {
 		return nil, nil, fmt.Errorf("render completion failed: %s", resp.ErrorMessage)
 	}
 
-	if len(resp.Items) == 0 {
-		return nil, nil, fmt.Errorf("render completion returned no items")
-	}
-
-	item := resp.Items[0]
-	if !item.Success {
-		return nil, nil, fmt.Errorf("render completion item failed: %s", item.ErrorMessage)
-	}
-
-	return item.TokenIds, nil, nil
+	return resp.TokenIds, nil, nil
 }
 
 // Encode tokenizes the input string and returns the token IDs and offsets.

api/tokenizerpb/tokenizer_grpc.pb.go

@@ -141,22 +141,15 @@ func (m *mockTokenizationServer) RenderCompletion(
 		}, nil
 	}
 
-	items := make([]*tokenizerpb.RenderChatCompletionResponse, 0, len(req.Prompts))
-	for _, prompt := range req.Prompts {
-		tokens := make([]uint32, 0, len(prompt))
-		for _, r := range prompt {
-			tokens = append(tokens, uint32(r))
-		}
-		items = append(items, &tokenizerpb.RenderChatCompletionResponse{
-			RequestId: "mock-request-id",
-			TokenIds:  tokens,
-			Success:   true,
-		})
+	tokens := make([]uint32, 0, len(req.Prompt))
+	for _, r := range req.Prompt {
+		tokens = append(tokens, uint32(r))
 	}
 
 	return &tokenizerpb.RenderCompletionResponse{
-		Items:   items,
-		Success: true,
+		RequestId: "mock-request-id",
+		TokenIds:  tokens,
+		Success:   true,
 	}, nil
 }
 

pkg/tokenization/uds_tokenizer.go

@@ -235,18 +235,20 @@ def RenderCompletion(
         try:
             completion_request = CompletionRequest(
                 model=request.model_name,
-                prompt=list(request.prompts),
+                prompt=request.prompt,
             )
             results = asyncio.run_coroutine_threadsafe(
                 self.renderer_service.render_completion(
                     completion_request, request.model_name
                 ),
                 self._loop,
             ).result()
-            items: list[tokenizer_pb2.RenderChatCompletionResponse] = [
-                self._generate_request_to_proto(r) for r in results
-            ]
-            return tokenizer_pb2.RenderCompletionResponse(items=items, success=True)
+            result = results[0]
+            return tokenizer_pb2.RenderCompletionResponse(
+                request_id=result.request_id,
+                token_ids=list(result.token_ids),
+                success=True,
+            )
         except Exception as e:
             logging.error(f"RenderCompletion failed: {e}", exc_info=True)
             context.abort(grpc.StatusCode.INTERNAL, str(e))