%load_ext autoreload
%autoreload 2
import mycode.vap as vap

# The goal is to prepare a Vespa app that models documents that have many chunks, and each chunk can have many entities.
# Then score the chunks based on how many entities it matched.
# Do not return chunks that match no entities.
# The tricks are:
# - using the "layered ranking" https://blog.vespa.ai/introducing-layered-ranking-for-rag-applications/
# - to discard the low scoring chunks we can use filter_subspaces https://docs.vespa.ai/en/reference/ranking/ranking-expressions.html#filter-subspaces

# To achieve that we need one document
# Lucene linguistics would be a definite way to achieve tokenization and matching requirements, but let's first try to achieve the goal with the default linguistics library.

from vespa.package import (ApplicationPackage, Field, Schema, Document, RankProfile, Function, Summary, FirstPhaseRanking, DocumentSummary)

doc_schema = Schema(
    name='docs',
    document=Document(
        fields=[
            Field(name='id', type='int', indexing=['attribute']),
            Field(
                name='chunks',
                type='array<string>',
                indexing=['index', 'summary'],
                summary=Summary(fields=[('select-elements-by', ['best_chunks'])]),
            ),
            Field(
                # stores ids per chunk,
                # each string encodes the entities mentioned per chunk
                # e.g. ["EID123 EID321", "", "EID9"]
                # doc has 3 chunks:
                # - first chunk has mentioned 'EID123' and 'EID321' entities
                # - second chunk has no entities
                # - third chunk has mentioned one entity "EID9"
                name='entities',
                type='array<string>',
                indexing=['index', 'summary'],
                match=['text', 'cased'],
                index={'stemming': 'none'},
                summary=Summary(fields=[('select-elements-by', ['best_chunks'])]),
            ),
        ]
    ),
    rank_profiles=[
        RankProfile(
            name='default',
            first_phase=FirstPhaseRanking(
                expression='reduce(best_entities, sum, i)'
            ),
            functions=[
                Function(
                    name='best_chunks',
                    expression='top(2, best_entities)',
                ),
                Function(
                    name='best_entities',
                    expression='elementwise(bm25(entities),i,float)'
                )
            ],
            summary_features=['best_chunks'],
            match_features=['best_entities'],
            rank_properties=[
                ('elementwise(bm25(entities),i,float).k1', '0')
            ]
        )
    ],
    document_summaries=[
        DocumentSummary(
            name='tokenization',
            from_disk=True,
            summary_fields=[
                Summary(
                    name='entities_tokens',
                    fields=[('source', 'entities'), 'tokens']
                )
            ],
        )
    ]
)
application_package = ApplicationPackage(
    name='test',
    schema=[doc_schema],
)

from vespa.deployment import VespaDocker

vespa_docker = VespaDocker(
    container_image="vespaengine/vespa:8.617.12",
)
vespa_docker.deploy(application_package=application_package)

client = vap.redeploy(vespa_docker, application_package)

docs = [
    {
        'id': f'{1}',
        'fields': {
            'id': 1,
            'chunks': [
                'chunk_1',
                'chunk_2',
                'chunk_3'
            ],
            'entities': [
                'eid1 eid2',
                '',
                'eid9'
            ],
        }
    }
]
client.feed_iterable(docs, schema="docs", namespace="doc", callback=vap.feed_callback)

client.query(
    body={
        # https://docs.vespa.ai/en/ranking/significance.html
        'yql': '''
               select *
               from sources docs
               where
                   {significance:1}
                   rank(
                   true
                   , entities contains "eid1"
                   , entities contains "eid2"
                   , entities contains "eid9"
                   )
               ''',
        # 'trace.level': 2,
        'summary': 'default'
    },
).json

# Fetching the exact tokens stored in the index
client.query(
    body={
        # https://docs.vespa.ai/en/ranking/significance.html
        'yql': '''
               select *
               from sources docs
               where
                   {significance:1}
                   rank(
                   true
                   , entities contains "eid1"
                   , entities contains "eid2"
                   , entities contains "eid9"
                   )
               ''',
        # 'trace.level': 2,
        'summary': 'tokenization'
    },
).json

# Let's use the entity ids for filtering docs
# The condition `id=1` is to show that filtering can be combined with anything, e.g. nearestNeighbor matching.
client.query(
    body={
        'yql': '''
               select *
               from sources docs
               where
                   id=1
                   AND
                   ({significance:1}
                   entities contains "eid9")
               ''',
        'summary': 'default'
    },
).json