From cf86d8feeb6750f5dfbb08e80a2859bdc13f216a Mon Sep 17 00:00:00 2001 From: Roland Minrui <114476598+RolandMinrui@users.noreply.github.com> Date: Wed, 12 Mar 2025 16:10:33 +0800 Subject: [PATCH] fix: return 1D embedding if create_embedding receive a string input (#670) * return single dimension embedding for string input of create_embedding * Update base.py * fix * fix * fix * fix --------- Co-authored-by: Xu --- rdagent/oai/backend/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/rdagent/oai/backend/base.py b/rdagent/oai/backend/base.py index c9a5789a..c15691f4 100644 --- a/rdagent/oai/backend/base.py +++ b/rdagent/oai/backend/base.py @@ -273,7 +273,7 @@ class APIBackend(ABC): logger.log_object({"system": system_prompt, "user": user_prompt, "resp": resp}, tag="debug_llm") return resp - def create_embedding(self, input_content: str | list[str], *args, **kwargs) -> list[list[float]]: # type: ignore[no-untyped-def] + def create_embedding(self, input_content: str | list[str], *args, **kwargs) -> list[float] | list[list[float]]: # type: ignore[no-untyped-def] input_content_list = [input_content] if isinstance(input_content, str) else input_content resp = self._try_create_chat_completion_or_embedding( # type: ignore[misc] input_content_list=input_content_list, @@ -281,6 +281,8 @@ class APIBackend(ABC): *args, **kwargs, ) + if isinstance(input_content, str): + return resp[0] # type: ignore[return-value] return resp # type: ignore[return-value] def build_messages_and_calculate_token(