Langchain embeddings with OVHcloud AI Endpoints

langchain_embedding.py

#!/usr/bin/env python3

### Langchain embeddings with OVHcloud AI Endpoints

import os
import sys
import time
import requests
import logging

from langchain_core.embeddings import Embeddings
from typing import List

#logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
#logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

"""
Usage:
    OVH_AI_ENDPOINTS_ACCESS_TOKEN="your-token" python3 langchain_embedding.py

NB: Make sure you are using a valid token. In the contrary, document indexing will be long due to rate-limite
"""

class OVHcloudAIEEmbeddings(Embeddings):
    def __init__(self, api_key: str = os.environ.get("OVH_AI_ENDPOINTS_ACCESS_TOKEN", None), api_url: str = "https://multilingual-e5-base.endpoints.kepler.ai.cloud.ovh.net/api/text2vec"):
        self.api_key = api_key
        self.api_url = api_url

    def _generate_embedding(self, text: str) -> List[float]:
        """Generate embeddings from OVHCLOUD AIE.
        Args:
            text: str. An input text sentence or document.
        Returns:
            embeddings: a list of float numbers. Embeddings correspond to your given text.
        """
        headers = {
            "content-type": "text/plain",
            "Authorization": f"Bearer {self.api_key}",
        }

        session = requests.session()
        while True:
            response = session.post(
                self.api_url,
                headers=headers,
                data=text,
            )
            if response.status_code != 200:
                if response.status_code == 429:
                    """Rate limit exceeded, wait for reset"""
                    reset_time = int(response.headers.get("RateLimit-Reset", 0))
                    logging.info("Rate limit exceeded. Waiting %d seconds.", reset_time)
                    if reset_time > 0:
                        time.sleep(reset_time)
                        continue
                    else:
                        """Rate limit reset time has passed, retry immediately"""
                        continue

                """ Handle other non-200 status codes """
                raise ValueError(
                    f"Request failed with status code {response.status_code}: {response.text}"
                )
            return response.json()

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return [self._generate_embedding(text) for text in texts]

    def embed_query(self, text: str) -> List[float]:
        return self._generate_embedding(text)



embeddings = OVHcloudAIEEmbeddings()

query_result = embeddings.embed_query("foo")
print(query_result)
doc_results = embeddings.embed_documents(["foo"])
print(doc_results)