Support embeddings models

[dmontagu]

Started this in collaboration with @DouweM, I'd like to ensure consensus on the API design before adding the remaining-providers/logfire-instrumentation/docs/tests.

This is inspired by the approach in haiku.rag, though we adapted it to be a bit closer to the Agent APIs are used (and how you can override model, settings, etc.).

Closes #58

Example:

import asyncio

from pydantic_ai.embeddings import Embedder

embedder = Embedder("openai:text-embedding-3-large")


async def main():
    result = await embedder.embed_documents(["hello", "world"])
    print(result)
    # (IsList, snapshot, and IsDatetime are testing helpers, but you get the point)
    # EmbeddingResult(
    #     embeddings=IsList(
    #         IsList(
    #             snapshot(0.01681816205382347),
    #             snapshot(-0.05579638481140137),
    #             snapshot(0.005661087576299906),
    #             length=1536,
    #         ),
    #         IsList(
    #             snapshot(-0.010592407546937466),
    #             snapshot(-0.03599696233868599),
    #             snapshot(0.030227113515138626),
    #             length=1536,
    #         ),
    #         length=2,
    #     ),
    #     inputs=['hello', 'world'],
    #     input_type='document',
    #     usage=RequestUsage(input_tokens=2),
    #     model_name='text-embedding-3-small',
    #     timestamp=IsDatetime(),
    #     provider_name='openai',
    # )



if __name__ == "__main__":
    asyncio.run(main())

To do:

• Embedder.embed_sync
• Distinguish between encoding query and document(s) Extension of embeddings draft implementation to support local models #3463
• Help with token counting (Support embeddings models #3252 (comment))
  • Add count_tokens
  • Add max_input_tokens
• Expose token usage and cost
• Test known model names completeness
• Instrumentation
  • Cost tracking
  • Verify Claude-generated attrs and metrics
  • Update logfire.instrument_pydantic_ai(): Instrument Pydantic AI embedders and embedding models logfire#1575
• Error handling (ModelAPIError)
• Test coverage
• Add prices for Cohere: Add Cohere Embed v4.0 genai-prices#240
• Docs. Examples: OpenAI (regular, Azure, Ollama, OpenRouter); Sentence Transformers; Cohere

[DouweM]

Add a test like this one to verify this is up to date:

pydantic-ai/tests/models/test_model_names.py

Line 52 in efa1e26

def test_known_model_names(): # pragma: lax no cover

[github-actions]

Docs Preview

commit:	35abadc
Preview URL:	https://8ee72db3-pydantic-ai-previews.pydantic.workers.dev

[ggozad]

Thanks for starting this and please do let me know if you need help :)
I went quickly through, looks like a great start!

One thing you might want to support from the start is having as part of the EmbeddingSettings is max_context_length and encoding.

Embedding models have a limit of how many tokens of input they can handle. Most providers will raise (openai.BadRequestError iirc for OpenAI, vLLM will return an ugly 500 omg) and then some will say nothing (looking at you Ollama) and just truncate the input so that it fits.

All this is well explained here

I would not necessarily truncate like in the cookbook and still just raise, but I would be grateful to have available from the model side the max_context_length and the encoding so that as a library I can quickly check if a chunk of text fits or not.
Even better if I could get the number of tokens used for some text by a given embedding model.

The only difficulty I see with this is that not all providers expose the tokenizers, for example Ollama does not. But still, would be nice to have it for the providers that do support it, as it's a crucial step when you are trying to chunk a document for embedding.

In haiku.rag, my focus is local models, and like I mentioned Ollama, the popular choice, does not expose a way to tokenize text. So I just do the dumb thing and guesstimate the tokens hoping they are not going to be all that different from some OpenAI model's encoder: I use tiktoken (which you would probably also want to use to support this) and gpt-4o as a "close" model and get an estimate. But I am sure we can do better that this here.

Edit: I am not suggesting that calling embed should calculate the tokens needed on every call. But I imagine that whoever used pydantic AI to embed, would need to also go through the process of chunking some large text, unless they only dealt with embedding queries or simple sentences. So it would be a missed opportunity to not have support for that.

[gvanrossum]

I would like to be able to comment on the API, but there are no tests showing how to call it.

[DouweM]

@gvanrossum I'll make some progress on the PR today, but this is the API as it stands today:

import asyncio

from pydantic_ai.embeddings import Embedder

embedder = Embedder("openai:text-embedding-3-large")


async def main():
    result = await embedder.embed("Hello, world!")
    print(result)


if __name__ == "__main__":
    asyncio.run(main())

With Azure OpenAI you currently have to create the model and provider manually, but we'll make Embedder('azure:text-embedding-3-large') work as well:

import asyncio

from pydantic_ai.embeddings import Embedder
from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
from pydantic_ai.providers.azure import AzureProvider

model = OpenAIEmbeddingModel("text-embedding-3-large", provider=AzureProvider())

embedder = Embedder(model)


async def main():
    result = await embedder.embed("Hello, world!")
    print(result)


if __name__ == "__main__":
    asyncio.run(main())

[gvanrossum]

Nice. Do you have a bulk API too? That's essential for typeagent.

[DouweM]

@gvanrossum Yep, the embed method is overloaded to take either a str and return list[float], or take Sequence[str] and return list[list[float]], so it's the same method for single and bulk usage. (I'm aware str is itself a Sequence[str], but type checkers appear to handle the overloads correctly.)

[DouweM]

@gvanrossum In case you'd like to give it a try pre-release, I've made some progress today, including support for Embedder('azure:...').

[DouweM]

Unfortunately I haven't managed to get to this this week. Next week should be better.

[ggozad]

@ggozad @gvanrossum Just curious, what vector DB are you using? I'll want to have a RAG example in our docs.

I use lancedb

[gvanrossum]

@gvanrossum Just curious, what vector DB are you using?

I am using something I wrote myself. Persistence is optional (currently embeddings are stored in sqlite). Here's the code:
https://github.com/microsoft/typeagent-py/blob/main/typeagent/aitools/vectorbase.py
It stores a list of N normalized embeddings as a 2D numpy array of float32 with dimension (N, embedding_size). This makes for a fast dot product as long as all the vectors fit in memory. :-)

[ggozad]

Thank you @DouweM & @dmontagu ❤️

[tomaarsen]

Nice work on this!

[paulocoutinhox]

Hi,

It has support for local embeddings?

@staticmethod
def generate_embeddings(texts):
    import os

    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    from sentence_transformers import SentenceTransformer

    model = SentenceTransformer("all-MiniLM-L6-v2")
    return model.encode(texts, convert_to_numpy=True).tolist()
    ```

[daikeren]

Hi,

It has support for local embeddings?

@staticmethod
def generate_embeddings(texts):
    import os

    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    from sentence_transformers import SentenceTransformer

    model = SentenceTransformer("all-MiniLM-L6-v2")
    return model.encode(texts, convert_to_numpy=True).tolist()
    ```

Based on document, I think it has.
https://ai.pydantic.dev/embeddings/#sentence-transformers-local

[paulocoutinhox]

Thanks @daikeren, implemented here and it works.

[Mazyod]

Tiktoken has a subtle footgun that bites us occasionally, which is its lazy loading for tokenizers. We run in an air-gapped environment and best case scenario, it just fails immediately, and the worst case scenario, it hangs the process until it times out.

I haven't tested this update yet, but wanted to share this bit as I do not see any particular support for offline/air-gapped environments.

[paulocoutinhox]

Thanks @Mazyod.

In my case, I use it with internet access, so it downloads normally.

Also, since you brought up the subject, how do I save it to a Dockerfile so I don't have to download it again when I need to use it in the application?

[Mazyod]

@paulocoutinhox I did this approach for a while, till it failed because tiktoken cache was outdated and it attempted to refetch.

# Cache tiktoken definition from the internet
RUN http_proxy=http://my-proxy \
    https_proxy=http://my-proxy \
    python -c "import tiktoken; tiktoken.encoding_for_model('gpt-4o')"

[DouweM]

@Mazyod Good catch about the offline environments; can you please file an issue for that? We can at the very least document a workaround like that one for Docker?

[paulocoutinhox]

Yeah. Will be nice a solution for this. 💯

[stuaxo]

Does this work with bedrock yet ?

I haven't quite work it out - or maybe something isn't implemented yet ?

By way of example on bedrock (not using pydantic-ai yet) I'm using the embedding models:

"cohere.embed-english-v3" and "amazon.titan-embed-text-v2"

My assumption was that I would prefix these with "bedrock:".

Have the embeddings models available on bedrock been added anywhere in the codebase ?

So far, I get unknown model using various combinations.

[DouweM]

@stuartaxonHO The supported providers are documented under https://ai.pydantic.dev/embeddings/#providers; Bedrock is not yet one of them but contribution welcome!

[stuaxo]

@DouweM the word provider is a little overloaded - what should happen, to handle embedding models from different providers on an API provider ?

e.g. Cohere models on bedrock, vs the Amazon models (the model parameters are different and one uses embeddings[0] and the other embedding)

Also - in the code, I noticed there's a list of LLM models, but I can't find one of the embedding models, would that need to be added ?

[DouweM]

@stuaxo I agree the provider/model terms are a bit overloaded, we use them as described in https://ai.pydantic.dev/models/overview/#models-and-providers. That means that in this context, "model" maps to "API format", and "provider" maps to "API client/base URL".

So in this case we'd need a BedrockEmbeddingModel that uses their embedding API/SDK, which can be used with the existing BedrockProvider that provide the base URL/API client, and should then support any model name their API does, no matter what company actually build that model.