diff options
Diffstat (limited to 'R2R/r2r/pipes/abstractions')
-rwxr-xr-x | R2R/r2r/pipes/abstractions/__init__.py | 0 | ||||
-rwxr-xr-x | R2R/r2r/pipes/abstractions/generator_pipe.py | 58 | ||||
-rwxr-xr-x | R2R/r2r/pipes/abstractions/search_pipe.py | 62 |
3 files changed, 120 insertions, 0 deletions
diff --git a/R2R/r2r/pipes/abstractions/__init__.py b/R2R/r2r/pipes/abstractions/__init__.py new file mode 100755 index 00000000..e69de29b --- /dev/null +++ b/R2R/r2r/pipes/abstractions/__init__.py diff --git a/R2R/r2r/pipes/abstractions/generator_pipe.py b/R2R/r2r/pipes/abstractions/generator_pipe.py new file mode 100755 index 00000000..002ebd23 --- /dev/null +++ b/R2R/r2r/pipes/abstractions/generator_pipe.py @@ -0,0 +1,58 @@ +import uuid +from abc import abstractmethod +from typing import Any, AsyncGenerator, Optional + +from r2r.base import ( + AsyncState, + KVLoggingSingleton, + LLMProvider, + PipeType, + PromptProvider, +) +from r2r.base.abstractions.llm import GenerationConfig +from r2r.base.pipes.base_pipe import AsyncPipe + + +class GeneratorPipe(AsyncPipe): + class Config(AsyncPipe.PipeConfig): + name: str + task_prompt: str + system_prompt: str = "default_system" + + def __init__( + self, + llm_provider: LLMProvider, + prompt_provider: PromptProvider, + type: PipeType = PipeType.GENERATOR, + config: Optional[Config] = None, + pipe_logger: Optional[KVLoggingSingleton] = None, + *args, + **kwargs, + ): + super().__init__( + type=type, + config=config or self.Config(), + pipe_logger=pipe_logger, + *args, + **kwargs, + ) + self.llm_provider = llm_provider + self.prompt_provider = prompt_provider + + @abstractmethod + async def _run_logic( + self, + input: AsyncPipe.Input, + state: AsyncState, + run_id: uuid.UUID, + rag_generation_config: GenerationConfig, + *args: Any, + **kwargs: Any, + ) -> AsyncGenerator[Any, None]: + pass + + @abstractmethod + def _get_message_payload( + self, message: str, *args: Any, **kwargs: Any + ) -> list: + pass diff --git a/R2R/r2r/pipes/abstractions/search_pipe.py b/R2R/r2r/pipes/abstractions/search_pipe.py new file mode 100755 index 00000000..bb0303e0 --- /dev/null +++ b/R2R/r2r/pipes/abstractions/search_pipe.py @@ -0,0 +1,62 @@ +import logging +import uuid +from abc import abstractmethod +from typing import Any, AsyncGenerator, Optional, Union + +from r2r.base import ( + AsyncPipe, + AsyncState, + KVLoggingSingleton, + PipeType, + VectorSearchResult, +) + +logger = logging.getLogger(__name__) + + +class SearchPipe(AsyncPipe): + class SearchConfig(AsyncPipe.PipeConfig): + name: str = "default_vector_search" + search_filters: dict = {} + search_limit: int = 10 + + class Input(AsyncPipe.Input): + message: Union[AsyncGenerator[str, None], str] + + def __init__( + self, + pipe_logger: Optional[KVLoggingSingleton] = None, + type: PipeType = PipeType.SEARCH, + config: Optional[AsyncPipe.PipeConfig] = None, + *args, + **kwargs, + ): + super().__init__( + pipe_logger=pipe_logger, + type=type, + config=config, + *args, + **kwargs, + ) + + @abstractmethod + async def search( + self, + query: str, + filters: dict[str, Any] = {}, + limit: int = 10, + *args: Any, + **kwargs: Any, + ) -> AsyncGenerator[VectorSearchResult, None]: + pass + + @abstractmethod + async def _run_logic( + self, + input: Input, + state: AsyncState, + run_id: uuid.UUID, + *args: Any, + **kwargs, + ) -> AsyncGenerator[VectorSearchResult, None]: + pass |