Draft implementation of support for embeddings APIs

[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

[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

[DouweM]

We'll have to check which of the OpenAI-compatible APIs also support embeddings

[DouweM]

https://github.com/ggozad/haiku.rag/tree/main/src/haiku/rag/embeddings has Ollama, vLLM and VoyageAI, which would be worth adding as well

[DouweM]

Suggested change

	@dataclass
	@dataclass(init=False)

[DouweM]

Update to https://docs.cohere.com/docs/cohere-embed

[DouweM]

I'd prefer to move this to __init__

[DouweM]

https://platform.openai.com/docs/guides/embeddings#embedding-models

[DouweM]

Drop

[DouweM]

Not sure how I feel about every model implementation needing to repeat this

[DouweM]

Following the pattern in ModelSettings, we should move any options only supported by one model to {Cohere}EmbeddingSettings with a {cohere}_ prefix

[github-actions]

Docs Preview

commit:	3dbad0d
Preview URL:	https://fda6de3f-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.