From cf5f75a228ecd8e8cd423f0ada61eca2d97fe96f Mon Sep 17 00:00:00 2001
From: Andrei Fajardo <92402603+nerdai@users.noreply.github.com>
Date: Fri, 15 Mar 2024 21:19:52 -0400
Subject: [PATCH] [FIX] - PydanticResponse Bug (#11986)
---
.../llama_index/core/base/response/schema.py | 10 ++++++
.../core/instrumentation/events/synthesis.py | 3 +-
.../instrumentation/span_handlers/simple.py | 33 +++++++++++++++----
.../core/response_synthesizers/base.py | 12 +++----
4 files changed, 44 insertions(+), 14 deletions(-)
diff --git a/llama-index-core/llama_index/core/base/response/schema.py b/llama-index-core/llama_index/core/base/response/schema.py
index f3239338bb..3689ed2cef 100644
--- a/llama-index-core/llama_index/core/base/response/schema.py
+++ b/llama-index-core/llama_index/core/base/response/schema.py
@@ -159,6 +159,16 @@ class AsyncStreamingResponse:
def __post_init__(self) -> None:
self._lock = asyncio.Lock()
+ def __str__(self) -> str:
+ """Convert to string representation."""
+ return asyncio.run(self._async_str)
+
+ async def _async_str(self) -> str:
+ """Convert to string representation."""
+ async for _ in self._yield_response():
+ ...
+ return self.response_txt or "None"
+
async def _yield_response(self) -> TokenAsyncGen:
"""Yield the string response."""
async with self._lock:
diff --git a/llama-index-core/llama_index/core/instrumentation/events/synthesis.py b/llama-index-core/llama_index/core/instrumentation/events/synthesis.py
index 71251eaec6..27a3a03dd2 100644
--- a/llama-index-core/llama_index/core/instrumentation/events/synthesis.py
+++ b/llama-index-core/llama_index/core/instrumentation/events/synthesis.py
@@ -1,4 +1,3 @@
-from llama_index.core.base.response.schema import RESPONSE_TYPE
from llama_index.core.instrumentation.events.base import BaseEvent
from llama_index.core.schema import QueryType
@@ -14,7 +13,7 @@ class SynthesizeStartEvent(BaseEvent):
class SynthesizeEndEvent(BaseEvent):
query: QueryType
- response: RESPONSE_TYPE
+ response: str
@classmethod
def class_name(cls):
diff --git a/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py b/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py
index 4cdee9d67d..1cea445b77 100644
--- a/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py
+++ b/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py
@@ -3,6 +3,7 @@ from llama_index.core.bridge.pydantic import Field
from llama_index.core.instrumentation.span.simple import SimpleSpan
from llama_index.core.instrumentation.span_handlers.base import BaseSpanHandler
from datetime import datetime
+import warnings
if TYPE_CHECKING:
from treelib import Tree
@@ -42,6 +43,7 @@ class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]):
"""Method for getting trace trees."""
try:
from treelib import Tree
+ from treelib.exceptions import NodeIDAbsentError
except ImportError as e:
raise ImportError(
"`treelib` package is missing. Please install it by using "
@@ -59,12 +61,31 @@ class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]):
# start new tree
tree = Tree()
- tree.create_node(
- tag=f"{span.id_} ({span.duration})",
- identifier=span.id_,
- parent=span.parent_id,
- data=span.start_time,
- )
+ try:
+ tree.create_node(
+ tag=f"{span.id_} ({span.duration})",
+ identifier=span.id_,
+ parent=span.parent_id,
+ data=span.start_time,
+ )
+ except NodeIDAbsentError:
+ warnings.warn("Parent with id {span.parent_id} missing from spans")
+ # create new tree and fake parent node
+ trees.append(tree)
+ tree = Tree()
+ tree.create_node(
+ tag=f"{span.parent_id} (MISSING)",
+ identifier=span.parent_id,
+ parent=None,
+ data=span.start_time,
+ )
+ tree.create_node(
+ tag=f"{span.id_} ({span.duration})",
+ identifier=span.id_,
+ parent=span.parent_id,
+ data=span.start_time,
+ )
+
trees.append(tree)
return trees
diff --git a/llama-index-core/llama_index/core/response_synthesizers/base.py b/llama-index-core/llama_index/core/response_synthesizers/base.py
index ebedd93aa6..b844080789 100644
--- a/llama-index-core/llama_index/core/response_synthesizers/base.py
+++ b/llama-index-core/llama_index/core/response_synthesizers/base.py
@@ -209,13 +209,13 @@ class BaseSynthesizer(ChainableMixin, PromptMixin):
response_gen=empty_response_generator()
)
dispatcher.event(
- SynthesizeEndEvent(query=query, response=empty_response)
+ SynthesizeEndEvent(query=query, response=str(empty_response))
)
return empty_response
else:
empty_response = Response("Empty Response")
dispatcher.event(
- SynthesizeEndEvent(query=query, response=empty_response)
+ SynthesizeEndEvent(query=query, response=str(empty_response))
)
return empty_response
@@ -240,7 +240,7 @@ class BaseSynthesizer(ChainableMixin, PromptMixin):
event.on_end(payload={EventPayload.RESPONSE: response})
- dispatcher.event(SynthesizeEndEvent(query=query, response=response))
+ dispatcher.event(SynthesizeEndEvent(query=query, response=str(response)))
return response
@dispatcher.span
@@ -258,13 +258,13 @@ class BaseSynthesizer(ChainableMixin, PromptMixin):
response_gen=empty_response_agenerator()
)
dispatcher.event(
- SynthesizeEndEvent(query=query, response=empty_response)
+ SynthesizeEndEvent(query=query, response=str(empty_response))
)
return empty_response
else:
empty_response = Response("Empty Response")
dispatcher.event(
- SynthesizeEndEvent(query=query, response=empty_response)
+ SynthesizeEndEvent(query=query, response=str(empty_response))
)
return empty_response
@@ -289,7 +289,7 @@ class BaseSynthesizer(ChainableMixin, PromptMixin):
event.on_end(payload={EventPayload.RESPONSE: response})
- dispatcher.event(SynthesizeEndEvent(query=query, response=response))
+ dispatcher.event(SynthesizeEndEvent(query=query, response=str(response)))
return response
def _as_query_component(self, **kwargs: Any) -> QueryComponent:
--
GitLab