Pydantic integration with the BaseESModel class

Author: miguelgrinberg

elasticsearch/dsl/__init__.py

@@ -90,6 +90,7 @@
     IndexTemplate,
 )
 from .mapping import AsyncMapping, Mapping
+from .pydantic import BaseESModel
 from .query import Q, Query
 from .response import AggResponse, Response, UpdateByQueryResponse
 from .search import (
@@ -120,6 +121,7 @@
     "AsyncUpdateByQuery",
     "AttrDict",
     "AttrList",
+    "BaseESModel",
     "Binary",
     "Boolean",
     "Byte",

elasticsearch/dsl/pydantic.py

@@ -0,0 +1,57 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from typing import Any, ClassVar, Dict, Self, Tuple, Type
+
+from pydantic import BaseModel, Field, PrivateAttr
+from typing_extensions import Annotated, dataclass_transform
+
+from elasticsearch import dsl
+
+
+class _BaseModel(BaseModel):
+    meta: Annotated[Dict[str, Any], dsl.mapped_field(exclude=True)] = Field(default={})
+
+
+class BaseESModelMetaclass(type(BaseModel)):  # type: ignore[misc]
+    def __new__(cls, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]) -> Any:
+        model = super().__new__(cls, name, bases, attrs)
+        dsl_attrs = {
+            attr: value
+            for attr, value in dsl.AsyncDocument.__dict__.items()
+            if not attr.startswith("__")
+        }
+        model._doc = type(dsl.AsyncDocument)(  # type: ignore[misc]
+            f"_ES{name}",
+            dsl.AsyncDocument.__bases__,
+            {**attrs, **dsl_attrs, "__qualname__": f"_ES{name}"},
+        )
+        return model
+
+
+@dataclass_transform(kw_only_default=True, field_specifiers=(Field, PrivateAttr))
+class BaseESModel(_BaseModel, metaclass=BaseESModelMetaclass):
+    _doc: ClassVar[Type[dsl.AsyncDocument]]
+
+    def to_doc(self) -> dsl.AsyncDocument:
+        data = self.model_dump()
+        meta = {f"_{k}": v for k, v in data.pop("meta", {}).items()}
+        return self._doc(**meta, **data)
+
+    @classmethod
+    def from_doc(cls, dsl_obj: dsl.AsyncDocument) -> Self:
+        return cls(meta=dsl_obj.meta.to_dict(), **dsl_obj.to_dict())

elasticsearch/dsl/response/aggs.py

@@ -63,7 +63,7 @@ def _wrap_bucket(self, data: Dict[str, Any]) -> Bucket[_R]:
         )
 
     def __iter__(self) -> Iterator["Agg"]:  # type: ignore[override]
-        return iter(self.buckets)
+        return iter(self.buckets)  # type: ignore[arg-type]
 
     def __len__(self) -> int:
         return len(self.buckets)

examples/quotes/.gitignore

@@ -0,0 +1,24 @@
+# Logs
+logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+pnpm-debug.log*
+lerna-debug.log*
+
+node_modules
+dist
+dist-ssr
+*.local
+
+# Editor directories and files
+.vscode/*
+!.vscode/extensions.json
+.idea
+.DS_Store
+*.suo
+*.ntvs*
+*.njsproj
+*.sln
+*.sw?

examples/quotes/README.md

@@ -0,0 +1,94 @@
+# Quotes
+Quotes database example, which demonstrates the Elasticsearch integration with
+Pydantic models. This example features a React frontend and a FastAPI back end.
+
+![Quotes app screenshot](screenshot.png)
+
+## What is this?
+
+This repository contains a small application that demonstrates how easy it is
+to set up a full-text and vector database using [Elasticsearch](https://www.elastic.co/elasticsearch),
+while defining the data model with [Pydantic](https://docs.pydantic.dev/latest/).
+
+The application includes a FastAPI back end and a React front end. The example
+ingests a dataset of famous quotes into in an Elasticsearch index, and for each quote
+it generates an embedding using the
+[all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
+Sentence Transformers model.
+
+The dataset used by this application has about 37,000 famous quotes, each with
+their author and tags. The data originates from a
+[Kaggle dataset](https://www.kaggle.com/datasets/akmittal/quotes-dataset) that
+appears to have been generated from quotes that were scraped from the Goodreads
+[popular quotes](https://www.goodreads.com/quotes) page.
+
+## Requirements
+
+Please make sure you have the following installed:
+
+- Python 3.10 or newer
+- Node.js 18 or newer
+
+## How To Run
+
+Follow these steps to install this demo on your computer:
+
+### Clone this repository
+
+Run the following command to install a copy of this project on your computer:
+
+```bash
+git clone https://github.com/elastic/elasticsearch-py
+cd examples/quotes
+```
+
+### Install the Node and Python dependencies
+
+Run the following command to set up the JavaScript and Python environment and
+install all the dependencies:
+
+```bash
+npm install
+```
+
+### Start a development Elasticsearch container
+
+You can use [start-local](https://www.elastic.co/docs/deploy-manage/deploy/self-managed/local-development-installation-quickstart)
+to start a small Elasticsearch instance.
+
+Use this command to launch the instance (Docker and Docker Compose are required):
+
+```bash
+curl -fsSL https://elastic.co/start-local | sh
+```
+
+### Create the quotes database
+
+Run this command in your terminal:
+
+```bash
+npm run ingest
+```
+
+This task may take a few minutes. How long it takes depends on your computer
+speed and wether you have a GPU, which is used to generate the embeddings if
+available.
+
+### Start the back end
+
+Run this command in your terminal:
+
+```bash
+npm run backend
+```
+
+### Start the front end
+
+Open a second terminal window and run this command:
+
+```bash
+npm run dev
+```
+
+You can now navigate to `http://localhost:5173` on your web browser to access
+the application.

examples/quotes/backend/pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "quotes"
+version = "0.1"
+dependencies = [
+    "elasticsearch[async]>=8,<9",
+    "fastapi",
+    "orjson",
+    "sentence-transformers",
+    "uvicorn",
+]
+
+[project.optional-dependencies]
+dev = [
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"