tjdify/api/services/hit_testing_service.py

import logging
import time
from typing import List

import numpy as np
from llama_index.data_structs.node_v2 import NodeWithScore
from llama_index.indices.query.schema import QueryBundle
from llama_index.indices.vector_store import GPTVectorStoreIndexQuery
from sklearn.manifold import TSNE

from core.docstore.empty_docstore import EmptyDocumentStore
from core.index.vector_index import VectorIndex
from extensions.ext_database import db
from models.account import Account
from models.dataset import Dataset, DocumentSegment, DatasetQuery
from services.errors.index import IndexNotInitializedError


class HitTestingService:
    @classmethod
    def retrieve(cls, dataset: Dataset, query: str, account: Account, limit: int = 10) -> dict:
        index = VectorIndex(dataset=dataset).query_index

        if not index:
            raise IndexNotInitializedError()

        index_query = GPTVectorStoreIndexQuery(
            index_struct=index.index_struct,
            service_context=index.service_context,
            vector_store=index.query_context.get('vector_store'),
            docstore=EmptyDocumentStore(),
            response_synthesizer=None,
            similarity_top_k=limit
        )

        query_bundle = QueryBundle(
            query_str=query,
            custom_embedding_strs=[query],
        )

        query_bundle.embedding = index.service_context.embed_model.get_agg_embedding_from_queries(
            query_bundle.embedding_strs
        )

        start = time.perf_counter()
        nodes = index_query.retrieve(query_bundle=query_bundle)
        end = time.perf_counter()
        logging.debug(f"Hit testing retrieve in {end - start:0.4f} seconds")

        dataset_query = DatasetQuery(
            dataset_id=dataset.id,
            content=query,
            source='hit_testing',
            created_by_role='account',
            created_by=account.id
        )

        db.session.add(dataset_query)
        db.session.commit()

        return cls.compact_retrieve_response(dataset, query_bundle, nodes)

    @classmethod
    def compact_retrieve_response(cls, dataset: Dataset, query_bundle: QueryBundle, nodes: List[NodeWithScore]):
        embeddings = [
            query_bundle.embedding
        ]

        for node in nodes:
            embeddings.append(node.node.embedding)

        tsne_position_data = cls.get_tsne_positions_from_embeddings(embeddings)

        query_position = tsne_position_data.pop(0)

        i = 0
        records = []
        for node in nodes:
            index_node_id = node.node.doc_id

            segment = db.session.query(DocumentSegment).filter(
                DocumentSegment.dataset_id == dataset.id,
                DocumentSegment.enabled == True,
                DocumentSegment.status == 'completed',
                DocumentSegment.index_node_id == index_node_id
            ).first()

            if not segment:
                i += 1
                continue

            record = {
                "segment": segment,
                "score": node.score,
                "tsne_position": tsne_position_data[i]
            }

            records.append(record)

            i += 1

        return {
            "query": {
                "content": query_bundle.query_str,
                "tsne_position": query_position,
            },
            "records": records
        }

    @classmethod
    def get_tsne_positions_from_embeddings(cls, embeddings: list):
        embedding_length = len(embeddings)
        if embedding_length <= 1:
            return [{'x': 0, 'y': 0}]

        concatenate_data = np.array(embeddings).reshape(embedding_length, -1)
        # concatenate_data = np.concatenate(embeddings)

        perplexity = embedding_length / 2 + 1
        if perplexity >= embedding_length:
            perplexity = max(embedding_length - 1, 1)

        tsne = TSNE(n_components=2, perplexity=perplexity, early_exaggeration=12.0)
        data_tsne = tsne.fit_transform(concatenate_data)

        tsne_position_data = []
        for i in range(len(data_tsne)):
            tsne_position_data.append({'x': float(data_tsne[i][0]), 'y': float(data_tsne[i][1])})

        return tsne_position_data