fix:reference bug

This commit is contained in:
aries_ckt 2023-11-13 23:43:52 +08:00
parent d7f677a59a
commit 6ea9c17aae
2 changed files with 5 additions and 1 deletions

View File

@ -174,6 +174,9 @@ class BaseChat(ABC):
def stream_plugin_call(self, text):
return text
def stream_call_reinforce_fn(self, text):
return text
async def check_iterator_end(iterator):
try:
await asyncio.anext(iterator)
@ -215,6 +218,7 @@ class BaseChat(ABC):
view_msg = view_msg.replace("\n", "\\n")
yield view_msg
self.current_message.add_ai_message(msg)
view_msg = self.stream_call_reinforce_fn(view_msg)
self.current_message.add_view_message(view_msg)
span.end()
except Exception as e:

View File

@ -89,7 +89,7 @@ class ChatKnowledge(BaseChat):
last_output = last_output + reference
yield last_output
def stream_plugin_call(self, text):
def stream_call_reinforce_fn(self, text):
"""return reference"""
return text + f"\n\n{self.parse_source_view(self.sources)}"