1
from r2r import R2R
2
3
app = R2R()
4
5
# Ingest a file, which will use the configured embedding model
6
response = app.ingest_files(["path/to/your/file.txt"])
7
print(f"Ingestion response: {response}")

1
# Perform a search, which will embed the query using the configured model
2
search_results = app.search("Your search query here")
3
print(f"Search results: {search_results}")

1
# Use RAG, which involves embedding for retrieval
2
rag_response = app.rag("Your question or prompt here")
3
print(f"RAG response: {rag_response}")

1
# Update existing files, which may involve re-embedding
2
update_response = app.update_files(["path/to/updated/file.txt"])
3
print(f"Update response: {update_response}")

1
from r2r.base import ProviderConfig
2
from typing import Optional
3
4
class EmbeddingConfig(ProviderConfig):
5
    provider: Optional[str] = None
6
    base_model: Optional[str] = None
7
    base_dimension: Optional[int] = None
8
    rerank_model: Optional[str] = None
9
    rerank_dimension: Optional[int] = None
10
    rerank_transformer_type: Optional[str] = None
11
    batch_size: int = 1
12
    prefixes: Optional[dict[str, str]] = None
13
14
    def validate(self) -> None:
15
        if self.provider not in self.supported_providers:
16
            raise ValueError(f"Provider '{self.provider}' is not supported.")
17
18
    @property
19
    def supported_providers(self) -> list[str]:
20
        return [None, "openai", "ollama", "sentence-transformers"]

1
from abc import abstractmethod
2
from enum import Enum
3
from r2r.base import Provider
4
from r2r.abstractions.embedding import EmbeddingPurpose
5
from r2r.abstractions.search import VectorSearchResult
6
7
class EmbeddingProvider(Provider):
8
    class PipeStage(Enum):
9
        BASE = 1
10
        RERANK = 2
11
12
    def __init__(self, config: EmbeddingConfig):
13
        if not isinstance(config, EmbeddingConfig):
14
            raise ValueError("EmbeddingProvider must be initialized with a `EmbeddingConfig`.")
15
        super().__init__(config)
16
17
    @abstractmethod
18
    def get_embedding(
19
        self,
20
        text: str,
21
        stage: PipeStage = PipeStage.BASE,
22
        purpose: EmbeddingPurpose = EmbeddingPurpose.INDEX,
23
    ):
24
        pass
25
26
    @abstractmethod
27
    def get_embeddings(
28
        self,
29
        texts: list[str],
30
        stage: PipeStage = PipeStage.BASE,
31
        purpose: EmbeddingPurpose = EmbeddingPurpose.INDEX,
32
    ):
33
        pass
34
35
    @abstractmethod
36
    def rerank(
37
        self,
38
        query: str,
39
        results: list[VectorSearchResult],
40
        stage: PipeStage = PipeStage.RERANK,
41
        limit: int = 10,
42
    ):
43
        pass
44
45
    @abstractmethod
46
    def tokenize_string(
47
        self, text: str, model: str, stage: PipeStage
48
    ) -> list[int]:
49
        pass
50
51
    def set_prefixes(self, config_prefixes: dict[str, str], base_model: str):
52
        # Implementation of prefix setting
53
        pass

1
import numpy as np
2
from r2r.base import EmbeddingProvider, EmbeddingConfig
3
from r2r.abstractions.embedding import EmbeddingPurpose
4
from r2r.abstractions.search import VectorSearchResult
5
6
class CustomEmbeddingProvider(EmbeddingProvider):
7
    def __init__(self, config: EmbeddingConfig):
8
        super().__init__(config)
9
        # Initialize any custom attributes or models here
10
        self.model = self._load_custom_model(config.base_model)
11
12
    def _load_custom_model(self, model_name):
13
        # Load your custom embedding model here
14
        pass
15
16
    def get_embedding(
17
        self,
18
        text: str,
19
        stage: EmbeddingProvider.PipeStage = EmbeddingProvider.PipeStage.BASE,
20
        purpose: EmbeddingPurpose = EmbeddingPurpose.INDEX,
21
    ) -> list[float]:
22
        # Apply prefix if available
23
        if purpose in self.prefixes:
24
            text = f"{self.prefixes[purpose]}{text}"
25
26
        # Generate embedding using your custom model
27
        embedding = self.model.encode(text)
28
        return embedding.tolist()
29
30
    def get_embeddings(
31
        self,
32
        texts: list[str],
33
        stage: EmbeddingProvider.PipeStage = EmbeddingProvider.PipeStage.BASE,
34
        purpose: EmbeddingPurpose = EmbeddingPurpose.INDEX,
35
    ) -> list[list[float]]:
36
        # Apply prefixes if available
37
        if purpose in self.prefixes:
38
            texts = [f"{self.prefixes[purpose]}{text}" for text in texts]
39
40
        # Generate embeddings in batches
41
        all_embeddings = []
42
        for i in range(0, len(texts), self.config.batch_size):
43
            batch = texts[i:i+self.config.batch_size]
44
            batch_embeddings = self.model.encode(batch)
45
            all_embeddings.extend(batch_embeddings.tolist())
46
        return all_embeddings
47
48
    def rerank(
49
        self,
50
        query: str,
51
        results: list[VectorSearchResult],
52
        stage: EmbeddingProvider.PipeStage = EmbeddingProvider.PipeStage.RERANK,
53
        limit: int = 10,
54
    ) -> list[VectorSearchResult]:
55
        if not self.config.rerank_model:
56
            return results[:limit]
57
58
        # Implement custom reranking logic here
59
        # This is a simple example using dot product similarity
60
        query_embedding = self.get_embedding(query, stage, EmbeddingPurpose.QUERY)
61
        for result in results:
62
            result.score = np.dot(query_embedding, result.embedding)
63
64
        reranked_results = sorted(results, key=lambda x: x.score, reverse=True)
65
        return reranked_results[:limit]
66
67
    def tokenize_string(
68
        self, text: str, model: str, stage: EmbeddingProvider.PipeStage
69
    ) -> list[int]:
70
        # Implement custom tokenization logic
71
        # This is a simple example using basic string splitting
72
        return [ord(char) for word in text.split() for char in word]
73
74
    # Optionally implement async versions of methods
75
    async def async_get_embedding(self, text: str, stage: EmbeddingProvider.PipeStage, purpose: EmbeddingPurpose):
76
        # Implement async version if needed
77
        return self.get_embedding(text, stage, purpose)
78
79
    async def async_get_embeddings(self, texts: list[str], stage: EmbeddingProvider.PipeStage, purpose: EmbeddingPurpose):
80
        # Implement async version if needed
81
        return self.get_embeddings(texts, stage, purpose)

1
class EmbeddingConfig(ProviderConfig):
2
    # ...existing code...
3
4
    @property
5
    def supported_providers(self) -> list[str]:
6
        return [None, "openai", "ollama", "sentence-transformers", "custom"]  # Add your custom provider here

1
[embedding]
2
provider = "custom"
3
base_model = "your-custom-model"
4
base_dimension = 768
5
batch_size = 32
6
7
[embedding.prefixes]
8
index = "Represent this document for retrieval: "
9
query = "Represent this query for retrieving relevant documents: "

1
from r2r import R2R
2
from r2r.base import EmbeddingConfig
3
from your_module import CustomEmbeddingProvider
4
5
def get_embedding_provider(config: EmbeddingConfig):
6
    if config.provider == "custom":
7
        return CustomEmbeddingProvider(config)
8
    # ... handle other providers ...
9
10
r2r = R2R(embedding_provider_factory=get_embedding_provider)

1
# Ingest documents (embeddings will be generated using your custom provider)
2
r2r.ingest_files(["path/to/document.txt"])
3
4
# Perform a search
5
results = r2r.search("Your search query")
6
7
# Use RAG
8
rag_response = r2r.rag("Your question here")