Spaces:
No application file
No application file
File size: 2,620 Bytes
ad73d17 |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import os
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_community.callbacks import StreamlitCallbackHandler
import streamlit as st
import requests
import time
model_service = os.getenv("MODEL_ENDPOINT", "http://localhost:8001")
model_service = f"{model_service}/v1"
model_service_bearer = os.getenv("MODEL_ENDPOINT_BEARER")
request_kwargs = {}
if model_service_bearer is not None:
request_kwargs = {"headers": {"Authorization": f"Bearer {model_service_bearer}"}}
@st.cache_resource(show_spinner=False)
def checking_model_service():
start = time.time()
print("Checking Model Service Availability...")
ready = False
while not ready:
try:
request = requests.get(f'{model_service}/models', **request_kwargs)
if request.status_code == 200:
ready = True
except:
pass
time.sleep(1)
print("Model Service Available")
print(f"{time.time()-start} seconds")
with st.spinner("Checking Model Service Availability..."):
checking_model_service()
st.title("Code Generation App")
if "messages" not in st.session_state:
st.session_state["messages"] = [{"role": "assistant",
"content": "How can I help you?"}]
for msg in st.session_state.messages:
st.chat_message(msg["role"]).write(msg["content"])
model_name = os.getenv("MODEL_NAME", "")
llm = ChatOpenAI(base_url=model_service,
model=model_name,
api_key="EMPTY" if model_service_bearer is None else model_service_bearer,
streaming=True)
# Define the Langchain chain
prompt = ChatPromptTemplate.from_template("""You are an helpful code assistant that can help developer to code for a given {input}.
Generate the code block at first, and explain the code at the end.
If the {input} is not making sense, please ask for more clarification.""")
chain = (
{"input": RunnablePassthrough()}
| prompt
| llm
)
if prompt := st.chat_input():
st.session_state.messages.append({"role": "user", "content": prompt})
st.chat_message("user").markdown(prompt)
st_callback = StreamlitCallbackHandler(st.container())
response = chain.invoke(prompt, {"callbacks": [st_callback]})
st.chat_message("assistant").markdown(response.content)
st.session_state.messages.append({"role": "assistant", "content": response.content})
st.rerun()
|