Update app.py

app.py

@@ -1,165 +1,81 @@
-from llama_index import ServiceContext, SimpleNodeParser, TokenTextSplitter, OpenAI, OpenAIEmbedding
-from llama_index.vector_stores import ChromaVectorStore
-from llama_index.storage.storage_context import StorageContext
-from llama_index import VectorStoreIndex, WikipediaReader
-from llama_index.tools import FunctionTool
-from llama_index.vector_stores.types import VectorStoreInfo, MetadataInfo, ExactMatchFilter, MetadataFilters
-from llama_index.retrievers import VectorIndexRetriever
-from llama_index.query_engine import RetrieverQueryEngine
-from typing import List, Tuple, Any
-from pydantic import BaseModel, Field
-import chromadb
-import pandas as pd
-from sqlalchemy import create_engine
-from llama_index import SQLDatabase, NLSQLTableQueryEngine, QueryEngineTool
-from llama_index.openai_agent import OpenAIAgent
-from chainlit import ChainLit
-
-# Embedding Model and Low-level model
-embed_model = OpenAIEmbedding()
-chunk_size = 1000
-chunk_overlap = 100
-llm = OpenAI(
-    temperature=0,
-    model="gpt-4-32k",
-    streaming=True
-)
-service_context = ServiceContext.from_defaults(
-    llm=llm,
-    chunk_size=chunk_size,
-    embed_model=embed_model
+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,
 )
-text_splitter = TokenTextSplitter(
-    chunk_size=chunk_size,
-    chunk_overlap=chunk_overlap
-)
-node_parser = SimpleNodeParser(text_splitter=text_splitter)
-
-# Vector Storage and Context
-chroma_client = chromadb.Client()
-chroma_collection = chroma_client.create_collection("wikipedia_barbie_opp")
-vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
-storage_context = StorageContext.from_defaults(vector_store=vector_store)
-
-# Your Wikipedia docs retrieval
-movie_list = ["Barbie (film)", "Oppenheimer (film)"]
-wiki_docs = WikipediaReader().load_data(pages=movie_list, auto_suggest=False)
-
-# Parsing and storing vectors
-wiki_vector_index = VectorStoreIndex([], storage_context=storage_context, service_context=service_context)
-for movie, wiki_doc in zip(movie_list, wiki_docs):
-    nodes = node_parser.get_nodes_from_documents([wiki_doc])
-    for node in nodes:
-        node.metadata = {"title": movie}
-    wiki_vector_index.insert_nodes(nodes)
-
-# Defining the tools for vector search and SQL query
-top_k = 3
-vector_store_info = VectorStoreInfo(
-    content_info="semantic information about movies",
-    metadata_info=[MetadataInfo(
-        name="title",
-        type="str",
-        description="title of the movie, one of [Barbie (film), Oppenheimer (film)]",
-    )]
-)
-
-# Create PyDantic model for auto retrieval
-class AutoRetrieveModel(BaseModel):
-    query: str = Field(..., description="natural language query string")
-    filter_key_list: List[str] = Field(
-        ..., description="List of metadata filter field names"
-    )
-    filter_value_list: List[str] = Field(
-        ...,
-        description=(
-            "List of metadata filter field values (corresponding to names specified in filter_key_list)"
-        )
+from langchain.chat_models import ChatOpenAI
+import chainlit as cl
+
+# Set up logging for debugging and monitoring
+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.")
+
+@cl.on_chat_start
+async def factory():
+    embed_model = OpenAIEmbedding()
+    chunk_size = 1000
+
+    llm_predictor = LLMPredictor(
+        llm=ChatOpenAI(
+            temperature=0,
+            model_name="gpt-4-32k",
+            streaming=True,
+        ),
     )
-
-def auto_retrieve_fn(query: str, filter_key_list: List[str], filter_value_list: List[str]):
-    exact_match_filters = [
-        ExactMatchFilter(key=k, value=v)
-        for k, v in zip(filter_key_list, filter_value_list)
-    ]
-    retriever = VectorIndexRetriever(
-        wiki_vector_index, filters=MetadataFilters(filters=exact_match_filters), top_k=top_k
+    service_context = ServiceContext.from_defaults(
+        llm_predictor=llm_predictor,
+        chunk_size=chunk_size,
+        callback_manager=CallbackManager([cl.LlamaIndexCallbackHandler()]),
     )
-    query_engine = RetrieverQueryEngine.from_args(retriever)
-    response = query_engine.query(query)
-    return str(response)
-
-description = f"""\
-Use this tool to look up semantic information about films.
-The vector database schema is given below:
-{vector_store_info.json()}
-"""
-
-auto_retrieve_tool = FunctionTool.from_defaults(
-    fn=auto_retrieve_fn,
-    name="Auto_Retriever",
-    description=description,
-    fn_schema=AutoRetrieveModel
-)
-
-# SQL setup and tool definition
-barbie_df = pd.read_csv("barbie_data/barbie.csv")
-oppenheimer_df = pd.read_csv("oppenheimer_data/oppenheimer.csv")
-engine = create_engine("sqlite+pysqlite:///:memory:")
-barbie_df.to_sql(name='barbie', con=engine)
-oppenheimer_df.to_sql(name='oppenheimer', con=engine)
-
-sql_database = SQLDatabase(
-    engine=engine,
-    include_tables=['barbie', 'oppenheimer']
-)
 
-sql_query_engine = NLSQLTableQueryEngine(
-    sql_database=sql_database,
-    tables=['barbie', 'oppenheimer']
-)
-
-sql_tool = QueryEngineTool.from_defaults(
-    query_engine=sql_query_engine,
-    name="Natural_Language_to_SQL_Tool",
-    description=(
-        "Useful for translating a natural language query into a SQL query."
+    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) 
 
-# Combining both tools into a single OpenAI Agent
-barbenheimer_agent = OpenAIAgent.from_tools(
-    tools=[auto_retrieve_tool, sql_tool]
-)
-
-# Initialize the ChainLit app
-cl = ChainLit()
-
-# On-Message Function
 @cl.on_message
-def handle_message(message: str, sender: str) -> Tuple[str, Any]:
-    query_result = barbenheimer_agent.query(
-        query=message,
-        user_id=sender
-    )
-    
-    # Extracting relevant information from the query result
-    tool_name = query_result.tool_name
-    response = query_result.response
-    
-    if tool_name == "Auto_Retriever":
-        # Processing for semantic information retrieval
-        return "Auto_Retriever", f"Semantic Information:\n{response}"
+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="")
         
-    elif tool_name == "Natural_Language_to_SQL_Tool":
-        # Processing for SQL-based information
-        return "Natural_Language_to_SQL_Tool", f"SQL Query Result:\n{response}"
-    
-    else:
-        # Handling unrecognized tool queries
-        return "Unknown", "I couldn't understand your request."
+        for token in response.response_gen:
+            await response_message.stream_token(token=token)
+
+        if response.response_txt:
+            response_message.content = response.response_txt
+
+        await response_message.send()
+        logger.info(f"Response sent: {response.response_txt}")
 
-# Running the app
-if __name__ == '__main__':
-    cl.run()
+    except Exception as e:
+        logger.error(f"An error occurred while processing the message: {e}")