def index()

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')