queryFromClient.py

"""
Purpose:

The cosine similarity in queryFromClient.py adds client-side semantic search to the system, complementing the servers vector search (POST /libraries/{library_id}/search/).

Also, provides example of cohere; you run prepCohereOnClient.sh and then you can use queryFromClient.py on its output.

The offline semantic search reduces server load.  A lightweight offline alternative to real time server updates.  Allows queries to run locally without updating server state.  Appropriate for smaller datasets.


Loads Embeddings:

Reads embeddings.npy (generated by prepCohereOnClient.sh via embed.py), containing an array of dictionaries with chunk_id, text, embedding, and metadata.


Generates Query Embedding:

Prompts the user for a query text.  Uses Coheres embed-english-v3.0 model with input_type=search_query to generate a 1024-dimensional embedding for the query.

Calculates similarity between the query embedding and each chunks embedding in embeddings.npy.
Produces a list of dictionaries with chunk_id, text, metadata, and similarity scores.


Sorts and Outputs:

Sorts results by similarity (descending) and limits to the top 5.  Prints JSON output and saves to query_results.json.


Example:
python3 queryFromClient.py 
Enter query text: Sample
[
  {
    "chunk_id": "b4032000-be36-4935-b620-7f73c47c1e07",
    "text": "Sample text",
    "metadata": {
      "source": "example"
    },
    "similarity": 0.5838617469023333
  }
]
"""

import numpy as np
import cohere
import os
import json

# Load embeddings
embeddings = np.load("embeddings.npy", allow_pickle=True)

# Initialize Cohere client
co = cohere.Client(os.getenv("COHERE_API_KEY"))

# Query
query = input("Enter query text: ")  # e.g., "Sample text"
try:
    query_embedding = co.embed(
        texts=[query],
        model="embed-english-v3.0",
        input_type="search_query"
    ).embeddings[0]
except cohere.CohereAPIError as e:
    print(f"Cohere API error: {e}")
    exit(1)

# Compute cosine similarity
def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

similarities = [
    {
        "chunk_id": emb["chunk_id"],
        "text": emb["text"],
        "metadata": emb["metadata"],
        "similarity": cosine_similarity(query_embedding, emb["embedding"])
    }
    for emb in embeddings
]

# Sort by similarity
similarities = sorted(similarities, key=lambda x: x["similarity"], reverse=True)[:5]

# Output
print(json.dumps(similarities, indent=2))
with open("query_results.json", "w") as f:
    json.dump(similarities, f, indent=2)