On-device Query Router with Chrome's Prompt API

So here I have done a uv init, its an empty project at the moment BUT I want you to you an API for the following:

Its a very simple RAG API, by simple I mean no unnecessary error handlers and such, very simple, I need to use the code for a demo blog thing so no complications as such.

In fact the RAG background should be just copied from this notebook here, its another utter simple RAG here @/home/alphons/project/OAISYS25/chrome_ai_demo/workbench/rag_synth.ipynb its just a reference to show how simple it could be.

for the vector database use something simple as FAISS BUT see that a file or some bin or some dumb is created locally AFTER all the documents are indexed into it because I don't want to re-index everything just because I restart the app, okay?

Now for the embedding model it would be a sentence transformer model but please for gods sake try NOT to download any nvidia libraries and stuff that comes as dependencies, they take GBs and endless times to download.
I guess its something like installing just the cpu versoin of torch first and then later installing sentence transfer and also setting device as cpu
`pip install torch --index-url https://download.pytorch.org/whl/cpu` (adapt to uv command as needed)


as for the model lets try this one: https://huggingface.co/sentence-transformers/embeddinggemma-300m-medical its a medical finetune of embeddinggemma-0.3b, its a new model and its around 1.3/4gb or so, YEs it will be slow on CPU but its okay, we are just creating a demo a>


Next is the chunking method, use a basic paragraph based chunking, no libraries needed as such, lets just do it manually detecting new lines and all.

A quick note on chunking:
When storing to the vector database create ids like this: doc_id_1_chunk_id_1, doc_id_1_chunk_id_2, doc_id_2_chunk_id_1, doc_id_2_chunk_id_2
Hope you got that pattern I'll tell you why we need that in a moment soon as I explain the RAG endpoints.

And as for the data set lets use https://huggingface.co/datasets/zxvix/MedicalTextbook its a data with text column filled with rows of TEXT about basic things. Don;t index complete database, just first 10 rows are enough.

And for inferencing we will use the same OpenRoute model and the way we are accessing it through openai library as seen in the previous notebook that I attached to you.

Now finally we will have 2 very important endpoints as following coming out of this FastAPI implementation:

"/search"
This is basically your vector search sort of thing.
Text in :: chunks out

"/chat"
this is your basic RAG chat thing, but the structure of this will be almost similar to OpenAI endpoints. WE WILL NOT store any history in the API instead we already sent an array of message as we usually do for such chatbots.
messages: [
    {
        "role": "system",
        "content": "some_content"
    },
    {
        "role": "user",
        "content": "tell me about xyz"
    },
    {
        "role": "assistant",
        "content": "xyz means abc and so on"
        "chunks_referred": [
            {
                "id": "doc_id_3_chunk_id_4",
                "content": "long text"
            },
            {
                "id": "doc_id_2_chunk_id_5",
                "content": "long text"
            }
        ]
    },
    {
        "role": "user",
        "content": "Oh thanks, that sounds nice"
    },
]

Now you might think how will the rag work in this case, so we pick the last element of the message array thats sent in, in this case the "oh thanks ..." and then do a quick retrieval on my function powering the search endpoint.

then pass it to the our openrouter api thing. note that since openrouter wont support custom params like "chunks_referred" don't send that into the openrouter api

The response need not be a streaming one for now. Since these are small models we will get the requires response quick so no worries there as well.

And a few more things to keep in mind:
- No need to write data models and stuff for the API, like I said this is a quick learning demo app.
- BUT yeah keep the logic code and stuff out of the routes file, every major action / function should be kept in a separate like how folks structure large projects, even tho this is not large keeping each file line of code to less than 100 lines would be much cleaner and >

I move all the uv init files to @api folder, we will use that folder from now on.

Now ultrathink, and first propose me a plan before coding anything.

Now next we need to create a very simple html + js + css site.
It would be a static site, no complex js frameworks OR anything.
we will create this in the empty @web folder here

Add a simple multi turn chat UI, no need for session managing and services and so on.

Keep the JS part dedicated into its own js file, I mean specifically the inference one and well make it talk with the python api that we just created.

btw keep the UI too very simple, like a multi-turn chat window in the middle, thats pretty much it.

Now we need to implement the Chrome API.
SO right now in this system I don't have access to WebGPU and therefore it wont work here but once you write the code I can test it on my secondary system which I previously already tested the demo sites of this new Chrome API.

So this is a very new API thing that Chrome is pushing for as web standards, so I need you to first do a bit of web search and study about Chrome in-built / on-device AI APIs.

We will be focusing specifically on the Prompt API here for our use case.
And here are some PDF copies of the documentations if you need:
1. @"workbench/developer.chrome.com-Get started with built-in AI.pdf"
2. @"workbench/developer.chrome.com-The Prompt API.pdf"
(Also my origin trials token can be found in ".origin_trail_token" in the workbench folder)

Now I will get down to our use case and idea properly here:
Am planning to use this like a prompt router / gateway sort of thing.

1. So when the user inputs there message in the chatbot (the chatbot can be found in the @web folder, as you can already see we have tried to keep all the code very slim, and simple as this is a demo app)

2. We will have a separate few-shot array of message for the "query_router".
The query_router thing will sort of classify the prompts into following buckets based on there types:

2.1 Follow up queries: "Oh, can you elaborate please" or "Ahh I see, whats the deal with xyz in there tho"
These need no calling to our python API, instead we send the same message array of user/assistants to the Prompt API itself

2.2 Needs more Context AND not complex or too simple to answer: "ok what does abc mean then" or "help me understand abc along with xyz"
So here mentioning any topic or keywords and such in the query mean that we need more context to answer, so we ping the /search endpoint (you can refer that in @api/routes.py) and then similar to how the rag works in @api, we simulate the behavior in the Prompt API itself.

2.3 Complex query: "Can you help me understand the complications of pqr in conjunction with xyz'
Just avoid using Prompt API in this case and route the query to the python endpoint.

---

Few things to keep in mind:

1. the docs suggest that we should implement error handlers and fallbacks and so on, but my whole objective here is to just implement this and use it on a supported platform directly
2. if it doest work THEN it doesn't work thats it, no fallbacks, no unnecessary try catch and stuff and so on.
3. This is a demo application so I just need it working ONE SINGLE TIME just to take a video explaining it and so on, so that Later I can blog post this whole experiment / demo thing.
4. So yeah all of this means that when writing code, just assume things need to work so no over-complications or unnecessary abstractions for future and so on, this will be pretty much the last version of this demo.

---
Alright now instead since I cannot do like a iterative development and test at the moment I need you to ultrathink, take time to think through and plan the whole thing before making any code changes, you will need to propose me the complete implementation plan after you think through the complete thing. And then i'll decide /modify and so as needed.

Component Updater → Optimization Guide → Model Executor → ChromeML Native Library
       ↓                    ↓                  ↓                    ↓
   Download            Manage Files      Create Model        Parse TFLite
   Model Data          & Metadata        Instance            & Initialize

// Chrome's Component Updater downloads:
// - Base model weights (weights.bin)
// - Sentencepiece model (for tokenization)
// - Cache files (optional, for acceleration)
// - Manifest with version info

void OnDeviceModelServiceController::Init() {
  // Register for base model availability
  base_model_asset_manager_ =
      std::make_unique<OnDeviceAssetManager>(model_provider_);

  // Wait for component to be ready
  base_model_asset_manager_->AddObserver(this);
}

void OnDeviceAssetManager::OnModelUpdated(
    proto::OptimizationTarget optimization_target,
    base::optional_ref<const ModelInfo> model_info) {

  if (!model_info.has_value()) {
    return;  // Model not available
  }

  // Extract model paths
  auto weights_path = model_info->GetModelFilePath();
  auto sp_model_path = model_info->GetAdditionalFileWithBaseName(
      kOnDeviceModelSentencePieceFile);

  // Create ModelAssetPaths
  ModelAssetPaths paths;
  paths.weights = weights_path;
  paths.sp_model = sp_model_path;
  paths.cache = GetCacheFilePath();
  paths.encoder_cache = GetEncoderCacheFilePath();
  paths.adapter_cache = GetAdapterCacheFilePath();

  // Notify controller
  NotifyAssetsAvailable(paths);
}

std::unique_ptr<OnDeviceModelExecutor> OnDeviceModelExecutor::Create(
    const ModelAssetPaths& paths) {

  // Load ChromeML native library
  ChromeML* chrome_ml = ChromeML::Get();
  if (!chrome_ml) {
    return nullptr;  // Library not available
  }

  // Prepare model descriptor
  ChromeMLModelDescriptor descriptor = {
    .backend_type = GetBackendType(),  // GPU or APU
    .model_data = &model_data,
    .max_tokens = kDefaultMaxTokens,
    .temperature = kDefaultTemperature,
    .top_k = kDefaultTopK,
    .adaptation_ranks = kSupportedLoRARanks,  // ⭐ [4,8,16,32,64]
    .adaptation_ranks_size = kSupportedLoRARanksSize,
    .prefer_texture_weights = true,
    .enable_host_mapped_pointer = true,
    // ... other config
  };

  // Prepare model data
  ChromeMLModelData model_data;
  model_data.weights_file = OpenModelFile(paths.weights);
  model_data.cache_file = OpenCacheFile(paths.cache);
  model_data.encoder_cache_file = OpenCacheFile(paths.encoder_cache);
  model_data.adapter_cache_file = OpenCacheFile(paths.adapter_cache);

  // ⭐ CREATE MODEL via native library
  ChromeMLModel model = chrome_ml->api().CreateModel(&descriptor);

  if (!model) {
    return nullptr;  // Model creation failed
  }

  return std::make_unique<OnDeviceModelExecutor>(chrome_ml, model);
}

// Native library function signature
typedef ChromeMLModel (*CreateModelFn)(
    const ChromeMLModelDescriptor* descriptor);

struct ChromeMLAPI {
  CreateModelFn CreateModel;
  DestroyModelFn DestroyModel;
  CreateSessionFn CreateSession;
  // ... other functions
};

void OnDeviceModelAdaptationController::MaybeRegisterAdaptation(
    ModelBasedCapabilityKey feature,
    bool was_recently_used) {

  // Check if base model is ready
  if (!base_model_available_) {
    return;
  }

  // Get base model spec
  const OnDeviceBaseModelSpec& spec = GetBaseModelSpec();

  // Register for adaptation download
  adaptation_loader_map_->MaybeRegisterModelDownload(
      feature, spec, was_recently_used);
}

std::unique_ptr<on_device_model::AdaptationAssetPaths> MaybeGetAdaptationPaths(
    const optimization_guide::ModelInfo& model_info) {

  // ⭐ Look for "adaptation_weights.bin"
  auto weights_file = model_info.GetAdditionalFileWithBaseName(
      kOnDeviceModelAdaptationWeightsFile);  // "adaptation_weights.bin"

  if (!weights_file) {
    return nullptr;  // No adaptation available
  }

  auto adaptation_assets =
      std::make_unique<on_device_model::AdaptationAssetPaths>();
  adaptation_assets->weights = *weights_file;  // ⭐ Just the path!

  return adaptation_assets;
}

// components/optimization_guide/core/optimization_guide_constants.cc:65-66
const base::FilePath::CharType kOnDeviceModelAdaptationWeightsFile[] =
    FILE_PATH_LITERAL("adaptation_weights.bin");  // ⭐ HARDCODED NAME

std::optional<OnDeviceModelAdaptationAvailability>
DetectBaseModelIncompatibility(
    const optimization_guide::ModelInfo& model_info,
    const OnDeviceBaseModelSpec& registered_spec) {

  // Parse adaptation metadata
  const std::optional<proto::Any>& metadata = model_info.GetModelMetadata();
  auto supported_model_spec =
      ParsedAnyMetadata<proto::OnDeviceBaseModelMetadata>(metadata.value());

  if (!supported_model_spec) {
    return OnDeviceModelAdaptationAvailability::kAdaptationModelInvalid;
  }

  // ⭐ CHECK 1: Base model name & version
  if (supported_model_spec->base_model_name() != registered_spec.model_name ||
      supported_model_spec->base_model_version() != registered_spec.model_version) {
    return OnDeviceModelAdaptationAvailability::kAdaptationModelIncompatible;
  }

  // ⭐ CHECK 2: Performance hints compatibility
  if (!ArePerformanceHintsCompatible(*supported_model_spec, registered_spec)) {
    return OnDeviceModelAdaptationAvailability::kAdaptationModelHintsIncompatible;
  }

  return std::nullopt;  // Compatible!
}

void OnDeviceModelAdaptationLoader::OnModelUpdated(
    proto::OptimizationTarget optimization_target,
    base::optional_ref<const ModelInfo> model_info) {

  // Get execution config file
  auto execution_config_file = model_info->GetAdditionalFileWithBaseName(
      kOnDeviceModelExecutionConfigFile);  // "on_device_model_execution_config.pb"

  if (!execution_config_file) {
    // Invalid - need config to know how to use the LoRA
    on_load_fn_.Run(base::unexpected(AdaptationUnavailability::kNotSupported));
    return;
  }

  // Read and parse protobuf config on background thread
  background_task_runner_->PostTaskAndReplyWithResult(
      FROM_HERE,
      base::BindOnce(&ReadOnDeviceModelExecutionConfig, *execution_config_file),
      base::BindOnce(&CreateAdaptationMetadataFromModelExecutionConfig,
                     feature_, MaybeGetAdaptationPaths(*model_info),
                     model_info->GetVersion())
          .Then(base::BindOnce(&OnDeviceModelAdaptationMetadataCreated, feature_))
          .Then(on_load_fn_));
}

void SessionAccessor::CreateInternal(
    on_device_model::mojom::SessionParamsPtr params,
    on_device_model::mojom::LoadAdaptationParamsPtr adaptation_params,
    std::optional<uint32_t> adaptation_id) {

  // Prepare base descriptor
  ChromeMLAdaptationDescriptor descriptor = {
    .max_tokens = params->max_tokens,
    .top_k = params->top_k,
    .temperature = params->temperature,
    .enable_image_input = params->capabilities.image_input,
    .enable_audio_input = params->capabilities.audio_input,
    .model_data = nullptr  // Will be set below if LoRA present
  };

  ChromeMLModelData data;
  std::string weights_path_str;

  // ⭐ IF ADAPTATION IS PRESENT
  if (adaptation_params) {
    weights_path_str = adaptation_params->assets.weights_path.AsUTF8Unsafe();

    if (adaptation_params->assets.weights.IsValid() || !weights_path_str.empty()) {

      // GPU backend: use file handle
      if (adaptation_params->assets.weights.IsValid()) {
        data.weights_file = adaptation_params->assets.weights.TakePlatformFile();
      }
      // APU backend: use file path
      else {
        data.model_path = weights_path_str.data();
      }

      data.file_id = adaptation_id;  // For caching
      descriptor.model_data = &data;  // ⭐ ATTACH LORA DATA
    }
  }

  // ⭐ CREATE SESSION with base model + optional LoRA
  session_ = chrome_ml_->api().CreateSession(model_, &descriptor);
}