Use extend method for embeddings concatenation in mlflow_gateway

[jamesev15]

Description

There is a bug in the concatenation of embeddings obtained from MLflow that does not conform to the type hint requested by the function.

def _query(self, texts: List[str]) -> List[List[float]]:

It is logical to expect a List[List[float]] for a List[str]. However, the append method encapsulates the response in a global List.
To avoid this, the extend method should be used, which will add the embeddings of all strings at the same list level.

Testing

I have tried using OpenAI-ADA to obtain the embeddings, and the result of executing this snippet is as follows:

embeds = await MlflowAIGatewayEmbeddings().aembed_documents(texts=["hi", "how are you?"])
print(embeds)

[[[-0.03512698, -0.020624293, -0.015343423, ...], [-0.021260535, -0.011461929, -0.00033121882, ...]]]

When in reality, the expected result should be:

[[-0.03512698, -0.020624293, -0.015343423, ...], [-0.021260535, -0.011461929, -0.00033121882, ...]]

The above result complies with the expected type hint: List[List[float]] . As I mentioned, we can achieve that by using the extend method instead of the append method.

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
langchain	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Aug 23, 2024 2:43pm

[jamesev15]

@hwchase17 I just want to use the extend method instead of the append method. Could you take a quick look at it? :)

[baskaryan]

cc @harupy