Spaces:
Runtime error
Runtime error
import os | |
import openai | |
import logging | |
from llama_index.query_engine.retriever_query_engine import RetrieverQueryEngine | |
from llama_index.callbacks.base import CallbackManager | |
from llama_index import ( | |
LLMPredictor, | |
ServiceContext, | |
StorageContext, | |
load_index_from_storage, | |
) | |
from langchain.chat_models import ChatOpenAI | |
import chainlit as cl | |
# Set up logging for debugging and monitoring of errors | |
logging.basicConfig(level=logging.INFO) | |
logger = logging.getLogger(__name__) | |
# Load OpenAI API key | |
openai.api_key = os.environ.get("OPENAI_API_KEY") | |
try: | |
# Attempt to rebuild storage context and load index | |
logger.info("Attempting to load index from storage.") | |
storage_context = StorageContext.from_defaults(persist_dir="./storage") | |
index = load_index_from_storage(storage_context) | |
except Exception as e: | |
# If index loading fails, create a new index | |
logger.warning(f"Failed to load index from storage: {e}. Creating a new index.") | |
from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader | |
documents = SimpleDirectoryReader("./data").load_data() | |
index = GPTVectorStoreIndex.from_documents(documents) | |
index.storage_context.persist() | |
logger.info("New index created and persisted.") | |
async def factory(): | |
#embed_model = OpenAIEmbedding() | |
chunk_size = 1000 | |
llm_predictor = LLMPredictor( | |
llm=ChatOpenAI( | |
temperature=0, | |
model_name="gpt-4", | |
streaming=True, | |
), | |
) | |
service_context = ServiceContext.from_defaults( | |
llm_predictor=llm_predictor, | |
chunk_size=chunk_size, | |
callback_manager=CallbackManager([cl.LlamaIndexCallbackHandler()]), | |
) | |
query_engine = index.as_query_engine( | |
service_context=service_context, | |
streaming=True, | |
) | |
logger.info("Query engine initialized.") # to facilitate debugging and monitoring | |
cl.user_session.set("query_engine", query_engine) | |
async def main(message): | |
try: | |
query_engine = cl.user_session.get("query_engine") # type: RetrieverQueryEngine | |
logger.info(f"Received message: {message}") | |
response = await cl.make_async(query_engine.query)(message) | |
response_message = cl.Message(content="") | |
# Logic to prepare answer and source_elements | |
for token in response.response_gen: | |
await response_message.stream_token(token=token) | |
if response.response_txt: | |
response_message.content = response.response_txt | |
# Integrated new message object | |
if answer: # conditional to when is not None | |
await cl.Message(content=answer, elements=source_elements).send() | |
await response_message.send() | |
logger.info(f"Response sent: {response.response_txt}") | |
except Exception as e: | |
logger.error(f"An error occurred while processing the message: {e}") | |