about summary refs log tree commit diff
path: root/R2R/r2r/pipes/abstractions
diff options
context:
space:
mode:
Diffstat (limited to 'R2R/r2r/pipes/abstractions')
-rwxr-xr-xR2R/r2r/pipes/abstractions/__init__.py0
-rwxr-xr-xR2R/r2r/pipes/abstractions/generator_pipe.py58
-rwxr-xr-xR2R/r2r/pipes/abstractions/search_pipe.py62
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