pip install -U pymilvus
pip install -U "pymilvus[model]"
pip install langchain
pip install langchain-community

import os

if not os.path.exists("chat_history.json"):
    !wget https://raw.githubusercontent.com/jina-ai/workshops/main/notebooks/embeddings/milvus/chat_history.json

import os
import getpass
os.environ["JINAAI_API_KEY"] = getpass.getpass(prompt="Jina AI API Key: ")

os.environ["HUGGINGFACEHUB_API_TOKEN"] = getpass.getpass(prompt="Hugging Face Token: ")

from pymilvus import MilvusClient, DataType

# Specify a local file name as uri parameter of MilvusClient to use Milvus Lite
client = MilvusClient("milvus_jina.db")

schema = MilvusClient.create_schema(
    auto_id=True,
    enable_dynamic_field=True,
)

schema.add_field(field_name="id", datatype=DataType.INT64, description="The Primary Key", is_primary=True)
schema.add_field(field_name="embedding", datatype=DataType.FLOAT_VECTOR, description="The Embedding Vector", dim=768)

index_params = client.prepare_index_params()
index_params.add_index(field_name="embedding", metric_type="COSINE", index_type="AUTOINDEX")

client.create_collection(collection_name="milvus_jina", schema=schema, index_params=index_params)

import json

with open("chat_history.json", "r", encoding="utf-8") as file:
    chat_data = json.load(file)
messages = []
metadatas = []

for channel in chat_data:
  chat_history = channel["chat_history"]
  chat_topic = channel["topic"]
  chat_expert = channel["expert_user"]for message in chat_history:
    text = f"""{message["user"]}: {message["message"]}"""
    messages.append(text)
    meta = {"time_stamp": message["time_stamp"],"file_name": message["file_name"],"parent_message_nr": message["parent_message_nr"],"channel": chat_topic,"expert": True if message["user"] == chat_expert else False
    }
    metadatas.append(meta)

from pymilvus.model.dense import JinaEmbeddingFunction
jina_ef = JinaEmbeddingFunction("jina-embeddings-v2-base-en")
embeddings = jina_ef.encode_documents(messages)

collection_data = [{"message": message,"embedding": embedding,"metadata": metadata
} for message, embedding, metadata in zip(messages, embeddings, metadatas)]
data = client.insert(
    collection_name="milvus_jina",
    data=collection_data
)

query = "Who knows the most about encryption protocols in my team?"

from pymilvus.model.reranker import JinaRerankFunction
query_vectors = jina_ef.encode_queries([query])
results = client.search(
    collection_name="milvus_jina",
    data=query_vectors,
    limit=5,
)
results = results[0]
ids = [results[i]["id"] for i in range(len(results))]
results = client.get(
    collection_name="milvus_jina",
    ids=ids,
    output_fields=["id", "message", "metadata"]
)
jina_rf = JinaRerankFunction("jina-reranker-v1-base-en")
documents = [results[i]["message"] for i in range(len(results))]
reranked_documents = jina_rf(query, documents)
reranked_messages = []
for reranked_document in reranked_documents:
  idx = reranked_document.index
  reranked_messages.append(results[idx])

from langchain.prompts import PromptTemplate
from langchain_community.llms import HuggingFaceEndpoint
llm = HuggingFaceEndpoint(repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1")
prompt = """<s>[INST] Context information is below.\\n
        It includes the five most relevant messages to the query, sorted based on their relevance to the query.\\n
        ---------------------\\n
        {context_str}\\\\n
        ---------------------\\n
        Given the context information and not prior knowledge,
        answer the query. Please be brief, concise, and complete.\\n
        If the context information does not contain an answer to the query,
        respond with \\"No information\\".\\n
        Query: {query_str}[/INST] </s>"""
prompt = PromptTemplate(template=prompt, input_variables=["query_str", "context_str"])
llm_chain = prompt | llm
answer = llm_chain.invoke({"query_str":query, "context_str":reranked_messages})

print(f"\n\nANSWER:\n\n{answer}")