Retain return type from @dispatcher.span

[GICodeWarrior]

Attempt to fix #16915

Version Bump?

Did I bump the version in the pyproject.toml file of the package I am updating? (Except for the llama-index-core package)

• Yes
• No

Type of Change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)

How Has This Been Tested?

Your pull-request will likely not be merged unless it is covered by some form of impactful unit testing.

• I added new unit tests to cover this change
• I believe this change is already covered by existing unit tests

Suggested Checklist:

• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• I have added Google Colab support for the newly added notebooks.
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I ran make format; make lint to appease the lint gods

[GICodeWarrior]

OK, I believe the type checking is working better now.

However, this has uncovered typing issues in 11 other files. 🤔

[GICodeWarrior]

Some of these might be actual bugs.

For example, RESPONSE_TYPE can be AsyncStreamingResponse, and if it is, the get_response call on line 47 below needs to be awaited, right?

llama_index/llama-index-core/llama_index/core/query_engine/retry_source_query_engine.py

Lines 42 to 48 in f43e337

	def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
	response = self._query_engine._query(query_bundle)
	if self.max_retries <= 0:
	return response
	typed_response = (
	response if isinstance(response, Response) else response.get_response()
	)

llama_index/llama-index-core/llama_index/core/base/response/schema.py

Lines 234 to 236 in f43e337

	RESPONSE_TYPE = Union[
	Response, StreamingResponse, AsyncStreamingResponse, PydanticResponse
	]

[GICodeWarrior]

Note, commit 9303917 is a pretty good branching point before I went down a deep rabbit hole. (8 mypy errors remaining there)

[GICodeWarrior]

d54996e gets back to 7 mypy errors (better than before the rabbit hole)

[GICodeWarrior]

For these remaining mypy errors, I'm not sure I know the code well enough to determine whether it's best to "fix" the types, squelch the error (with a cast or type ignore), or if there are potential bugs in the code.

llama_index/core/indices/common_tree/base.py:222: error: Incompatible types in assignment (expression has type "tuple[str]", variable has type "list[tuple[str, str]]")  [assignment]
llama_index/core/tools/retriever_tool.py:84: error: "BaseNode" has no attribute "text_template"  [attr-defined]
llama_index/core/tools/retriever_tool.py:109: error: "BaseNode" has no attribute "text_template"  [attr-defined]
llama_index/core/extractors/metadata_extractors.py:503: error: Item "list[Any]" of "Any | list[Any]" has no attribute "dict"  [union-attr]
llama_index/core/evaluation/retrieval/evaluator.py:48: error: "BaseNode" has no attribute "text"  [attr-defined]
llama_index/core/evaluation/retrieval/evaluator.py:87: error: "BaseNode" has no attribute "text"  [attr-defined]
llama_index/core/evaluation/retrieval/evaluator.py:88: error: Argument 1 to "append" of "list" has incompatible type "BaseNode"; expected "TextNode"  [arg-type]

https://github.com/run-llama/llama_index/actions/runs/13335762370/job/37250483650?pr=17817

[GICodeWarrior]

Reworked commits to be a little more cohesive.

[GICodeWarrior]

Fixed the remaining type checking issues.

The llama-index-node-parser-docling unit test failures doesn't make sense to me, so that will take some digging.

[GICodeWarrior]

The llama-index-node-parser-docling unit test failure appears to be due to a change in the docling project.

https://github.com/DS4SD/docling/blob/e1436a8b0574e6bb2bb89bd65e98221e418d7142/CHANGELOG.md#v2210---2025-02-10

[GICodeWarrior]

• Fixed unrelated unit test failure.
• Re-reviewed changes.
• Tidied commit history again

[logan-markewich]

I'm not sure why your type checking uncovered so many errors. Our CICD runs type checking on core in every PR, and these haven't been issues.

For example, RESPONSE_TYPE can be AsyncStreamingResponse, and if it is, the get_response call on line 47 below needs to be awaited, right?

Actually no. More-so a symptom of lazy typing lol. a synchronous query will never return an async response

I appreciate the type fixes, but since they weren't an issue in CICD before, I'd feel better about this change if we just split out the actual changes to the dispatcher 🤔 Lots of changes flying around at the moment

[GICodeWarrior]

The @dispatcher.span fix uncovered most of these typing issues.

With some additional effort, it should be possible to shrink the changes in 742c866 if necessary. There were a few typing issues uncovered by the @dispatcher.span fix that veered into that realm, but it should be possible to cut those off with some # type: ignore, etc.

The unit test issue was an artifact of the build system choosing a newer version of docling. It should be possible to fix that in a separate PR, as long as that PR is merged before this one.

[GICodeWarrior]

In fact, your PR #17831 is hitting the same unit test failure I fixed here.

[logan-markewich]

Yea I've been ignoring that docling test (the test itself is a little silly lol)

Ok fair, that makes sense if the dispatcher was swallowing types. Let me take a deeper look at these changes to make sure there isn't a breaking change...