Skip to content

Commit

Permalink
fix chat and knowledge_base_chat
Browse files Browse the repository at this point in the history
  • Loading branch information
liunux4odoo authored and hzg0601 committed Aug 14, 2023
1 parent 90c92f8 commit 0704eab
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 5 deletions.
14 changes: 11 additions & 3 deletions server/chat/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ def chat(query: str = Body(..., description="用户输入", examples=["恼羞成
{"role": "user", "content": "我们来玩成语接龙,我先来,生龙活虎"},
{"role": "assistant", "content": "虎头虎脑"}]]
),
stream: bool = Body(False, description="流式输出"),
):
history = [History(**h) if isinstance(h, dict) else h for h in history]

Expand Down Expand Up @@ -46,9 +47,16 @@ async def chat_iterator(query: str,
callback.done),
)

async for token in callback.aiter():
# Use server-sent-events to stream the response
yield token
if stream:
async for token in callback.aiter():
# Use server-sent-events to stream the response
yield token
else:
answer = ""
async for token in callback.aiter():
answer += token
yield answer

await task

return StreamingResponse(chat_iterator(query, history),
Expand Down
5 changes: 3 additions & 2 deletions server/chat/knowledge_base_chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from server.chat.utils import History
from server.knowledge_base.kb_service.base import KBService, KBServiceFactory
import json
import os


def knowledge_base_chat(query: str = Body(..., description="用户输入", examples=["你好"]),
Expand Down Expand Up @@ -64,7 +65,7 @@ async def knowledge_base_chat_iterator(query: str,
)

source_documents = [
f"""出处 [{inum + 1}] [{doc.metadata["source"]}]({doc.metadata["source"]}) \n\n{doc.page_content}\n\n"""
f"""出处 [{inum + 1}] [{os.path.split(doc.metadata["source"])[-1]}] \n\n{doc.page_content}\n\n"""
for inum, doc in enumerate(docs)
]

Expand All @@ -78,7 +79,7 @@ async def knowledge_base_chat_iterator(query: str,
answer = ""
async for token in callback.aiter():
answer += token
yield json.dumps({"answer": token,
yield json.dumps({"answer": answer,
"docs": source_documents},
ensure_ascii=False)

Expand Down

0 comments on commit 0704eab

Please sign in to comment.