Update app.py

app.py

@@ -1,50 +1,58 @@
-import asyncio
-
-import chainlit as cl
-
-from chain import Chain
 import os
 
-from typing import Dict, Optional
+from langchain.llms.huggingface_hub import HuggingFaceHub
+from langchain.prompts import ChatPromptTemplate
+from langchain.schema import StrOutputParser
+from langchain.schema.runnable import Runnable
+from langchain.schema.runnable.config import RunnableConfig
+
+from chainlit.playground.config import add_llm_provider
+from chainlit.playground.providers.langchain import LangchainGenericProvider
 import chainlit as cl
 
+# Instantiate the LLM
+llm = HuggingFaceHub(
+    model_kwargs={"max_length": 500},
+    repo_id="google/flan-t5-xxl",
+    huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"],
+)
+
+# Add the LLM provider
+add_llm_provider(
+    LangchainGenericProvider(
+        # It is important that the id of the provider matches the _llm_type
+        id=llm._llm_type,
+        # The name is not important. It will be displayed in the UI.
+        name="HuggingFaceHub",
+        # This should always be a Langchain llm instance (correctly configured)
+        llm=llm,
+        # If the LLM works with messages, set this to True
+        is_chat=False,
+    )
+)
 
-@cl.oauth_callback
-def oauth_callback(
-  provider_id: str,
-  token: str,
-  raw_user_data: Dict[str, str],
-  default_user: cl.User,
-) -> Optional[cl.User]:
-  return default_user
 
 @cl.on_chat_start
-async def start_chat():
-    app_user = cl.user_session.get("user")
-    print(app_user)
-    print(os.getenv("OAUTH_CLIENT_ID"))
-    print('HERE-------------------------------------------------')
-    chain = Chain(None)
-    await chain.text("I will count to 5. How many concurrent times should I count?")
+async def on_chat_start():
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("human", "{question}"),
+        ]
+    )
+    runnable = prompt | llm | StrOutputParser()
+    cl.user_session.set("runnable", runnable)
 
 
 @cl.on_message
-async def on_message(message: str, message_id: str):
-    chain = Chain(message_id)
-
-    try:
-        num = int(message)
-    except ValueError:
-        await chain.text_stream("Sorry, that doesn't look like an integer to me.", final=True)
-        return
-
-    if num > 10:
-        await chain.text_stream("Whoa, let's try a smaller number. (Max 10.)", final=True)
-        return
-
-    await chain.text("Alright, here we go:")
-    coroutines = []
-    for i in range(num):
-        coroutines.append(chain.text_stream("1 2 3 4 5", delay=1, name=f"Counter {i + 1}"))
-    await asyncio.gather(*coroutines)
-    await chain.text_stream("Okay, I'm done counting now.", final=True)
+async def on_message(message: cl.Message):
+    runnable = cl.user_session.get("runnable")  # type: Runnable
+
+    msg = cl.Message(content="")
+
+    async for chunk in runnable.astream(
+        {"question": message.content},
+        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
+    ):
+        await msg.stream_token(chunk)
+
+    await msg.send()