File size: 1,088 Bytes
3723d1a
 
 
 
 
 
 
 
729f0f9
3723d1a
 
 
 
 
 
307474b
3723d1a
 
 
 
 
 
 
 
 
 
 
307474b
 
 
 
 
3e85355
307474b
 
87ff9a1
 
 
 
3723d1a
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
from langchain import PromptTemplate, LLMChain
import chainlit as cl
from custom_llm import CustomLLM
from langchain.prompts import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
)

template = """ Write a python code for the following problem :
{question} 

Code:
"""


@cl.on_chat_start
def factory():
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)

    prompt = ChatPromptTemplate.from_messages([system_message_prompt])
    llm = CustomLLM()

    llm_chain = LLMChain(prompt=prompt, llm=llm, verbose=True,)

    cl.user_session.set("llm_chain", llm_chain)



@cl.on_message
async def main(message):
    llm_chain = cl.user_session.get("llm_chain")

    res = await llm_chain.acall(message.content, callbacks=[cl.AsyncLangchainCallbackHandler()])

    await cl.Message(content=res["text"]).send()



@cl.author_rename   # This will be particularly useful when we want to customize this thing for production.
def rename(orig_author):
    rename_dict = {
        'LLMChain': 'Scooby'
    }
    return rename_dict.get(orig_author, orig_author)