Migrate files

app.py

@@ -1,74 +1,118 @@
-# You can find this code for Chainlit python streaming here (https://docs.chainlit.io/concepts/streaming/python)
-
-# OpenAI Chat completion
-
-import openai #importing openai for API usage
-import chainlit as cl #importing chainlit for our app
-from chainlit.input_widget import Select, Switch, Slider #importing chainlit settings selection tools
-from chainlit.prompt import Prompt, PromptMessage #importing prompt tools 
-from chainlit.playground.providers import ChatOpenAI #importing ChatOpenAI tools
-
-# You only need the api key inserted here if it's not in your .env file
-#openai.api_key = "YOUR_API_KEY"
-
-# ChatOpenAI Templates
-system_template = """You are a helpful assistant who always speaks in a pleasant tone!
-"""
-
-user_template = """{input}
-Think through your response step by step.
-"""
-
-@cl.on_chat_start # marks a function that will be executed at the start of a user session
-async def start_chat():
-    settings = {
-        "model": "gpt-3.5-turbo",
-        "temperature": 0,
-        "max_tokens": 500,
-        "top_p": 1,
-        "frequency_penalty": 0,
-        "presence_penalty": 0,
-    }
-
-    cl.user_session.set("settings", settings)
-
-@cl.on_message # marks a function that should be run each time the chatbot receives a message from a user
-async def main(message: str):
-
-    settings = cl.user_session.get("settings")
-
-    prompt = Prompt(
-        provider=ChatOpenAI.id,
-        messages=[
-            PromptMessage(
-                role="system",
-                template=system_template,
-                formatted=system_template,
+import os
+
+import chainlit as cl
+import openai
+from chainlit.input_widget import Select, Slider, Switch
+from langchain.chat_models import ChatOpenAI
+from llama_index import (
+    LLMPredictor,
+    ServiceContext,
+    StorageContext,
+    TrafilaturaWebReader,
+    VectorStoreIndex,
+    load_index_from_storage,
+)
+from llama_index.callbacks.base import CallbackManager
+from llama_index.llms import ChatMessage, HuggingFaceLLM, MessageRole, OpenAI
+
+
+def get_api_key():
+    api_key = os.getenv("OPENAI_API_KEY")
+    if api_key is None:
+        print("OPENAI_API_KEY missing from environment variables")
+        api_key = input("Please enter your OPENAI_API_KEY: ")
+    return api_key
+
+
+openai.api_key = get_api_key()
+
+
+def load_index():
+    try:
+        storage_context = StorageContext.from_defaults(persist_dir="./storage")
+        index = load_index_from_storage(storage_context)
+    except FileNotFoundError:
+        print("Storage file not found. Loading from web.")
+        documents = TrafilaturaWebReader().load_data(["https://bit.ly/45BncJA"])
+        index = VectorStoreIndex.from_documents(documents)
+        index.storage_context.persist()
+    return index
+
+
+index = load_index()
+
+welcome_msg = (
+    "Hi there! I’m your China Life chatbot, specialising in answering "
+    "[frequently asked questions](https://bit.ly/45BncJA). "
+    "How may I assist you today? "
+    "Feel free to ask questions like, "
+    "“Is there any action required after receiving the policy?” or "
+    "“Can I settle using a demand draft?”"
+)
+
+
+@cl.on_chat_start
+async def start():
+    chat_profile = cl.user_session.get("chat_profile")
+    msg = cl.Message(content="")
+    for token in list(welcome_msg):
+        await cl.sleep(0.01)
+        await msg.stream_token(token)
+
+    await msg.send()
+
+    settings = await cl.ChatSettings(
+        [
+            Select(
+                id="Model",
+                label="Model",
+                values=["gpt-3.5-turbo", "gpt-4"],
+                initial_index=1,
             ),
-            PromptMessage(
-                role="user",
-                template=user_template, 
-                formatted=user_template.format(input=message),
-            )
-        ],
-        inputs = {"input" : message},
-        settings=settings
+            Slider(
+                id="Temperature",
+                label="Temperature",
+                initial=0,
+                min=0,
+                max=2,
+                step=0.1,
+            ),
+        ]
+    ).send()
+    await setup_query_engine(settings)
+
+
+@cl.on_settings_update
+async def setup_query_engine(settings):
+    print("on_settings_update", settings)
+
+    llm = OpenAI(model=settings["Model"], temperature=settings["Temperature"])
+
+    service_context = ServiceContext.from_defaults(
+        llm=llm, callback_manager=CallbackManager([cl.LlamaIndexCallbackHandler()])
     )
 
-    print([m.to_openai() for m in prompt.messages])
+    query_engine = index.as_query_engine(
+        service_context=service_context,
+        streaming=True,
+    )
 
-    msg = cl.Message(content="")
+    cl.user_session.set("query_engine", query_engine)
 
-    # Call OpenAI
-    async for stream_resp in await openai.ChatCompletion.acreate(
-        messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
-    ):
-        token = stream_resp.choices[0]["delta"].get("content", "")
-        await msg.stream_token(token)
 
-    # Update the prompt object with the completion
-    prompt.completion = msg.content
-    msg.prompt = prompt
+@cl.on_message
+async def main(message: cl.Message):
+    query_engine = cl.user_session.get("query_engine")
 
-    # Send and close the message stream
-    await msg.send()
+    if query_engine is None:
+        await start()
+        query_engine = cl.user_session.get("query_engine")
+
+    if query_engine:
+        query_result = await cl.make_async(query_engine.query)(message.content)
+        response_message = cl.Message(content=query_result.response_txt or "")
+
+        for token in query_result.response_gen:
+            await response_message.stream_token(token=token)
+
+        await response_message.send()