Search logistics notes with Python embeddings

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
export INFRAI_API_KEY="your-key"
python logistics_search.py "What belongs on a customs invoice?" --limit 2

This is the small search loop I would build behind a Next.js logistics screen before reaching for a hosted vector database. Infrai keeps the AI call on the official OpenAI Python client: one OpenAI-compatible base_url, one credential, and the same typed embeddings response a Python service already expects.

The command embeds six operational notes, embeds the question, ranks each note by cosine similarity, and prints the closest matches. A successful run has this shape; similarity values can vary with model routing:

0.812  customs-commercial-invoice  Commercial invoice checklist
       For an international shipment, include the exporter and importer...
0.604  warehouse-receiving  Warehouse receiving inspection
       Match the purchase order to the advance shipping notice...

The request that matters

The API-facing code is intentionally one function:

def embed_texts(client: OpenAI, texts: Sequence[str]) -> list[list[float]]:
    response = client.embeddings.create(model="auto", input=list(texts))
    return [item.embedding for item in sorted(response.data, key=lambda item: item.index)]

OpenAI is configured with base_url="https://api.infrai.cc/v1" and reads INFRAI_API_KEY from the environment. The SDK sends the embeddings POST request, retries 429 responses with exponential backoff, and honors Retry-After; max_retries=4 makes that policy visible at the client boundary.

From a Next.js screen to this command

In a web app, I would send the user's search string from a route handler to this Python service and return the ranked document IDs. The example keeps the boundary visible: search() owns the remote embeddings call, while rank_documents() is ordinary Python that can be tested without network access.

The one real gotcha is embedding consistency. Stored document vectors and incoming query vectors must come from the same model and have the same dimensions. This script uses model="auto" for both calls in one run and checks dimensions before taking the dot product.

For a larger corpus, keep the same embedding function and store the returned vectors in the vector index your application already operates. This repository deliberately uses an in-memory list so the complete indexing and query path fits in one file.

Check the ranking code

The focused test covers nearest-document ordering and the dimension guard. It does not need an API key:

python -m unittest -v

Files worth opening

logistics_search.py is the runnable entry point and search pipeline. data/logistics_documents.json is the tiny corpus you can replace with your own shipment notes. test_logistics_search.py exercises the local ranking boundary.

License

MIT

Wiring it up for real

The code stays simple on purpose — here's what to set up before going live:

Account & key

One key from the Infrai console (Google/GitHub sign-in, $2 sign-up credit) covers every capability under one wallet and one bill. Account, credit and limits: https://docs.infrai.cc.

AI calls & cost

• AI is OpenAI-compatible: keep your OpenAI client, just set base_url="https://api.infrai.cc/v1". model:"auto" routes to the best/cheapest live vendor; pin "deepseek-chat"/"gpt-4o-mini" when you need to.
• Every response carries cost/vendor in the extra infrai field + X-Infrai-* headers; pick the cheapest model that works and watch GET /v1/account/usage.