! pip3 install pyvespa llama-index

from typing import List

def synthetic_mail_data_generator() -> List[dict]:
    synthetic_mails = [
        {
            "id": 1,
            "groupname": "bergum@vespa.ai",
            "fields": {
                "subject": "LlamaIndex news, 2023-11-14",
                "to": "bergum@vespa.ai",
                "body": """Hello Llama Friends 🦙 LlamaIndex is 1 year old this week! 🎉 To celebrate, we're taking a stroll down memory 
                    lane on our blog with twelve milestones from our first year. Be sure to check it out.""",
                "from": "news@llamaindex.ai",
                "display_date": "2023-11-15T09:00:00Z"
            }
        },
        {
            "id": 2,
            "groupname": "bergum@vespa.ai",
            "fields": {
                "subject": "Dentist Appointment Reminder",
                "to": "bergum@vespa.ai",
                "body": "Dear Jo Kristian ,\nThis is a reminder for your upcoming dentist appointment on 2023-12-04 at 09:30. Please arrive 15 minutes early.\nBest regards,\nDr. Dentist",
                "from": "dentist@dentist.no",
                "display_date": "2023-11-15T15:30:00Z"
            }
        },
        {
            "id": 1,
            "groupname": "giraffe@wildlife.ai",
            "fields": {
                "subject": "Wildlife Update: Giraffe Edition",
                "to": "giraffe@wildlife.ai",
                "body": "Dear Wildlife Enthusiasts 🦒, We're thrilled to share the latest insights into giraffe behavior in the wild. Join us on an adventure as we explore their natural habitat and learn more about these majestic creatures.",
                "from": "updates@wildlife.ai",
                "display_date": "2023-11-12T14:30:00Z"
        }
        },
        {
            "id": 1,
            "groupname": "penguin@antarctica.ai",
            "fields": {
                "subject": "Antarctica Expedition: Penguin Chronicles",
                "to": "penguin@antarctica.ai",
                "body": "Greetings Explorers 🐧, Our team is embarking on an exciting expedition to Antarctica to study penguin colonies. Stay tuned for live updates and behind-the-scenes footage as we dive into the world of these fascinating birds.",
                "from": "expedition@antarctica.ai",
                "display_date": "2023-11-11T11:45:00Z"
            }
        },
        {
            "id": 1,
            "groupname": "space@exploration.ai",
            "fields": {
                "subject": "Space Exploration News: November Edition",
                "to": "space@exploration.ai",
                "body": "Hello Space Enthusiasts 🚀, Join us as we highlight the latest discoveries and breakthroughs in space exploration. From distant galaxies to new technologies, there's a lot to explore!",
                "from": "news@exploration.ai",
                "display_date": "2023-11-01T16:20:00Z"
            }
         },
        {
            "id": 1,
            "groupname": "ocean@discovery.ai",
            "fields": {
                "subject": "Ocean Discovery: Hidden Treasures Unveiled",
                "to": "ocean@discovery.ai",
                "body": "Dear Ocean Explorers 🌊, Dive deep into the secrets of the ocean with our latest discoveries. From undiscovered species to underwater landscapes, our team is uncovering the wonders of the deep blue.",
                "from": "discovery@ocean.ai",
                "display_date": "2023-10-01T10:15:00Z"
            }
        }
    ]
    for mail in synthetic_mails:
        yield mail  

from typing import List

def synthetic_calendar_data_generator() -> List[dict]:
    calendar_data = [
       
        {
            "id": 1,
            "groupname": "bergum@vespa.ai",
            "fields": {
                "subject": "Dentist Appointment",
                "to": "bergum@vespa.ai",
                "body": "Dentist appointment at 2023-12-04 at 09:30 - 1 hour duration",
                "from": "dentist@dentist.no",
                "display_date": "2023-11-15T15:30:00Z",
                "duration": 60,
            }
        },
         {
            "id": 2,
            "groupname": "bergum@vespa.ai",
            "fields": {
                "subject": "Public Cloud Platform Events",
                "to": "bergum@vespa.ai",
                "body": "The cloud team continues to push new features and improvements to the platform. Join us for a live demo of the latest updates",
                "from": "public-cloud-platform-events",
                "display_date": "2023-11-21T09:30:00Z",
                "duration": 60,
            }
        }
    ]
    for event in calendar_data:
        yield event

from vespa.package import Schema, Document, Field, FieldSet, HNSW
mail_schema = Schema(
            name="mail",
            mode="streaming",
            document=Document(
                fields=[
                    Field(name="id", type="string", indexing=["summary", "index"]),
                    Field(name="subject", type="string", indexing=["summary", "index"]),
                    Field(name="to", type="string", indexing=["summary", "index"]),
                    Field(name="from", type="string", indexing=["summary", "index"]),
                    Field(name="body", type="string", indexing=["summary", "index"]),
                    Field(name="display_date", type="string", indexing=["summary"]),
                    Field(name="timestamp", type="long", indexing=["input display_date", "to_epoch_second", "summary", "attribute"], is_document_field=False),
                    Field(name="embedding", type="tensor<bfloat16>(x[384])",
                        indexing=["\"passage: \" . input subject .\" \". input body", "embed e5", "attribute", "index"],
                        ann=HNSW(distance_metric="angular"),
                        is_document_field=False
                    )
                ],
            ),
            fieldsets=[
                FieldSet(name = "default", fields = ["subject", "body", "to", "from"])
            ]
)

from vespa.package import Schema, Document, Field, FieldSet, HNSW
calendar_schema = Schema(
            name="calendar",
            inherits="mail",
            mode="streaming",
            document=Document(inherits="mail",
                fields=[
                    Field(name="duration", type="int", indexing=["summary", "index"]),
                    Field(name="guests", type="array<string>", indexing=["summary", "index"]),
                    Field(name="location", type="string", indexing=["summary", "index"]),
                    Field(name="url", type="string", indexing=["summary", "index"]),
                    Field(name="address", type="string", indexing=["summary", "index"])
                ]
            )
)

from vespa.package import ApplicationPackage, Component, Parameter

vespa_app_name = "assistant"
vespa_application_package = ApplicationPackage(
        name=vespa_app_name,
        schema=[mail_schema, calendar_schema],
        components=[Component(id="e5", type="hugging-face-embedder",
            parameters=[
                Parameter("transformer-model", {"url": "https://github.com/vespa-engine/sample-apps/raw/master/simple-semantic-search/model/e5-small-v2-int8.onnx"}),
                Parameter("tokenizer-model", {"url": "https://raw.githubusercontent.com/vespa-engine/sample-apps/master/simple-semantic-search/model/tokenizer.json"})
            ]
        )]
) 

from vespa.package import RankProfile, Function, GlobalPhaseRanking, FirstPhaseRanking

keywords_and_freshness = RankProfile(
    name="default", 
    functions=[Function(
        name="my_function", expression="nativeRank(subject) + nativeRank(body) + freshness(timestamp)"
    )],
    first_phase=FirstPhaseRanking(
        expression="my_function",
        rank_score_drop_limit=0.02
    ),
    match_features=["nativeRank(subject)", "nativeRank(body)", "my_function", "freshness(timestamp)"],
)

semantic = RankProfile(
    name="semantic", 
    functions=[Function(
        name="cosine", expression="max(0,cos(distance(field, embedding)))"
    )],
    inputs=[("query(q)", "tensor<float>(x[384])"), ("query(threshold)","", "0.75")],
    first_phase=FirstPhaseRanking(
        expression="if(cosine > query(threshold), cosine, -1)",
        rank_score_drop_limit=0.1
    ),
    match_features=["cosine", "freshness(timestamp)", "distance(field, embedding)", "query(threshold)"],
)

fusion = RankProfile(
    name="fusion",
    inherits="semantic",
    functions=[
        Function(
            name="keywords_and_freshness", expression=" nativeRank(subject) + nativeRank(body) + freshness(timestamp)"
        ),
        Function(
            name="semantic", expression="cos(distance(field,embedding))"
        )

    ],
    inputs=[("query(q)", "tensor<float>(x[384])"), ("query(threshold)", "", "0.75")],
    first_phase=FirstPhaseRanking(
        expression="if(cosine > query(threshold), cosine, -1)",
        rank_score_drop_limit=0.1
    ),
    match_features=["nativeRank(subject)", "keywords_and_freshness", "freshness(timestamp)", "cosine", "query(threshold)"],
    global_phase=GlobalPhaseRanking(
        rerank_count=1000,
        expression="reciprocal_rank_fusion(semantic, keywords_and_freshness)"
    )
)

application_directory="my-assistant-vespa-app"
vespa_application_package.to_files(application_directory)
import os

def print_files_in_directory(directory):
    for root, _, files in os.walk(directory):
        for file in files:
            print(os.path.join(root, file))
print_files_in_directory(application_directory)

from vespa.io import VespaResponse

def callback(response:VespaResponse, id:str):
    if not response.is_successful():
        print(f"Error {response.url} : {response.get_json()}")
    else:
        print(f"Success {response.url}")

app.feed_iterable(synthetic_mail_data_generator(), schema="mail", namespace="assistant", callback=callback)
app.feed_iterable(synthetic_calendar_data_generator(), schema="calendar", namespace="assistant", callback=callback)

from vespa.io import VespaQueryResponse
import json

response:VespaQueryResponse = app.query(
    yql="select subject, display_date, to from sources mail where userQuery()",
    query="when is my dentist appointment", 
    groupname="bergum@vespa.ai", 
    ranking="default"
)
assert(response.is_successful())
print(json.dumps(response.hits[0], indent=2))

from vespa.io import VespaQueryResponse
import json

response:VespaQueryResponse = app.query(
    yql="select subject, display_date from mail where {targetHits:10}nearestNeighbor(embedding,q)",
    groupname="bergum@vespa.ai", 
    ranking="semantic",
    body={
        "input.query(q)": "embed(e5, \"when is my dentist appointment\")",
    }
)
assert(response.is_successful())
print(json.dumps(response.hits[0], indent=2))

from llama_index.core import BaseRetriever
from llama_index.schema import NodeWithScore, QueryBundle, TextNode
from llama_index.callbacks.base import CallbackManager

from vespa.application import Vespa
from vespa.io import VespaQueryResponse

from typing import List, Union, Optional

class PersonalAssistantVespaRetriever(BaseRetriever):

   def __init__(
      self,
      app: Vespa,
      user: str,
      hits: int = 5,
      vespa_rank_profile: str = "default",
      vespa_score_cutoff: float = 0.70,
      sources: List[str] = ["mail"],
      fields: List[str] = ["subject", "body"],
      callback_manager: Optional[CallbackManager] = None
   ) -> None:
      """Sample Retriever for a personal assistant application.
      Args:
      param: app: Vespa application object
      param: user: user id to retrieve documents for (used for Vespa streaming groupname)
      param: hits: number of hits to retrieve from Vespa app
      param: vespa_rank_profile: Vespa rank profile to use
      param: vespa_score_cutoff: Vespa score cutoff to use during first-phase ranking
      param: sources: sources to retrieve documents from
      param: fields: fields to retrieve
      """
 
      self.app = app
      self.hits = hits
      self.user = user
      self.vespa_rank_profile = vespa_rank_profile
      self.vespa_score_cutoff = vespa_score_cutoff
      self.fields = fields
      self.summary_fields = ",".join(fields)
      self.sources = ",".join(sources)
      super().__init__(callback_manager)

   def _retrieve(self, query:Union[str,QueryBundle]) -> List[NodeWithScore]:
      """Retrieve documents from Vespa application.
      """
      if isinstance(query, QueryBundle):
         query = query.query_str
      
      if self.vespa_rank_profile == 'default':
         yql:str = f"select {self.summary_fields} from mail where userQuery()"
      else:
         yql = f"select {self.summary_fields} from sources {self.sources} where {targetHits:10}nearestNeighbor(embedding,q) or userQuery()"
      vespa_body_request = {
         "yql" : yql,
         "query": query,
         "hits": self.hits,
         "ranking.profile": self.vespa_rank_profile,
         "timeout": "1s",
         "input.query(threshold)": self.vespa_score_cutoff,
      }
      if self.vespa_rank_profile != "default":
         vespa_body_request["input.query(q)"] = f"embed(e5, \"{query}\")"

      with self.app.syncio(connections=1) as session:
         response:VespaQueryResponse = session.query(body=vespa_body_request, groupname=self.user)
         if not response.is_successful():
            raise ValueError(f"Query request failed: {response.status_code}, response payload: {response.get_json()}")

      nodes: List[NodeWithScore] = []
      for hit in response.hits:
         response_fields:dict = hit.get('fields', {})
         text: str = ""
         for field in response_fields.keys():
            if isinstance(response_fields[field], str) and field in self.fields:
                  text += response_fields[field] + " "
         id = hit['id']
         # 
         doc = TextNode(id_=id, text=text, 
            metadata=response_fields,    
         )
         nodes.append(NodeWithScore(node=doc, score=hit['relevance']))    
      return nodes                  

retriever = PersonalAssistantVespaRetriever(
    app=app, 
    user="bergum@vespa.ai", 
    vespa_rank_profile="default"
)
retriever.retrieve("When is my dentist appointment?")

retriever = PersonalAssistantVespaRetriever(
    app=app, 
    user="bergum@vespa.ai", 
    vespa_rank_profile="semantic",
    vespa_score_cutoff=0.6,
    hits=20
)
retriever.retrieve("When is my dentist appointment?")

retriever = PersonalAssistantVespaRetriever(
    app=app, 
    user="bergum@vespa.ai", 
    vespa_rank_profile="fusion",
    sources=["calendar", "mail"],
    vespa_score_cutoff=0.80
)
retriever.retrieve("When is my dentist appointment?")