Spaces:
Sleeping
Sleeping
File size: 2,122 Bytes
9002555 0743bb0 9002555 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
# Experimental
from typing import Dict, Any, Optional, List
import asyncio
import logging
from uuid import uuid4
from anyio import ClosedResourceError
from anyio.streams.memory import MemoryObjectSendStream
from llama_index.core.callbacks.base import BaseCallbackHandler, CallbackManager
from llama_index.core.callbacks import CBEventType, EventPayload
from llama_index.core.query_engine.sub_question_query_engine import (
SubQuestionAnswerPair,
)
from llama_index.core.chat_engine.types import StreamingAgentChatResponse
from pydantic import BaseModel
from core.chat import schema
from db.db import MessageSubProcessSourceEnum
from core.chat.schema import SubProcessMetadataKeysEnum, SubProcessMetadataMap
from core.chat.engine import Engine
from script.vector_db import IndexManager
from service.dto import UserPromptRequest
logger = logging.getLogger(__name__)
class StreamedMessage(BaseModel):
content: str
async def handle_chat_message(
user_message: str,
send_chan: MemoryObjectSendStream,
) -> None:
async with send_chan:
engine = Engine()
index_manager = IndexManager()
index = index_manager.load_existing_indexes()
# Retrieve the chat engine with the loaded index
chat_engine = await engine.get_chat_engine(index)
logger.debug("Engine received")
streaming_chat_response: StreamingAgentChatResponse = (
await chat_engine.astream_chat(user_message)
)
response_str = ""
async for text in streaming_chat_response.async_response_gen():
response_str += text
if send_chan._closed:
logger.debug(
"Received streamed token after send channel closed. Ignoring."
)
return
await send_chan.send(StreamedMessage(content=response_str))
if response_str.strip() == "":
await send_chan.send(
StreamedMessage(
content="Sorry, I either wasn't able to understand your question or I don't have an answer for it."
)
)
|