pinecone_utils.py

import functools
import pinecone
import openai
import pandas as pd

from local_secrets import secrets

pinecone.init(api_key=secrets.PINECONE_KEY, environment="us-east1-gcp")
openai.api_key = secrets.OPENAI_KEY
TARGET_MODEL = "text-embedding-ada-002"


def pretty_timestamp(time_ms):
    return pd.to_datetime(time_ms, unit="ms").strftime("%H:%M:%S")


@functools.cache
def _get_index() -> pinecone.Index:
    return pinecone.Index("sense")


def ask_pinecone(question: str, max_results=20) -> list[dict]:
    index = _get_index()
    question_vec = openai.Embedding.create(input=question, engine=TARGET_MODEL)["data"][
        0
    ]["embedding"]

    import pdb

    results = index.query(
        vector=question_vec,
        top_k=max_results,
        include_metadata=True,
        include_values=False,
    )["matches"]

    results = [result["metadata"] for result in results]
    for result in results:
        result["start"] = pretty_timestamp(result["start"])
        result["end"] = pretty_timestamp(result["end"])

    return results