Repository: ray-project/ray
File: python/ray/llm/examples/sglang/modules/sglang_engine.py (Python)

Code before:
            all_choices.append(choice_data)

        usage_data = {
            "prompt_tokens": total_prompt_tokens,
            "completion_tokens": total_completion_tokens,
            "total_tokens": total_prompt_tokens + total_completion_tokens,
        }

        # Use metadata from the last generation for shared fields (id, created)
        resp = CompletionResponse(
            id=last_metadata.get("id", f"sglang-batch-gen-{int(time.time())}"),
            object="text_completion",
            created=last_metadata.get("created", int(time.time())),
            model=getattr(request, "model", "default_model"),
            choices=all_choices,
            usage=usage_data,
        )

        yield resp

    async def embeddings(
        self, request: EmbeddingRequest, raw_request: Optional[Any] = None
    ) -> AsyncGenerator[EmbeddingResponse, None]:

        # EmbeddingRequest is Union[EmbeddingCompletionRequest, EmbeddingChatRequest]
        if isinstance(request, EmbeddingCompletionRequest):
            input_data = request.input
            if isinstance(input_data, str):
                texts = [input_data]
            elif isinstance(input_data, list):
                if not input_data:
                    texts = []
                elif isinstance(input_data[0], str):
                    texts = list(input_data)
                else:
                    tokenizer = self.engine.get_tokenizer()
                    if isinstance(input_data[0], list):
                        # Batched token IDs - decode each list
                        texts = [tokenizer.decode(ids) for ids in input_data]
                    else:
                        # Single list of token IDs
                        texts = [tokenizer.decode(input_data)]
            else:
                texts = [str(input_data)]
        else:
            # Chat embedding request - convert messages to prompt
            texts = [format_messages_to_prompt(request.messages)]

        data = []
        total_prompt_tokens = 0


Review comment (suggestion):
Please re-use some logic SGLang `OpenAIServingEmbedding` instead:
https://github.com/sgl-project/sglang/blob/b2573fe4267ae3a6e3cda521fad15e8efb07dc0c/python/sglang/srt/entrypoints/openai/serving_embedding.py#L25

Code after:
            all_choices.append(choice_data)

        usage_data = {
            "prompt_tokens": total_prompt_tokens,
            "completion_tokens": total_completion_tokens,
            "total_tokens": total_prompt_tokens + total_completion_tokens,
        }

        # Use metadata from the last generation for shared fields (id, created)
        resp = CompletionResponse(
            id=last_metadata.get("id", f"sglang-batch-gen-{int(time.time())}"),
            object="text_completion",
            created=last_metadata.get("created", int(time.time())),
            model=getattr(request, "model", "default_model"),
            choices=all_choices,
            usage=usage_data,
        )

        yield resp

    async def embeddings(
        self, request: EmbeddingRequest, raw_request: Optional[Any] = None
    ) -> AsyncGenerator[EmbeddingResponse, None]:
        # Input handling follows SGLang's OpenAIServingEmbedding pattern:
        # https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/openai/serving_embedding.py
        if isinstance(request, EmbeddingCompletionRequest):
            prompt = request.input
        else:
            # Chat embedding request - convert messages to prompt
            prompt = format_messages_to_prompt(request.messages)

        # async_encode handles both single strings and lists of strings
        results = await self.engine.async_encode(prompt)
        if not isinstance(results, list):
            results = [results]

        if not results:
            raise RuntimeError(
                "SGLang engine returned an empty response for embedding request."
            )

        # Build response following SGLang's _build_embedding_response pattern
        data = []
        total_prompt_tokens = 0

        for idx, ret_item in enumerate(results):
            data.append(
                {
                    "index": idx,
                    "object": "embedding",
                    "embedding": ret_item.get("embedding", []),