utils/retriever.py

def query_pinecone_sparse(
    dense_vec,
    sparse_vec,
    top_k,
    index,
    year,
    quarter,
    ticker,
    participant_type,
    threshold=0.25,
):
    if participant_type == "Company Speaker":
        participant = "Answer"
    else:
        participant = "Question"

    if year == "All":
        if quarter == "All":
            xc = index.query(
                vector=dense_vec,
                sparse_vector=sparse_vec,
                top_k=top_k,
                filter={
                    "Year": {
                        "$in": [
                            int("2020"),
                            int("2019"),
                            int("2018"),
                            int("2017"),
                            int("2016"),
                        ]
                    },
                    "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
                    "Ticker": {"$eq": ticker},
                    "QA_Flag": {"$eq": participant},
                },
                include_metadata=True,
            )
        else:
            xc = index.query(
                vector=dense_vec,
                sparse_vector=sparse_vec,
                top_k=top_k,
                filter={
                    "Year": {
                        "$in": [
                            int("2020"),
                            int("2019"),
                            int("2018"),
                            int("2017"),
                            int("2016"),
                        ]
                    },
                    "Quarter": {"$eq": quarter},
                    "Ticker": {"$eq": ticker},
                    "QA_Flag": {"$eq": participant},
                },
                include_metadata=True,
            )
    else:
        # search pinecone index for context passage with the answer
        xc = index.query(
            vector=dense_vec,
            sparse_vector=sparse_vec,
            top_k=top_k,
            filter={
                "Year": int(year),
                "Quarter": {"$eq": quarter},
                "Ticker": {"$eq": ticker},
                "QA_Flag": {"$eq": participant},
            },
            include_metadata=True,
        )
    # filter the context passages based on the score threshold
    filtered_matches = []
    for match in xc["matches"]:
        if match["score"] >= threshold:
            filtered_matches.append(match)
    xc["matches"] = filtered_matches
    return xc


def query_pinecone(
    dense_vec,
    top_k,
    index,
    year,
    quarter,
    ticker,
    participant_type,
    threshold=0.25,
):
    if participant_type == "Company Speaker":
        participant = "Answer"
    else:
        participant = "Question"

    if year == "All":
        if quarter == "All":
            xc = index.query(
                vector=dense_vec,
                top_k=top_k,
                filter={
                    "Year": {
                        "$in": [
                            int("2020"),
                            int("2019"),
                            int("2018"),
                            int("2017"),
                            int("2016"),
                        ]
                    },
                    "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
                    "Ticker": {"$eq": ticker},
                    "QA_Flag": {"$eq": participant},
                },
                include_metadata=True,
            )
        else:
            xc = index.query(
                vector=dense_vec,
                top_k=top_k,
                filter={
                    "Year": {
                        "$in": [
                            int("2020"),
                            int("2019"),
                            int("2018"),
                            int("2017"),
                            int("2016"),
                        ]
                    },
                    "Quarter": {"$eq": quarter},
                    "Ticker": {"$eq": ticker},
                    "QA_Flag": {"$eq": participant},
                },
                include_metadata=True,
            )
    else:
        # search pinecone index for context passage with the answer
        xc = index.query(
            vector=dense_vec,
            top_k=top_k,
            filter={
                "Year": int(year),
                "Quarter": {"$eq": quarter},
                "Ticker": {"$eq": ticker},
                "QA_Flag": {"$eq": participant},
            },
            include_metadata=True,
        )
    # filter the context passages based on the score threshold
    filtered_matches = []
    for match in xc["matches"]:
        if match["score"] >= threshold:
            filtered_matches.append(match)
    xc["matches"] = filtered_matches
    return xc


def format_query(query_results):
    # extract passage_text from Pinecone search result
    context = [
        result["metadata"]["Text"] for result in query_results["matches"]
    ]
    return context


def sentence_id_combine(data, query_results, lag=1):
    # Extract sentence IDs from query results
    ids = [
        result["metadata"]["Sentence_id"]
        for result in query_results["matches"]
    ]
    # Generate new IDs by adding a lag value to the original IDs
    new_ids = [id + i for id in ids for i in range(-lag, lag + 1)]
    # Remove duplicates and sort the new IDs
    new_ids = sorted(set(new_ids))
    # Create a list of lookup IDs by grouping the new IDs in groups of lag*2+1
    lookup_ids = [
        new_ids[i : i + (lag * 2 + 1)]
        for i in range(0, len(new_ids), lag * 2 + 1)
    ]
    # Create a list of context sentences by joining the sentences
    #  corresponding to the lookup IDs
    context_list = [
        " ".join(
            data.loc[data["Sentence_id"].isin(lookup_id), "Text"].to_list()
        )
        for lookup_id in lookup_ids
    ]
    return context_list


def text_lookup(data, sentence_ids):
    context = ". ".join(data.iloc[sentence_ids].to_list())
    return context


def year_quarter_range(start_quarter, start_year, end_quarter, end_year):
    """Creates a list of all (year, quarter) pairs that lie in the range including the start and end quarters."""
    start_year = int(start_year)
    end_year = int(end_year)

    quarters = (
        [("Q1", "Q2", "Q3", "Q4")] * (end_year - start_year)
        + [("Q1", "Q2", "Q3" if end_quarter == "Q4" else "Q4")]
        * (end_quarter == "Q4")
        + [
            (
                "Q1"
                if start_quarter == "Q1"
                else "Q2"
                if start_quarter == "Q2"
                else "Q3"
                if start_quarter == "Q3"
                else "Q4",
            )
            * (end_year - start_year)
        ]
    )
    years = list(range(start_year, end_year + 1))
    list_year_quarter = [
        (y, q) for y in years for q in quarters[years.index(y)]
    ]
    # Remove duplicate pairs
    seen = set()
    list_year_quarter_cleaned = []
    for tup in list_year_quarter:
        if tup not in seen:
            seen.add(tup)
            list_year_quarter_cleaned.append(tup)
    return list_year_quarter_cleaned


def multi_document_query(
    dense_query_embedding,
    sparse_query_embedding,
    num_results,
    pinecone_index,
    start_quarter,
    start_year,
    end_quarter,
    end_year,
    ticker,
    participant_type,
    threshold,
):
    pass