diff options
Diffstat (limited to 'R2R/r2r/main/abstractions.py')
-rwxr-xr-x | R2R/r2r/main/abstractions.py | 58 |
1 files changed, 58 insertions, 0 deletions
diff --git a/R2R/r2r/main/abstractions.py b/R2R/r2r/main/abstractions.py new file mode 100755 index 00000000..3622b22d --- /dev/null +++ b/R2R/r2r/main/abstractions.py @@ -0,0 +1,58 @@ +from typing import Optional + +from pydantic import BaseModel + +from r2r.base import ( + AsyncPipe, + EmbeddingProvider, + EvalProvider, + KGProvider, + LLMProvider, + PromptProvider, + VectorDBProvider, +) +from r2r.pipelines import ( + EvalPipeline, + IngestionPipeline, + RAGPipeline, + SearchPipeline, +) + + +class R2RProviders(BaseModel): + vector_db: Optional[VectorDBProvider] + embedding: Optional[EmbeddingProvider] + llm: Optional[LLMProvider] + prompt: Optional[PromptProvider] + eval: Optional[EvalProvider] + kg: Optional[KGProvider] + + class Config: + arbitrary_types_allowed = True + + +class R2RPipes(BaseModel): + parsing_pipe: Optional[AsyncPipe] + embedding_pipe: Optional[AsyncPipe] + vector_storage_pipe: Optional[AsyncPipe] + vector_search_pipe: Optional[AsyncPipe] + rag_pipe: Optional[AsyncPipe] + streaming_rag_pipe: Optional[AsyncPipe] + eval_pipe: Optional[AsyncPipe] + kg_pipe: Optional[AsyncPipe] + kg_storage_pipe: Optional[AsyncPipe] + kg_agent_search_pipe: Optional[AsyncPipe] + + class Config: + arbitrary_types_allowed = True + + +class R2RPipelines(BaseModel): + eval_pipeline: EvalPipeline + ingestion_pipeline: IngestionPipeline + search_pipeline: SearchPipeline + rag_pipeline: RAGPipeline + streaming_rag_pipeline: RAGPipeline + + class Config: + arbitrary_types_allowed = True |