# %pip install cohere hnswlib unstructured

import cohere

co_chat = cohere.ClientV2(
    api_key="AZURE_API_KEY_CHAT",
    base_url="AZURE_ENDPOINT_CHAT",  # example: "https://cohere-command-r-plus-08-2024-xyz.eastus.models.ai.azure.com/"
)

co_embed = cohere.ClientV2(
    api_key="AZURE_API_KEY_EMBED",
    base_url="AZURE_ENDPOINT_EMBED",  # example: "https://embed-v-4-0-xyz.eastus.models.ai.azure.com/"
)

co_rerank = cohere.ClientV2(
    api_key="AZURE_API_KEY_RERANK",
    base_url="AZURE_ENDPOINT_RERANK",  # example: "https://cohere-rerank-v3-multilingual-xyz.eastus.models.ai.azure.com/"
)
```

## A quick example

Let's begin with a simple example to explore how RAG works.

The foundation of RAG is having a set of documents for the LLM to reference. Below, we'll work with a small collection of basic documents. While RAG systems usually involve retrieving relevant documents based on the user's query (which we'll explore later), for now we'll keep it simple and use this entire small set of documents as context for the LLM.

We have seen how to use the Chat endpoint in the text generation chapter. To use the RAG feature, we simply need to add one additional parameter, `documents`, to the endpoint call. These are the documents we want to provide as the context for the model to use in its response.

```python PYTHON
documents = [
    {
        "title": "Tall penguins",
        "text": "Emperor penguins are the tallest.",
    },
    {
        "title": "Penguin habitats",
        "text": "Emperor penguins only live in Antarctica.",
    },
    {
        "title": "What are animals?",
        "text": "Animals are different from plants.",
    },
]
```

Let's see how the model responds to the question "What are the tallest living penguins?"

The model leverages the provided documents as context for its response. Specifically, when mentioning that Emperor penguins are the tallest species, it references `doc_0` - the document which states that "Emperor penguins are the tallest."

```python PYTHON
message = "What are the tallest living penguins?"

response = co_chat.chat(
    model="model",  # Pass a dummy string
    messages=[{"role": "user", "content": message}],
    documents=[{"data": doc} for doc in documents],
)

print("\nRESPONSE:\n")
print(response.message.content[0].text)

if response.message.citations:
    print("\nCITATIONS:\n")
    for citation in response.message.citations:
        print(citation)
```

```mdx
RESPONSE:

The tallest living penguins are the Emperor penguins. They only live in Antarctica.

CITATIONS:

start=36 end=53 text='Emperor penguins.' sources=[DocumentSource(type='document', id='doc:0', document={'id': 'doc:0', 'text': 'Emperor penguins are the tallest.', 'title': 'Tall penguins'})] type=None
start=59 end=83 text='only live in Antarctica.' sources=[DocumentSource(type='document', id='doc:1', document={'id': 'doc:1', 'text': 'Emperor penguins only live in Antarctica.', 'title': 'Penguin habitats'})] type=None
```

## A more comprehensive example

Now that we’ve covered a basic RAG implementation, let’s look at a more comprehensive example of RAG that includes:

* Creating a retrieval system that converts documents into text embeddings and stores them in an index
* Building a query generation system that transforms user messages into optimized search queries
* Implementing a chat interface to handle LLM interactions with users
* Designing a response generation system capable of handling various query types

First, let’s import the necessary libraries for this project. This includes `hnswlib` for the vector library and `unstructured` for chunking the documents (more details on these later).

```python PYTHON
import uuid
import yaml
import hnswlib
from typing import List, Dict
from unstructured.partition.html import partition_html
from unstructured.chunking.title import chunk_by_title
```

## Define documents

Next, we’ll define the documents we’ll use for RAG. We’ll use a few pages from the Cohere documentation that discuss prompt engineering. Each entry is identified by its title and URL.

```python PYTHON
raw_documents = [
    {
        "title": "Crafting Effective Prompts",
        "url": "https://docs.cohere.com/docs/crafting-effective-prompts",
    },
    {
        "title": "Advanced Prompt Engineering Techniques",
        "url": "https://docs.cohere.com/docs/advanced-prompt-engineering-techniques",
    },
    {
        "title": "Prompt Truncation",
        "url": "https://docs.cohere.com/docs/prompt-truncation",
    },
    {
        "title": "Preambles",
        "url": "https://docs.cohere.com/docs/preambles",
    },
]
```

## Create vectorstore

The Vectorstore class handles the ingestion of documents into embeddings (or vectors) and the retrieval of relevant documents given a query.

It includes a few methods:

* `load_and_chunk`: Loads the raw documents from the URL and breaks them into smaller chunks
* `embed`: Generates embeddings of the chunked documents
* `index`: Indexes the document chunk embeddings to ensure efficient similarity search during retrieval
* `retrieve`: Uses semantic search to retrieve relevant document chunks from the index, given a query. It involves two steps: first, dense retrieval from the index via the Embed endpoint, and second, a reranking via the Rerank endpoint to boost the search results further.

```python PYTHON
class Vectorstore:

    def __init__(self, raw_documents: List[Dict[str, str]]):
        self.raw_documents = raw_documents
        self.docs = []
        self.docs_embs = []
        self.retrieve_top_k = 10
        self.rerank_top_k = 3
        self.load_and_chunk()
        self.embed()
        self.index()

    def load_and_chunk(self) -> None:
        """
        Loads the text from the sources and chunks the HTML content.
        """
        print("Loading documents...")

        for raw_document in self.raw_documents:
            elements = partition_html(url=raw_document["url"])
            chunks = chunk_by_title(elements)
            for chunk in chunks:
                self.docs.append(
                    {
                        "data": {
                            "title": raw_document["title"],
                            "text": str(chunk),
                            "url": raw_document["url"],
                        }
                    }
                )

    def embed(self) -> None:
        """
        Embeds the document chunks using the Cohere API.
        """
        print("Embedding document chunks...")

        batch_size = 90
        self.docs_len = len(self.docs)
        for i in range(0, self.docs_len, batch_size):
            batch = self.docs[i : min(i + batch_size, self.docs_len)]
            texts = [item["data"]["text"] for item in batch]
            docs_embs_batch = co_embed.embed(
                texts=texts,
                model="embed-v4.0",
                input_type="search_document",
                embedding_types=["float"],
            ).embeddings.float
            self.docs_embs.extend(docs_embs_batch)

    def index(self) -> None:
        """
        Indexes the document chunks for efficient retrieval.
        """
        print("Indexing document chunks...")

        self.idx = hnswlib.Index(space="ip", dim=1024)
        self.idx.init_index(
            max_elements=self.docs_len, ef_construction=512, M=64
        )
        self.idx.add_items(
            self.docs_embs, list(range(len(self.docs_embs)))
        )

        print(
            f"Indexing complete with {self.idx.get_current_count()} document chunks."
        )

    def retrieve(self, query: str) -> List[Dict[str, str]]:
        """
        Retrieves document chunks based on the given query.

        Parameters:
        query (str): The query to retrieve document chunks for.

        Returns:
        List[Dict[str, str]]: A list of dictionaries representing the retrieved document chunks, with 'title', 'text', and 'url' keys.
        """

        # Dense retrieval
        query_emb = co_embed.embed(
            texts=[query],
            model="embed-v4.0",
            input_type="search_query",
            embedding_types=["float"],
        ).embeddings.float

        doc_ids = self.idx.knn_query(
            query_emb, k=self.retrieve_top_k
        )[0][0]

        # Reranking
        docs_to_rerank = [
            self.docs[doc_id]["data"] for doc_id in doc_ids
        ]
        yaml_docs = [
            yaml.dump(doc, sort_keys=False) for doc in docs_to_rerank
        ]
        rerank_results = co_rerank.rerank(
            query=query,
            documents=yaml_docs,
            model="model",  # Pass a dummy string
            top_n=self.rerank_top_k,
        )

        doc_ids_reranked = [
            doc_ids[result.index] for result in rerank_results.results
        ]

        docs_retrieved = []
        for doc_id in doc_ids_reranked:
            docs_retrieved.append(self.docs[doc_id]["data"])

        return docs_retrieved
```

## Process documents

With the Vectorstore set up, we can process the documents, which will involve chunking, embedding, and indexing.

```python PYTHON