End-to-end RAG using Elasticsearch and Cohere

1
pip install elasticsearch_serverless cohere

1
from elasticsearch_serverless import Elasticsearch, helpers
2
from getpass import getpass
3
import cohere
4
import json
5
import requests

1
ELASTICSEARCH_ENDPOINT = getpass("Elastic Endpoint: ")
2
ELASTIC_API_KEY = getpass("Elastic encoded API key: ") # Use the encoded API key
3
4
client = Elasticsearch(
5
  ELASTICSEARCH_ENDPOINT,
6
  api_key=ELASTIC_API_KEY
7
)

1
print(client.info())

1
COHERE_API_KEY = getpass("Enter Cohere API key:  ")
2
3
client.options(ignore_status=[404]).inference.delete_model(inference_id="cohere_embeddings")
4
5
client.inference.put_model(
6
    task_type="text_embedding",
7
    inference_id="cohere_embeddings",
8
    body={
9
        "service": "cohere",
10
        "service_settings": {
11
            "api_key": COHERE_API_KEY,
12
            "model_id": "embed-english-v3.0",
13
            "embedding_type": "int8",
14
            "similarity": "cosine"
15
        },
16
        "task_settings": {},
17
    },
18
)

1
client.options(ignore_status=[404]).ingest.delete_pipeline(id="cohere_embeddings")
2
3
client.ingest.put_pipeline(
4
    id="cohere_embeddings",
5
    description="Ingest pipeline for Cohere inference.",
6
    processors=[
7
        {
8
            "inference": {
9
                "model_id": "cohere_embeddings",
10
                "input_output": {
11
                    "input_field": "text",
12
                    "output_field": "text_embedding",
13
                },
14
            }
15
        }
16
    ],
17
)

1
client.indices.delete(index="cohere-wiki-embeddings", ignore_unavailable=True)
2
client.indices.create(
3
    index="cohere-wiki-embeddings",
4
    settings={"index": {"default_pipeline": "cohere_embeddings"}},
5
    mappings={
6
        "properties": {
7
            "text_embedding": {
8
                "type": "dense_vector",
9
                "dims": 1024,
10
                "element_type": "byte"
11
            },
12
            "text": {"type": "text"},
13
            "wiki_id": {"type": "integer"},
14
            "url": {"type": "text"},
15
            "views": {"type": "float"},
16
            "langs": {"type": "integer"},
17
            "title": {"type": "text"},
18
            "paragraph_id": {"type": "integer"},
19
            "id": {"type": "integer"}
20
        }
21
    },
22
)

1
url = "https://raw.githubusercontent.com/cohere-ai/cohere-developer-experience/main/notebooks/data/embed_jobs_sample_data.jsonl"
2
response = requests.get(url)
3
4
jsonl_data = response.content.decode('utf-8').splitlines()
5
6
documents = []
7
for line in jsonl_data:
8
    data_dict = json.loads(line)
9
    documents.append({
10
        "_index": "cohere-wiki-embeddings",
11
        "_source": data_dict,
12
        }
13
      )
14
15
helpers.bulk(client, documents)
16
17
print("Done indexing documents into `cohere-wiki-embeddings` index!")

1
query = "When were the semi-finals of the 2022 FIFA world cup played?"
2
3
response = client.search(
4
    index="cohere-wiki-embeddings",
5
    size=100,
6
    knn={
7
        "field": "text_embedding",
8
        "query_vector_builder": {
9
            "text_embedding": {
10
                "model_id": "cohere_embeddings",
11
                "model_text": query,
12
            }
13
        },
14
        "k": 10,
15
        "num_candidates": 50,
16
    },
17
    query={
18
      "multi_match": {
19
          "query": query,
20
          "fields": ["text", "title"]
21
        }
22
      }
23
)
24
25
raw_documents = response["hits"]["hits"]
26
27
for document in raw_documents[0:10]:
28
  print(f'Title: {document["_source"]["title"]}\nText: {document["_source"]["text"]}\n')
29
30
documents = []
31
for hit in response["hits"]["hits"]:
32
    documents.append(hit["_source"]["text"])

1
client.options(ignore_status=[404]).inference.delete_model(inference_id="cohere_rerank")
2
3
client.inference.put_model(
4
    task_type="rerank",
5
    inference_id="cohere_rerank",
6
    body={
7
        "service": "cohere",
8
        "service_settings":{
9
            "api_key": COHERE_API_KEY,
10
            "model_id": "rerank-english-v3.0"
11
           },
12
        "task_settings": {
13
            "top_n": 10,
14
        },
15
    }
16
)

1
response = client.inference.inference(
2
    inference_id="cohere_rerank",
3
    body={
4
        "query": query,
5
        "input": documents,
6
        "task_settings": {
7
            "return_documents": False
8
            }
9
        }
10
)
11
12
ranked_documents = []
13
for document in response.body["rerank"]:
14
  ranked_documents.append({
15
      "title": raw_documents[int(document["index"])]["_source"]["title"],
16
      "text": raw_documents[int(document["index"])]["_source"]["text"]
17
  })
18
19
for document in ranked_documents[0:10]:
20
  print(f"Title: {document['title']}\nText: {document['text']}\n")

1
co = cohere.Client(COHERE_API_KEY)

1
response = co.chat(
2
    message=query,
3
    documents=ranked_documents,
4
    model='command-r-plus'
5
)
6
7
source_documents = []
8
for citation in response.citations:
9
  for document_id in citation.document_ids:
10
    if document_id not in source_documents:
11
      source_documents.append(document_id)
12
13
print(f"Query: {query}")
14
print(f"Response: {response.text}")
15
print("Sources:")
16
for document in response.documents:
17
  if document['id'] in source_documents:
18
    print(f"{document['title']}: {document['text']}")