Language Models (LLMs)

[completions]
provider = "litellm"
[completions.generation_config]
model = "gpt-4"
temperature = 0.7
max_tokens = 150

1
from r2r.base import ProviderConfig
2
from r2r.base.abstractions.llm import GenerationConfig
3
from typing import Optional
4
5
class LLMConfig(ProviderConfig):
6
    provider: Optional[str] = None
7
    generation_config: Optional[GenerationConfig] = None
8
9
    def validate(self) -> None:
10
        if not self.provider:
11
            raise ValueError("Provider must be set.")
12
        if self.provider and self.provider not in self.supported_providers:
13
            raise ValueError(f"Provider '{self.provider}' is not supported.")
14
15
    @property
16
    def supported_providers(self) -> list[str]:
17
        return ["litellm", "openai"]

1
from abc import abstractmethod
2
from r2r.base import Provider
3
from r2r.base.abstractions.llm import GenerationConfig, LLMChatCompletion, LLMChatCompletionChunk
4
5
class LLMProvider(Provider):
6
    def __init__(self, config: LLMConfig) -> None:
7
        if not isinstance(config, LLMConfig):
8
            raise ValueError("LLMProvider must be initialized with a `LLMConfig`.")
9
        super().__init__(config)
10
11
    @abstractmethod
12
    def get_completion(
13
        self,
14
        messages: list[dict],
15
        generation_config: GenerationConfig,
16
        **kwargs,
17
    ) -> LLMChatCompletion:
18
        pass
19
20
    @abstractmethod
21
    def get_completion_stream(
22
        self,
23
        messages: list[dict],
24
        generation_config: GenerationConfig,
25
        **kwargs,
26
    ) -> LLMChatCompletionChunk:
27
        pass

1
import logging
2
from typing import Generator
3
from r2r.base import LLMProvider, LLMConfig, LLMChatCompletion, LLMChatCompletionChunk
4
from r2r.base.abstractions.llm import GenerationConfig
5
6
logger = logging.getLogger(__name__)
7
8
class CustomLLMProvider(LLMProvider):
9
    def __init__(self, config: LLMConfig) -> None:
10
        super().__init__(config)
11
        # Initialize any custom attributes or connections here
12
        self.custom_client = self._initialize_custom_client()
13
14
    def _initialize_custom_client(self):
15
        # Initialize your custom LLM client here
16
        pass
17
18
    def get_completion(
19
        self,
20
        messages: list[dict],
21
        generation_config: GenerationConfig,
22
        **kwargs,
23
    ) -> LLMChatCompletion:
24
        # Implement the logic to get a completion from your custom LLM
25
        response = self.custom_client.generate(messages, **generation_config.dict(), **kwargs)
26
27
        # Convert the response to LLMChatCompletion format
28
        return LLMChatCompletion(
29
            id=response.id,
30
            choices=[
31
                {
32
                    "message": {
33
                        "role": "assistant",
34
                        "content": response.text
35
                    },
36
                    "finish_reason": response.finish_reason
37
                }
38
            ],
39
            usage={
40
                "prompt_tokens": response.usage.prompt_tokens,
41
                "completion_tokens": response.usage.completion_tokens,
42
                "total_tokens": response.usage.total_tokens
43
            }
44
        )
45
46
    def get_completion_stream(
47
        self,
48
        messages: list[dict],
49
        generation_config: GenerationConfig,
50
        **kwargs,
51
    ) -> Generator[LLMChatCompletionChunk, None, None]:
52
        # Implement the logic to get a streaming completion from your custom LLM
53
        stream = self.custom_client.generate_stream(messages, **generation_config.dict(), **kwargs)
54
55
        for chunk in stream:
56
            yield LLMChatCompletionChunk(
57
                id=chunk.id,
58
                choices=[
59
                    {
60
                        "delta": {
61
                            "role": "assistant",
62
                            "content": chunk.text
63
                        },
64
                        "finish_reason": chunk.finish_reason
65
                    }
66
                ]
67
            )
68
69
    # Add any additional methods specific to your custom provider
70
    def custom_method(self, *args, **kwargs):
71
        # Implement custom functionality
72
        pass

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

1
[completions]
2
provider = "custom"
3
4
[completions.generation_config]
5
model = "your-custom-model"
6
temperature = 0.7
7
max_tokens = 150

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

1
messages = [
2
    {"role": "system", "content": "You are a helpful assistant."},
3
    {"role": "user", "content": "What is the capital of France?"}
4
]
5
6
response = r2r.get_completion(messages)
7
print(response.choices[0].message.content)

1
class R2RConfig:
2
    REQUIRED_KEYS: dict[str, list] = {
3
        # ... other keys ...
4
        "completions": ["provider"],
5
        # ... other keys ...
6
    }
7
8
    def __init__(self, config_data: dict[str, Any]):
9
        # Load and validate configuration
10
        # ...
11
12
        # Set LLM configuration
13
        self.completions = LLMConfig.create(**self.completions)
14
15
        # Override GenerationConfig defaults
16
        GenerationConfig.set_default(**self.completions.get("generation_config", {}))
17
18
        # ... other initialization ...