Repository: run-llama/llama_index

Issue:
[Bug]: Gemini Client does not have callbacks for async methods (achat, acomplete)
### Bug Description

Hi,

We noticed that we were not receiving instrumentation when using Gemini models. 

After digging a bit in the code, it seems the `.achat` and `.acomplete` method in the Gemini client ([/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/base.py](https://github.com/run-llama/llama_index/blob/d22fdcd2c1d09bfa66b2370b8ea2d5680d89cd63/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/base.py)) do not have the `@llm_completion_callback()` and `@llm_chat_callback()`. What is strange is that the `.complete` and `. chat` method have the decorators.

Is there a particular reason why this has not beed added? Maybe the original reason is not valid anymore and they can be simply added?

Happy to open a PR if that is a simple fix

### Version

0.12.14

### Steps to Reproduce

1. Setup some instrumentation (e.g. a custom EventHandler following [the doc](https://docs.llamaindex.ai/en/stable/module_guides/observability/instrumentation/#using-the-instrumentation-module-for-observability)) 
2. Make an async query with a Gemini model
``` 
from llama_index.llms.gemini import Gemini

flash2 = Gemini(model="models/gemini-2.0-flash-exp")
await flash2.acomplete("Write your best joke about pirates.")
```

3. Run and see that the instrumentation is never called. Models like OpenAI trigger the instrumentation.

### Relevant Logs/Tracbacks

(no failures, just nothing happens)

Hints from issue comments:
@dosubot: <!-- Greeting -->
Hey @Digma! 👋 I'm here to help with bugs, questions, and contributions. Let's squash some bugs together!

<!-- Answer -->
The `.achat` and `.acomplete` methods in the Gemini client are indeed decorated with `@llm_chat_callback()` and `@llm_completion_callback()`, respectively. This ensures that the callback functionality is applied to these asynchronous methods, allowing them to handle chat and completion tasks with the necessary callback support <sup>[[1]](https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/llms/llama-index-llms-ollama/llama_index/llms/ollama/base.py)</sup><sup>[[2]](https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/llms/llama-index-llms-clarifai/llama_index/llms/clarifai/base.py)</sup>. 

If you are not receiving instrumentation, it might be worth checking if the instrumentation setup is correctly configured or if there are any issues specific to the version you are using. Additionally, ensure that the callbacks are properly registered and that there are no errors in the instrumentation logic itself.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*

---
<sup>Help Dosu learn! Give it feedback: [Great Response](https://app.dosu.dev/response-feedback/0a180c3c-da07-4f67-aab4-db870bd7e2c7?feedback_type=great_response) | [Irrelevant Answer](https://app.dosu.dev/response-feedback/0a180c3c-da07-4f67-aab4-db870bd7e2c7?feedback_type=irrelevant_answer) | [Incorrect Sources](https://app.dosu.dev/response-feedback/0a180c3c-da07-4f67-aab4-db870bd7e2c7?feedback_type=incorrect_sources) | [Too Verbose](https://app.dosu.dev/response-feedback/0a180c3c-da07-4f67-aab4-db870bd7e2c7?feedback_type=too_verbose) | [Hallucination](https://app.dosu.dev/response-feedback/0a180c3c-da07-4f67-aab4-db870bd7e2c7?feedback_type=hallucination) | [Other](https://app.dosu.dev/response-feedback/0a180c3c-da07-4f67-aab4-db870bd7e2c7?feedback_type=other)</sup>
@Digma: @dosu 
```
The .achat and .acomplete methods in the Gemini client are indeed decorated with @llm_chat_callback() and @llm_completion_callback()
```
This is incorrect, the decorators are missing. Here is an extract of the latest version of the file:
``` 
    @llm_completion_callback()
    def complete(
        self, prompt: str, formatted: bool = False, **kwargs: Any
    ) -> CompletionResponse:
...
        return completion_from_gemini_response(result)

    async def acomplete(
        self, prompt: str, formatted: bool = False, **kwargs: Any
    ) -> CompletionResponse:
...
        return completion_from_gemini_response(result)
```
Where we clearly see the missing decorator@dosubot: <!-- Answer -->
The absence of the `@llm_completion_callback()` and `@llm_chat_callback()` decorators on the `acomplete` and `achat` methods in the Gemini client seems to be an oversight, especially since the synchronous methods have these decorators. These decorators are important for enabling callback functionality, which is crucial for instrumentation and logging purposes <sup>[[1]](https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/llms/llama-index-llms-gigachat/llama_index/llms/gigachat/base.py)</sup><sup>[[2]](https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/llms/structured_llm.py)</sup><sup>[[3]](https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/multi_modal_llms/base.py)</sup>.

If there are no specific reasons documented for their absence, it might be worth considering adding them to ensure consistent behavior with other methods. You could proceed with opening a pull request to add these decorators, assuming there are no underlying reasons for their exclusion.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*

---
<sup>Help Dosu learn! Give it feedback: [Great Response](https://app.dosu.dev/response-feedback/8094bcb4-fcbd-4bf8-9ef2-180215c448e2?feedback_type=great_response) | [Irrelevant Answer](https://app.dosu.dev/response-feedback/8094bcb4-fcbd-4bf8-9ef2-180215c448e2?feedback_type=irrelevant_answer) | [Incorrect Sources](https://app.dosu.dev/response-feedback/8094bcb4-fcbd-4bf8-9ef2-180215c448e2?feedback_type=incorrect_sources) | [Too Verbose](https://app.dosu.dev/response-feedback/8094bcb4-fcbd-4bf8-9ef2-180215c448e2?feedback_type=too_verbose) | [Hallucination](https://app.dosu.dev/response-feedback/8094bcb4-fcbd-4bf8-9ef2-180215c448e2?feedback_type=hallucination) | [Other](https://app.dosu.dev/response-feedback/8094bcb4-fcbd-4bf8-9ef2-180215c448e2?feedback_type=other)</sup>
@logan-markewich: I think this was just missed tbh -- feel free to make a pr!@logan-markewich: I think this was just missed tbh -- feel free to make a pr!@Riddhimaan-Senapati: I created a PR to solve this issue.@Digma: @Riddhimaan-Senapati You were faster than me :)

Solution patch:
diff --git a/docs/docs/examples/retrievers/bm25_retriever.ipynb b/docs/docs/examples/retrievers/bm25_retriever.ipynb
index 27443b4ec..7c8cbb47d 100644
--- a/docs/docs/examples/retrievers/bm25_retriever.ipynb
+++ b/docs/docs/examples/retrievers/bm25_retriever.ipynb
@@ -621,7 +621,7 @@
    "source": [
     "storage_context.docstore.persist(\"./docstore.json\")\n",
     "\n",
-    "# or, we could ignore the docstore and just persist the bm25 retriever as shown above\n",
+    "# or, we could ignore the docstore and just persist the bm25 retriever as shown below\n",
     "# bm25_retriever.persist(\"./bm25_retriever\")"
    ]
   },
diff --git a/llama-index-core/llama_index/core/tools/utils.py b/llama-index-core/llama_index/core/tools/utils.py
index 73349fe6e..22e7e8adb 100644
--- a/llama-index-core/llama_index/core/tools/utils.py
+++ b/llama-index-core/llama_index/core/tools/utils.py
@@ -42,6 +42,8 @@ def create_schema_from_function(
             param_type = args[0]
             if isinstance(args[1], str):
                 description = args[1]
+            elif isinstance(args[1], FieldInfo):
+                description = args[1].description
 
         if param_type is params[param_name].empty:
             param_type = Any
diff --git a/llama-index-core/llama_index/core/workflow/context.py b/llama-index-core/llama_index/core/workflow/context.py
index 10ad0eca8..f46f05b8a 100644
--- a/llama-index-core/llama_index/core/workflow/context.py
+++ b/llama-index-core/llama_index/core/workflow/context.py
@@ -1,14 +1,14 @@
 import asyncio
 import json
-import warnings
 import uuid
+import warnings
 from collections import defaultdict
-from typing import Dict, Any, Optional, List, Type, TYPE_CHECKING, Set, Tuple, TypeVar
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Type, TypeVar
 
 from .context_serializers import BaseSerializer, JsonSerializer
 from .decorators import StepConfig
-from .events import Event
 from .errors import WorkflowRuntimeError
+from .events import Event
 
 if TYPE_CHECKING:  # pragma: no cover
     from .workflow import Workflow
@@ -60,7 +60,7 @@ class Context:
         self._lock = asyncio.Lock()
         self._globals: Dict[str, Any] = {}
         # Step-specific instance
-        self._events_buffer: Dict[Type[Event], List[Event]] = defaultdict(list)
+        self._events_buffer: Dict[str, List[Event]] = defaultdict(list)
 
     def _serialize_queue(self, queue: asyncio.Queue, serializer: BaseSerializer) -> str:
         queue_items = list(queue._queue)  # type: ignore
@@ -240,14 +240,17 @@ class Context:
         warnings.warn(msg, DeprecationWarning)
         return self
 
+    def _get_full_path(self, ev_type: Type[Event]) -> str:
+        return f"{ev_type.__module__}.{ev_type.__name__}"
+
     def collect_events(
         self, ev: Event, expected: List[Type[Event]]
     ) -> Optional[List[Event]]:
-        self._events_buffer[type(ev)].append(ev)
+        self._events_buffer[self._get_full_path(type(ev))].append(ev)
 
         retval: List[Event] = []
         for e_type in expected:
-            e_instance_list = self._events_buffer.get(e_type)
+            e_instance_list = self._events_buffer.get(self._get_full_path(e_type))
             if e_instance_list:
                 retval.append(e_instance_list.pop(0))
 
@@ -256,7 +259,7 @@ class Context:
 
         # put back the events if unable to collect all
         for ev in retval:
-            self._events_buffer[type(ev)].append(ev)
+            self._events_buffer[self._get_full_path(type(ev))].append(ev)
 
         return None
 
diff --git a/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/base.py b/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/base.py
index 0a133894c..21ce63db0 100644
--- a/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/base.py
+++ b/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/base.py
@@ -192,6 +192,7 @@ class Gemini(CustomLLM):
         )
         return completion_from_gemini_response(result)
 
+    @llm_completion_callback()
     async def acomplete(
         self, prompt: str, formatted: bool = False, **kwargs: Any
     ) -> CompletionResponse:
@@ -201,6 +202,7 @@ class Gemini(CustomLLM):
         )
         return completion_from_gemini_response(result)
 
+    @llm_completion_callback()
     def stream_complete(
         self, prompt: str, formatted: bool = False, **kwargs: Any
     ) -> CompletionResponseGen:
@@ -223,6 +225,7 @@ class Gemini(CustomLLM):
         )
         return chat_from_gemini_response(response)
 
+    @llm_chat_callback()
     async def achat(
         self, messages: Sequence[ChatMessage], **kwargs: Any
     ) -> ChatResponse:
@@ -268,6 +271,7 @@ class Gemini(CustomLLM):
 
         return gen()
 
+    @llm_chat_callback()
     async def astream_chat(
         self, messages: Sequence[ChatMessage], **kwargs: Any
     ) -> ChatResponseAsyncGen:


Test patch:
diff --git a/llama-index-core/tests/tools/test_utils.py b/llama-index-core/tests/tools/test_utils.py
index 6ca980ba7..8b13708e3 100644
--- a/llama-index-core/tests/tools/test_utils.py
+++ b/llama-index-core/tests/tools/test_utils.py
@@ -73,3 +73,25 @@ def test_create_schema_from_function_with_typing_annotated() -> None:
 
     instance = schema(x=5)
     assert instance.x == 5  # type: ignore
+
+
+def test_create_schema_from_function_with_field_annotated() -> None:
+    """Test create_schema_from_function with Annotated[pydantic.Field]."""
+
+    def tmp_function(x: Annotated[int, Field(description="An integer")] = 3) -> str:
+        return str(x)
+
+    schema = create_schema_from_function("TestSchema", tmp_function)
+    actual_schema = schema.model_json_schema()
+
+    assert "x" in actual_schema["properties"]
+    assert actual_schema["properties"]["x"]["type"] == "integer"
+    assert actual_schema["properties"]["x"]["default"] == 3
+    assert actual_schema["properties"]["x"]["description"] == "An integer"
+
+    # Test the created schema
+    instance = schema()
+    assert instance.x == 3  # type: ignore
+
+    instance = schema(x=5)
+    assert instance.x == 5  # type: ignore
diff --git a/llama-index-core/tests/workflow/test_context.py b/llama-index-core/tests/workflow/test_context.py
index 1cb35c237..41afa9476 100644
--- a/llama-index-core/tests/workflow/test_context.py
+++ b/llama-index-core/tests/workflow/test_context.py
@@ -1,18 +1,18 @@
 import asyncio
+import json
+from typing import Optional, Union
 from unittest import mock
-from typing import Union, Optional
 
 import pytest
+from llama_index.core.workflow.decorators import step
+from llama_index.core.workflow.errors import WorkflowRuntimeError
+from llama_index.core.workflow.events import Event, StartEvent, StopEvent
 from llama_index.core.workflow.workflow import (
-    Workflow,
     Context,
+    Workflow,
 )
-from llama_index.core.workflow.decorators import step
-from llama_index.core.workflow.errors import WorkflowRuntimeError
-from llama_index.core.workflow.events import StartEvent, StopEvent, Event
-from llama_index.core.workflow.workflow import Workflow
 
-from .conftest import OneTestEvent, AnotherTestEvent
+from .conftest import AnotherTestEvent, OneTestEvent
 
 
 @pytest.mark.asyncio()
@@ -115,6 +115,11 @@ def test_get_result(ctx):
     assert ctx.get_result() == 42
 
 
+def test_to_dict_with_events_buffer(ctx):
+    ctx.collect_events(OneTestEvent(), [OneTestEvent, AnotherTestEvent])
+    assert json.dumps(ctx.to_dict())
+
+
 @pytest.mark.asyncio()
 async def test_deprecated_params(ctx):
     with pytest.warns(
diff --git a/llama-index-core/tests/workflow/test_workflow.py b/llama-index-core/tests/workflow/test_workflow.py
index 32cde7168..91d6375c8 100644
--- a/llama-index-core/tests/workflow/test_workflow.py
+++ b/llama-index-core/tests/workflow/test_workflow.py
@@ -220,6 +220,9 @@ async def test_workflow_num_workers():
             ctx.session.send_event(OneTestEvent(test_param="test2"))
             ctx.session.send_event(OneTestEvent(test_param="test3"))
 
+            # send one extra event
+            ctx.session.send_event(AnotherTestEvent(another_test_param="test4"))
+
             return LastEvent()
 
         @step(num_workers=3)
@@ -241,11 +244,21 @@ async def test_workflow_num_workers():
     workflow = NumWorkersWorkflow()
 
     start_time = time.time()
-    result = await workflow.run()
+    handler = workflow.run()
+    result = await handler
     end_time = time.time()
 
     assert workflow.is_done()
-    assert set(result) == {"test1", "test2", "test3"}
+    assert set(result) == {"test1", "test2", "test4"}
+
+    # ctx should have 1 extra event
+    assert (
+        len(handler.ctx._events_buffer["tests.workflow.conftest.AnotherTestEvent"]) == 1
+    )
+
+    # ensure ctx is serializable
+    ctx = handler.ctx
+    ctx.to_dict()
 
     # Check if the execution time is close to 1 second (with some tolerance)
     execution_time = end_time - start_time