Ho, ho, ho. Merry QA Christmas, one and all! 🎅

Santa @leycec has answers – but nobody is gonna like them. Sadly, my answers are all basically the exact same as @JWCS's. You have multiple options here. In descending order of badness:

• [EASY-PEASY] Just use standard dict[...] type hints subscripted by unions over all possible values that those dictionaries can contain. This approach is easy, because (A) it's perfectly supported by the most recent release of @beartype and (B) you don't need to define any additional classes. It looks something like this:

  def retrieval(
      self, question, embd_mdl, tenant_ids, kb_ids, page, page_size,
      similarity_threshold=0.2, vector_similarity_weight=0.3, top=1024,
      doc_ids=None, aggs=True, rerank_mdl=None, highlight=False
  
  # Something like this is what you want. A return dictionary type hint
  # mapping from strings to any of several possible different values,
  # including nested dictionaries also mapping from strings to integers.
  # I just made these nested type hints up, because I couldn't quite
  # parse the actual types from your example. But... something like this.
  ) -> dict[str, int | dict[str, int] | list[int] | set[list[str]]]:

• [HUMDRUM] The next less easy (but still viable) approach would be to bite the bullet and just begin encapsulating your data as proper @dataclasses.dataclass-decorated types. This approach is the next less easy, because (A) it's perfectly supported by the most recent release of @beartype but (B) you do need to define additional classes. It looks something like this:

  from dataclasses import dataclass
  
  @dataclass
  class DocAggs(object):
      count: int
      doc_id: int
      doc_name: int
  
  @dataclass
  class RetrievalRank(object):
      total: int
      chunks: list[int]
      doc_aggs: list[DocAggs]
  
  def retrieval(
      self, question, embd_mdl, tenant_ids, kb_ids, page, page_size,
      similarity_threshold=0.2, vector_similarity_weight=0.3, top=1024,
      doc_ids=None, aggs=True, rerank_mdl=None, highlight=False
  ) -> RetrievalRank:

• [SUPER-BAD] The worst possible approach would probably be to redefine your dictionaries as typing.TypedDict objects, because (A) @beartype doesn't deeply type-check typing.TypeDict objects at the moment and (B) @dataclasses.dataclass objects are a lot more useful than typing.TypedDict objects (even ignoring @beartype's poor support for typing.TypedDict).

Personally, I favour the @dataclasses.dataclass approach for this sort of thing. The syntax is pretty concise (albeit not quite as concise as the first approach using raw dict[...] type hints) and unambiguous (unlike the first approach using raw dict[...] type hints, which invites ambiguity over dictionary values).

@dataclasses.dataclass: when you don't know what else to do, it's Christmas Eve already, you're all hopped up on eggnog, and you begin to hear ominous sleighbells on the roof. 😆