fix: simplify ContextVar and fix sub-span attribution when delegate…

…d to thread (#13700)

* fix: simplify contextvar

* test trees in parallel

* fix Token.MISSING

llama-index-core/llama_index/core/agent/legacy/react/base.py

@@ -1,6 +1,5 @@
 import asyncio
 from itertools import chain
-from threading import Thread
 from typing import (
     Any,
     AsyncGenerator,
@@ -42,6 +41,7 @@
 from llama_index.core.settings import Settings
 from llama_index.core.tools import BaseTool, ToolOutput, adapt_to_async_tool
 from llama_index.core.tools.types import AsyncBaseTool
+from llama_index.core.types import Thread
 from llama_index.core.utils import print_text, unit_generator
 
 

llama-index-core/llama_index/core/agent/react/step.py

@@ -3,7 +3,6 @@
 import asyncio
 import uuid
 from functools import partial
-from threading import Thread
 from typing import (
     Any,
     AsyncGenerator,
@@ -53,6 +52,7 @@
 from llama_index.core.settings import Settings
 from llama_index.core.tools import BaseTool, ToolOutput, adapt_to_async_tool
 from llama_index.core.tools.types import AsyncBaseTool
+from llama_index.core.types import Thread
 from llama_index.core.utils import print_text
 
 

llama-index-core/llama_index/core/agent/runner/base.py

@@ -392,9 +392,9 @@ def _run_step(
         **kwargs: Any,
     ) -> TaskStepOutput:
         """Execute step."""
-        dispatch_event = dispatcher.get_dispatch_event()
-
-        dispatch_event(AgentRunStepStartEvent(task_id=task_id, step=step, input=input))
+        dispatcher.event(
+            AgentRunStepStartEvent(task_id=task_id, step=step, input=input)
+        )
         task = self.state.get_task(task_id)
         step_queue = self.state.get_step_queue(task_id)
         step = step or step_queue.popleft()
@@ -421,7 +421,7 @@ def _run_step(
         completed_steps = self.state.get_completed_steps(task_id)
         completed_steps.append(cur_step_output)
 
-        dispatch_event(AgentRunStepEndEvent(step_output=cur_step_output))
+        dispatcher.event(AgentRunStepEndEvent(step_output=cur_step_output))
         return cur_step_output
 
     @dispatcher.span
@@ -434,9 +434,9 @@ async def _arun_step(
         **kwargs: Any,
     ) -> TaskStepOutput:
         """Execute step."""
-        dispatch_event = dispatcher.get_dispatch_event()
-
-        dispatch_event(AgentRunStepStartEvent(task_id=task_id, step=step, input=input))
+        dispatcher.event(
+            AgentRunStepStartEvent(task_id=task_id, step=step, input=input)
+        )
         task = self.state.get_task(task_id)
         step_queue = self.state.get_step_queue(task_id)
         step = step or step_queue.popleft()
@@ -462,7 +462,7 @@ async def _arun_step(
         completed_steps = self.state.get_completed_steps(task_id)
         completed_steps.append(cur_step_output)
 
-        dispatch_event(AgentRunStepEndEvent(step_output=cur_step_output))
+        dispatcher.event(AgentRunStepEndEvent(step_output=cur_step_output))
         return cur_step_output
 
     @dispatcher.span
@@ -561,14 +561,12 @@ def _chat(
         mode: ChatResponseMode = ChatResponseMode.WAIT,
     ) -> AGENT_CHAT_RESPONSE_TYPE:
         """Chat with step executor."""
-        dispatch_event = dispatcher.get_dispatch_event()
-
         if chat_history is not None:
             self.memory.set(chat_history)
         task = self.create_task(message)
 
         result_output = None
-        dispatch_event(AgentChatWithStepStartEvent(user_msg=message))
+        dispatcher.event(AgentChatWithStepStartEvent(user_msg=message))
         while True:
             # pass step queue in as argument, assume step executor is stateless
             cur_step_output = self._run_step(
@@ -586,7 +584,7 @@ def _chat(
             task.task_id,
             result_output,
         )
-        dispatch_event(AgentChatWithStepEndEvent(response=result))
+        dispatcher.event(AgentChatWithStepEndEvent(response=result))
         return result
 
     @dispatcher.span
@@ -598,14 +596,12 @@ async def _achat(
         mode: ChatResponseMode = ChatResponseMode.WAIT,
     ) -> AGENT_CHAT_RESPONSE_TYPE:
         """Chat with step executor."""
-        dispatch_event = dispatcher.get_dispatch_event()
-
         if chat_history is not None:
             self.memory.set(chat_history)
         task = self.create_task(message)
 
         result_output = None
-        dispatch_event(AgentChatWithStepStartEvent(user_msg=message))
+        dispatcher.event(AgentChatWithStepStartEvent(user_msg=message))
         while True:
             # pass step queue in as argument, assume step executor is stateless
             cur_step_output = await self._arun_step(
@@ -623,7 +619,7 @@ async def _achat(
             task.task_id,
             result_output,
         )
-        dispatch_event(AgentChatWithStepEndEvent(response=result))
+        dispatcher.event(AgentChatWithStepEndEvent(response=result))
         return result
 
     @dispatcher.span
@@ -778,16 +774,14 @@ def _chat(
         mode: ChatResponseMode = ChatResponseMode.WAIT,
     ) -> AGENT_CHAT_RESPONSE_TYPE:
         """Chat with step executor."""
-        dispatch_event = dispatcher.get_dispatch_event()
-
         if chat_history is not None:
             self.memory.set(chat_history)
 
         # create initial set of tasks
         plan_id = self.create_plan(message)
 
         results = []
-        dispatch_event(AgentChatWithStepStartEvent(user_msg=message))
+        dispatcher.event(AgentChatWithStepStartEvent(user_msg=message))
         while True:
             # EXIT CONDITION: check if all sub-tasks are completed
             next_task_ids = self.get_next_tasks(plan_id)
@@ -812,7 +806,7 @@ def _chat(
             # refine the plan
             self.refine_plan(message, plan_id)
 
-        dispatch_event(
+        dispatcher.event(
             AgentChatWithStepEndEvent(
                 response=results[-1] if len(results) > 0 else None
             )
@@ -828,16 +822,14 @@ async def _achat(
         mode: ChatResponseMode = ChatResponseMode.WAIT,
     ) -> AGENT_CHAT_RESPONSE_TYPE:
         """Chat with step executor."""
-        dispatch_event = dispatcher.get_dispatch_event()
-
         if chat_history is not None:
             self.memory.set(chat_history)
 
         # create initial set of tasks
         plan_id = self.create_plan(message)
 
         results = []
-        dispatch_event(AgentChatWithStepStartEvent(user_msg=message))
+        dispatcher.event(AgentChatWithStepStartEvent(user_msg=message))
         while True:
             # EXIT CONDITION: check if all sub-tasks are completed
             next_task_ids = self.get_next_tasks(plan_id)
@@ -862,7 +854,7 @@ async def _achat(
             # refine the plan
             await self.arefine_plan(message, plan_id)
 
-        dispatch_event(
+        dispatcher.event(
             AgentChatWithStepEndEvent(
                 response=results[-1] if len(results) > 0 else None
             )

llama-index-core/llama_index/core/async_utils.py

@@ -120,10 +120,9 @@ async def run_jobs(
         List[Any]:
             List of results.
     """
-    parent_span_id = dispatcher.current_span_id
     semaphore = asyncio.Semaphore(workers)
 
-    @dispatcher.async_span_with_parent_id(parent_id=parent_span_id)
+    @dispatcher.span
     async def worker(job: Coroutine) -> Any:
         async with semaphore:
             return await job

llama-index-core/llama_index/core/base/base_query_engine.py

@@ -44,26 +44,26 @@ def _update_prompts(self, prompts: PromptDictType) -> None:
 
     @dispatcher.span
     def query(self, str_or_query_bundle: QueryType) -> RESPONSE_TYPE:
-        dispatch_event = dispatcher.get_dispatch_event()
-
-        dispatch_event(QueryStartEvent(query=str_or_query_bundle))
+        dispatcher.event(QueryStartEvent(query=str_or_query_bundle))
         with self.callback_manager.as_trace("query"):
             if isinstance(str_or_query_bundle, str):
                 str_or_query_bundle = QueryBundle(str_or_query_bundle)
             query_result = self._query(str_or_query_bundle)
-        dispatch_event(QueryEndEvent(query=str_or_query_bundle, response=query_result))
+        dispatcher.event(
+            QueryEndEvent(query=str_or_query_bundle, response=query_result)
+        )
         return query_result
 
     @dispatcher.span
     async def aquery(self, str_or_query_bundle: QueryType) -> RESPONSE_TYPE:
-        dispatch_event = dispatcher.get_dispatch_event()
-
-        dispatch_event(QueryStartEvent(query=str_or_query_bundle))
+        dispatcher.event(QueryStartEvent(query=str_or_query_bundle))
         with self.callback_manager.as_trace("query"):
             if isinstance(str_or_query_bundle, str):
                 str_or_query_bundle = QueryBundle(str_or_query_bundle)
             query_result = await self._aquery(str_or_query_bundle)
-        dispatch_event(QueryEndEvent(query=str_or_query_bundle, response=query_result))
+        dispatcher.event(
+            QueryEndEvent(query=str_or_query_bundle, response=query_result)
+        )
         return query_result
 
     def retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:

llama-index-core/llama_index/core/base/base_retriever.py

@@ -224,10 +224,8 @@ def retrieve(self, str_or_query_bundle: QueryType) -> List[NodeWithScore]:
                 a QueryBundle object.
 
         """
-        dispatch_event = dispatcher.get_dispatch_event()
-
         self._check_callback_manager()
-        dispatch_event(
+        dispatcher.event(
             RetrievalStartEvent(
                 str_or_query_bundle=str_or_query_bundle,
             )
@@ -246,7 +244,7 @@ def retrieve(self, str_or_query_bundle: QueryType) -> List[NodeWithScore]:
                 retrieve_event.on_end(
                     payload={EventPayload.NODES: nodes},
                 )
-        dispatch_event(
+        dispatcher.event(
             RetrievalEndEvent(
                 str_or_query_bundle=str_or_query_bundle,
                 nodes=nodes,
@@ -257,9 +255,8 @@ def retrieve(self, str_or_query_bundle: QueryType) -> List[NodeWithScore]:
     @dispatcher.span
     async def aretrieve(self, str_or_query_bundle: QueryType) -> List[NodeWithScore]:
         self._check_callback_manager()
-        dispatch_event = dispatcher.get_dispatch_event()
 
-        dispatch_event(
+        dispatcher.event(
             RetrievalStartEvent(
                 str_or_query_bundle=str_or_query_bundle,
             )
@@ -280,7 +277,7 @@ async def aretrieve(self, str_or_query_bundle: QueryType) -> List[NodeWithScore]
                 retrieve_event.on_end(
                     payload={EventPayload.NODES: nodes},
                 )
-        dispatch_event(
+        dispatcher.event(
             RetrievalEndEvent(
                 str_or_query_bundle=str_or_query_bundle,
                 nodes=nodes,