object and nested fields

Author: miguelgrinberg

elasticsearch/dsl/pydantic.py

@@ -15,36 +15,122 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from typing import Any, ClassVar, Dict, Tuple, Type
+from typing import Any, ClassVar, Dict, List, Tuple, Type
 
 from pydantic import BaseModel, Field, PrivateAttr
 from typing_extensions import Annotated, Self, dataclass_transform
 
 from elasticsearch import dsl
 
 
+class ESMeta(BaseModel):
+    id: str = ""
+    index: str = ""
+    primary_term: int = 0
+    seq_no: int = 0
+    version: int = 0
+
+
 class _BaseModel(BaseModel):
-    meta: Annotated[Dict[str, Any], dsl.mapped_field(exclude=True)] = Field(default={})
+    meta: Annotated[ESMeta, dsl.mapped_field(exclude=True)] = Field(
+        default=ESMeta(), init=False
+    )
 
 
-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)
+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)
+    #     model._doc = cls.make_dsl_class(cls, dsl.AsyncDocument, model, attrs)
+    #     # dsl_attrs = {
+    #     #     attr: value
+    #     #     for attr, value in dsl.AsyncDocument.__dict__.items()
+    #     #     if not attr.startswith("__")
+    #     # }
+    #     # pydantic_attrs = {
+    #     #     **attrs,
+    #     #     "__annotations__": cls.process_annotations(cls, attrs["__annotations__"]),
+    #     # }
+    #     # model._doc = type(dsl.AsyncDocument)(  # type: ignore[misc]
+    #     #     f"_ES{name}",
+    #     #     dsl.AsyncDocument.__bases__,
+    #     #     {**pydantic_attrs, **dsl_attrs, "__qualname__": f"_ES{name}"},
+    #     # )
+    #     return model
+
+    @staticmethod
+    def process_annotations(cls, annotations):
+        updated_annotations = {}
+        for var, ann in annotations.items():
+            if isinstance(ann, type(BaseModel)):
+                # an inner Pydantic model is transformed into an Object field
+                updated_annotations[var] = cls.make_dsl_class(cls, dsl.InnerDoc, ann)
+            elif (
+                hasattr(ann, "__origin__")
+                and ann.__origin__ in [list, List]
+                and isinstance(ann.__args__[0], type(BaseModel))
+            ):
+                # an inner list of Pydantic models is transformed into a Nested field
+                updated_annotations[var] = list[
+                    cls.make_dsl_class(cls, dsl.InnerDoc, ann.__args__[0])
+                ]
+            else:
+                updated_annotations[var] = ann
+        return updated_annotations
+
+    @staticmethod
+    def make_dsl_class(cls, dsl_class, pydantic_model, pydantic_attrs=None):
         dsl_attrs = {
             attr: value
-            for attr, value in dsl.AsyncDocument.__dict__.items()
+            for attr, value in dsl_class.__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}"},
+        pydantic_attrs = {
+            **(pydantic_attrs or {}),
+            "__annotations__": cls.process_annotations(
+                cls, pydantic_model.__annotations__
+            ),
+        }
+        return type(dsl_class)(  # type: ignore[misc]
+            f"_ES{pydantic_model.__name__}",
+            (dsl_class,),
+            {
+                **pydantic_attrs,
+                **dsl_attrs,
+                "__qualname__": f"_ES{pydantic_model.__name__}",
+            },
         )
+
+
+class BaseESModelMetaclass(_BaseESModelMetaclass):  # type: ignore[misc]
+    def __new__(cls, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]) -> Any:
+        model = super().__new__(cls, name, bases, attrs)
+        model._doc = cls.make_dsl_class(cls, dsl.Document, model, attrs)
         return model
 
 
 @dataclass_transform(kw_only_default=True, field_specifiers=(Field, PrivateAttr))
 class BaseESModel(_BaseModel, metaclass=BaseESModelMetaclass):
+    _doc: ClassVar[Type[dsl.Document]]
+
+    def to_doc(self) -> dsl.Document:
+        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.Document) -> Self:
+        return cls(meta=ESMeta(**dsl_obj.meta.to_dict()), **dsl_obj.to_dict())
+
+
+class AsyncBaseESModelMetaclass(_BaseESModelMetaclass):  # type: ignore[misc]
+    def __new__(cls, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]) -> Any:
+        model = super().__new__(cls, name, bases, attrs)
+        model._doc = cls.make_dsl_class(cls, dsl.AsyncDocument, model, attrs)
+        return model
+
+
+@dataclass_transform(kw_only_default=True, field_specifiers=(Field, PrivateAttr))
+class AsyncBaseESModel(_BaseModel, metaclass=AsyncBaseESModelMetaclass):
     _doc: ClassVar[Type[dsl.AsyncDocument]]
 
     def to_doc(self) -> dsl.AsyncDocument:
@@ -54,4 +140,9 @@ def to_doc(self) -> dsl.AsyncDocument:
 
     @classmethod
     def from_doc(cls, dsl_obj: dsl.AsyncDocument) -> Self:
-        return cls(meta=dsl_obj.meta.to_dict(), **dsl_obj.to_dict())
+        return cls(meta=ESMeta(**dsl_obj.meta.to_dict()), **dsl_obj.to_dict())
+
+
+# TODO
+# - object and nested fields
+# - tests

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)  # type: ignore[arg-type]
+        return iter(self.buckets)
 
     def __len__(self) -> int:
         return len(self.buckets)

examples/quotes/README.md

@@ -6,18 +6,18 @@ Pydantic models. This example features a React frontend and a FastAPI back end.
 
 ## What is this?
 
-This repository contains a small application that demonstrates how easy it is
+This directory 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
+The application includes a FastAPI back end and a React front end. It 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
+The dataset 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.
@@ -62,6 +62,14 @@ Use this command to launch the instance (Docker and Docker Compose are required)
 curl -fsSL https://elastic.co/start-local | sh
 ```
 
+Once your Elasticsearch instead is deployed, create an environment variable called
+`ELASTICSEARCH_URL`, making sure it includes the password generated by start-local.
+Example:
+
+```bash
+export ELASTICSEARCH_URL=http://elastic:your-password-here@localhost:9200
+```
+
 ### Create the quotes database
 
 Run this command in your terminal:
@@ -70,6 +78,9 @@ Run this command in your terminal:
 npm run ingest
 ```
 
+Note that the `ELASTICSEARCH_URL` variable must be defined in the terminal
+session in which you run this command.
+
 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.
@@ -82,6 +93,9 @@ Run this command in your terminal:
 npm run backend
 ```
 
+Note that the `ELASTICSEARCH_URL` variable must be defined in the terminal
+session in which you run this command.
+
 ### Start the front end
 
 Open a second terminal window and run this command:

examples/quotes/backend/quotes.py

@@ -4,18 +4,19 @@
 from time import time
 from typing import Annotated
 
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, Field
 from sentence_transformers import SentenceTransformer
 
-from elasticsearch.dsl.pydantic import BaseESModel
+from elasticsearch import NotFoundError
+from elasticsearch.dsl.pydantic import AsyncBaseESModel
 from elasticsearch import dsl
 
 model = SentenceTransformer("all-MiniLM-L6-v2")
 dsl.async_connections.create_connection(hosts=[os.environ['ELASTICSEARCH_URL']])
 
 
-class Quote(BaseESModel):
+class Quote(AsyncBaseESModel):
     quote: str
     author: Annotated[str, dsl.Keyword()]
     tags: Annotated[list[str], dsl.Keyword()]
@@ -44,11 +45,49 @@ class SearchResponse(BaseModel):
     total: int
 
 
-app = FastAPI()
+app = FastAPI(
+    title="Quotes API",
+    version="1.0.0",
+)
+
+@app.get("/api/quotes/{id}")
+async def get_quote(id: str) -> Quote:
+    doc = None
+    try:
+        doc = await Quote._doc.get(id)
+    except NotFoundError:
+        pass
+    if not doc:
+        raise HTTPException(status_code=404, detail="Item not found")
+    return Quote.from_doc(doc)
+
+
+@app.post("/api/quotes", status_code=201)
+async def create_quote(req: Quote) -> Quote:
+    doc = req.to_doc()
+    doc.meta.id = ""
+    await doc.save(refresh=True)
+    return Quote.from_doc(doc)
+
+
+@app.put("/api/quotes/{id}")
+async def update_quote(id: str, req: Quote) -> Quote:
+    doc = req.to_doc()
+    doc.meta.id = id
+    await doc.save(refresh=True)
+    return Quote.from_doc(doc)
+
+
+@app.delete("/api/quotes/{id}", status_code=204)
+async def delete_quote(id: str, req: Quote) -> None:
+    doc = await Quote._doc.get(id)
+    if not doc:
+        raise HTTPException(status_code=404, detail="Item not found")
+    await doc.delete(refresh=True)
 
 
 @app.post('/api/search')
-async def search(req: SearchRequest) -> SearchResponse:
+async def search_quotes(req: SearchRequest) -> SearchResponse:
     quotes, tags, total = await search_quotes(req.query, req.filters, use_knn=req.knn, start=req.start)
     return SearchResponse(
         quotes=quotes,

examples/quotes/vite.config.ts

@@ -7,6 +7,9 @@ export default defineConfig({
   server: {
     proxy: {
       '/api': 'http://localhost:5000',
+      '/docs': 'http://localhost:5000',
+      '/redoc': 'http://localhost:5000',
+      '/openapi.json': 'http://localhost:5000',
     },
   },
 })