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()
• Add prices for Cohere: https://github.com/pydantic/genai-prices/blob/main/prices/providers/cohere.yml
• Error handling (ModelAPIError)
• Test coverage
• Docs. Examples: OpenAI (regular, Azure, Ollama); 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

[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

[stuartaxonHO]

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

Nice.

For people doing local development OLlama and LlamaCPP (if you don't have it already) are essentials.

[github-actions]

Docs Preview

commit:	7392c38
Preview URL:	https://43fa6bd2-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.

[Kludex]

This should also be in Model - We don't have it right now.

[Kludex]

Not related to this PR.

[DouweM]

@Kludex What about models that don't take a model name, like TestModel, FunctionModel, WrapperModel?

[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.

[DouweM]

@ggozad You'll be happy to hear that EmbeddingModel has count_tokens and max_input_tokens methods that have been implemented for all of OpenAI (using tiktoken), Cohere, and SentenceTransformers! Does that give you what you're looking for?

[ggozad]

@ggozad You'll be happy to hear that EmbeddingModel has count_tokens and max_input_tokens methods that have been implemented for all of OpenAI (using tiktoken), Cohere, and SentenceTransformers! Does that give you what you're looking for?

I have been following along the PR and I just can't wait for a release! This looks really good, looking forward to removing a bunch of my code :)

[tomaarsen]

Following this PR now!
The modifications re. count_tokens and max_input_tokens look solid on the SentenceTransformer side, they match what I'd have done. Exciting!

• Tom Aarsen

[stuartaxonHO]

It might be nice to be able to do this as a single function call so you don't always need to create the embedder ahead of time, but I'm not sure if this fits with the rest of pydantic.ai ?

[DouweM]

@stuartaxonHO I personally don't think it's worth adding a helper function when the "verbose" version is just await Embedder('openai:text-embedding-3-large').embed_documents(['hello', 'world'])

[stuartaxonHO]

Think I was spoiled by the litellm version embeddings("somemodel", ["sometext", "some more text"]), you're right there's not much in it though.