Update app.py

app.py

@@ -1,136 +1,71 @@
 import os
-import textwrap
-
-#from llama_index.query_engine.retriever_query_engine import RetrieverQueryEngine
-from llama_index.core.query_engine import RetrieverQueryEngine
-
-
-#from llama_index import VectorStoreIndex, SimpleDirectoryReader, Document
-# new (v0.10+)
-from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
-from llama_index.core.schema import Document  # safest place for Document
-
-
-from llama_index.vector_stores import DeepLakeVectorStore
-from langchain.chat_models import ChatOpenAI
 import chainlit as cl
-
-
-import os
-os.environ["OPENAI_API_KEY"]=  os.environ.get("open_ai")
-
-
-
-from llama_index import SimpleDirectoryReader, Document, StorageContext, OpenAIEmbedding, ServiceContext, PromptHelper, VectorStoreIndex
-from llama_index.vector_stores import PineconeVectorStore, QdrantVectorStore, SimpleVectorStore, DeepLakeVectorStore
-from transformers import BertTokenizerFast
-import openai
-from llama_index.llms import OpenAI
-from llama_index import ServiceContext
-from llama_index.embeddings import OpenAIEmbedding
-from llama_index.node_parser import SimpleNodeParser
-from llama_index.text_splitter import TokenTextSplitter
-
-from llama_index import StorageContext, load_index_from_storage
-from llama_index import load_index_from_storage, load_indices_from_storage, load_graph_from_storage
-
-
-
-#dataset_path ="hub://cxcxxaaaaaz/text_embedding" # if we comment this out and don't pass the path then GPTDeepLakeIndex will create dataset in memory
-from llama_index.storage.storage_context import StorageContext
-from llama_index import VectorStoreIndex, SimpleDirectoryReader, Document
-from llama_index.vector_stores import ChromaVectorStore
-
-# Create an index over the documnts
-#vector_store = DeepLakeVectorStore(dataset_path=dataset_path
-
-
-
-
-
-
-
 import chromadb
 
-db2 = chromadb.PersistentClient(path="./chroma_db")
-chroma_collection = db2.get_or_create_collection("vector")
-
-
+# --- LlamaIndex (v0.10+) imports ---
+from llama_index.core import VectorStoreIndex, Settings, StorageContext
+from llama_index.llms.openai import OpenAI
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.vector_stores.chroma import ChromaVectorStore
+
+# --------------------
+# OpenAI credentials
+# --------------------
+# Prefer OPENAI_API_KEY; fall back to `open_ai` if that’s how your Space stores it.
+_openai_key = os.getenv("OPENAI_API_KEY") or os.getenv("open_ai")
+if _openai_key:
+    os.environ["OPENAI_API_KEY"] = _openai_key
+else:
+    raise RuntimeError(
+        "Missing OPENAI_API_KEY (or 'open_ai'). "
+        "Set it in your Space secrets."
+    )
 
+# --------------------
+# LlamaIndex global settings (recommended in v0.10+)
+# --------------------
+Settings.llm = OpenAI(model="gpt-4o-mini", temperature=0.1)
+Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
+
+# --------------------
+# Vector store: Chroma (persistent)
+# --------------------
+chroma_client = chromadb.PersistentClient(path="./chroma_db")
+chroma_collection = chroma_client.get_or_create_collection("vector")
 vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
-storage_context = StorageContext.from_defaults(vector_store=vector_store)
-
-
-
-
-
-#vector_store = LanceDBVectorStore1(uri="./sample_data/")
-
-
-#storage_context = StorageContext.from_defaults(vector_store=vector_store)              
-
-
-llm = OpenAI(model='gpt-3.5-turbo', temperature=0.1)
-embed_model = OpenAIEmbedding()
-
-#node_parser = SimpleNodeParser(text_splitter=TokenTextSplitter(chunk_size=2924, chunk_overlap=20))
-
-prompt_helper = PromptHelper(
-  context_window=2000,
-  num_output=256,
-  chunk_overlap_ratio=0.1,
-  chunk_size_limit=200
-)
-
-import tiktoken
-from llama_index.callbacks import CallbackManager, TokenCountingHandler
-from llama_index import load_index_from_storage, load_indices_from_storage, load_graph_from_storage
-
-token_counter = TokenCountingHandler(tokenizer=tiktoken.encoding_for_model("gpt-3.5-turbo").encode)
-
-callback_manager = CallbackManager([token_counter])
 
+storage_context = StorageContext.from_defaults(vector_store=vector_store)
 
-service_context = ServiceContext.from_defaults(
-  llm=llm,
-  embed_model=embed_model
+# Create an index “shell” over the vector store.
+# (Use `index.insert_nodes(...)` or build from docs later when you ingest.)
+index = VectorStoreIndex.from_vector_store(
+    vector_store=vector_store,
+    storage_context=storage_context,
 )
-from llama_index import set_global_service_context
-
-
-
-index = VectorStoreIndex.from_documents([], vectorstore=vector_store, storage_context=storage_context, service_context=service_context)
-
-
 
+# --------------------
+# Chainlit app
+# --------------------
 @cl.on_chat_start
-async def factory():
-
-        
-    
-# Substitute your connection string here
-   
-    query_engine = index.as_query_engine(
-        service_context=service_context,
-        streaming=True,
-    )
-
-    
+async def on_chat_start():
+    # Build a streaming query engine
+    query_engine = index.as_query_engine(streaming=True)
     cl.user_session.set("query_engine", query_engine)
-
-
+    await cl.Message(content="Hi! Ask me anything.").send()
 
 @cl.on_message
-async def main(message: cl.Message):
-    query_engine = cl.user_session.get("query_engine")  # type: RetrieverQueryEngine
+async def on_message(message: cl.Message):
+    query_engine = cl.user_session.get("query_engine")
+    # Run query in a thread pool (Chainlit helper)
     response = await cl.make_async(query_engine.query)(message.content)
 
-    response_message = cl.Message(content="")
-
+    # Stream tokens as they arrive
+    reply = cl.Message(content="")
     for token in response.response_gen:
-        await response_message.stream_token(token=token)
+        await reply.stream_token(token)
 
-    if response.response_txt:
-        response_message.content = response.response_txt
+    # Ensure final text is set (in case engine also returns response_txt)
+    if getattr(response, "response_txt", None):
+        reply.content = response.response_txt
 
-    await response_message.send()
+    await reply.send()