vectors.py

from embedding import strings_to_vectors
import pinecone
import os

PINECONE_API = os.getenv("PINECONE_API")

pinecone.init(api_key=PINECONE_API, environment="us-west4-gcp-free")

vector_index = pinecone.Index("quesmed")


def scored_vector_todict(scored_vector):
    x = {
        "id": scored_vector["id"],
        "metadata": {
            "topicId": int(scored_vector["metadata"]["topicId"]),
            "chapterId": int(scored_vector["metadata"]["chapterId"]),
            "conceptId": int(scored_vector["metadata"]["conceptId"]),
        },
        "score": scored_vector["score"] * 100,
        "values": scored_vector["values"],
    }
    for k, v in x["metadata"].items():
        x[k] = int(v)
    x["passage_idx"] = int(x["id"][-1])
    return x


def match_query(query: str, n_res=3):
    queries = [f"query: {query.replace('?','').lower()}"]
    query_embeddings = strings_to_vectors(queries)
    result = vector_index.query(
        query_embeddings[0].tolist(),
        top_k=n_res,
        include_metadata=True,
        namespace="quesbook",
    )
    return list(map(scored_vector_todict, result["matches"]))