in ai-ml/spark-gemini-rag/main.py [0:0]
def index():
async def update_prompt():
print("prompt received")
input_time = time.time()
user_input = user_input_raw.value
with chat_container:
ui.chat_message(user_input, name='Me')
print("spinner")
spinner = ui.spinner('audio', size='lg', color='green')
client_id = str(uuid.uuid4())
request_id = f"{client_id}-{str(uuid.uuid4())[:8]}"
prompt, prompt_version = await run.cpu_bound(prompt_maker, user_input)
app.storage.client["count"] = app.storage.client.get("count", 0) + 1
app.storage.client["history"] = app.storage.client.get("history", "") + "### User: " + prompt
with TRACER.start_as_current_span("child") as span:
span.set_attribute(
"operation.count", app.storage.client["count"])
span.set_attribute("prompt", user_input)
span.set_attribute("prompt_id", prompt_version)
span.set_attribute("client_id", client_id)
span.set_attribute("request_id", request_id)
request_time = time.time()
response = await run.io_bound(make_gemini_prediction, prompt)
# response = make_prediction(user_input)
response_time = time.time()
app.storage.client["history"] = app.storage.client.get("history") + "### Agent: " + response
span.set_attribute("response", response)
spinner.delete()
ui.chat_message(response,
name='Robot',
stamp='now',
avatar='https://robohash.org/ui',) \
.style('font-family: Comic Sans, sans-serif; font-size: 16px;')
query = {
"request_id": request_id,
"prompt": user_input,
"response": response,
"input_time": input_time,
"request_time": request_time,
"response_time": response_time,
"prompt_version": prompt_version
}
print(f"Count: {app.storage.client['count']}")
write_to_database(client_id, query)
# print(multiturn_quality(
# app.storage.client.get("history"),
# prompt,
# response
# ))
ui.markdown("<h2>Welcome to predictions bot!</h2>")
with ui.row().classes('flex flex-col h-screen'):
chat_container = ui.column().classes('w-full max-w-3xl mx-auto my-6')
with ui.footer().classes('bg-black'), ui.column().classes('w-full max-w-3xl mx-auto my-6'):
with ui.row().classes('w-full no-wrap items-center'):
user_input_raw = ui.input("Prompt").on('keydown.enter', update_prompt) \
.props('rounded outlined input-class=mx-3').classes('flex-grow')