first commit for rd

r2r/core/__init__.py

@@ -1,6 +1,7 @@
 from .abstractions.document import BasicDocument
 from .abstractions.output import RAGPipelineOutput
 from .pipelines.embedding import EmbeddingPipeline
+from .pipelines.extraction import EntityExtractionPipeline
 from .pipelines.eval import EvalPipeline
 from .pipelines.ingestion import IngestionPipeline
 from .pipelines.rag import RAGPipeline
@@ -20,6 +21,7 @@
     "DefaultPromptProvider",
     "RAGPipelineOutput",
     "EmbeddingPipeline",
+    "EntityExtractionPipeline",
     "EvalPipeline",
     "IngestionPipeline",
     "RAGPipeline",

r2r/core/pipelines/extraction.py

@@ -0,0 +1,37 @@
+from abc import abstractmethod
+from typing import Iterator, Optional
+
+from ..providers.prompt import PromptProvider
+from ..providers.llm import LLMProvider
+from ..providers.logging import LoggingDatabaseConnection, log_execution_to_db
+from r2r.core import BasicDocument, GenerationConfig
+from r2r.pipelines import Pipeline
+
+class EntityExtractionPipeline(Pipeline):
+    def __init__(
+        self,
+        llm: LLMProvider,
+        prompt_provider: PromptProvider,
+        logging_connection: Optional[LoggingDatabaseConnection] = None,
+        *args,
+        **kwargs,
+    ):
+        self.llm = llm
+        self.prompt_provider = prompt_provider
+        super().__init__(logging_connection=logging_connection, **kwargs)
+
+    @abstractmethod
+    def preprocess_text(self, text: str) -> str:
+        pass
+
+    @abstractmethod
+    def extract_entities(self, text: str, generation_config: GenerationConfig) -> list[str]:
+        pass
+
+    @abstractmethod
+    def postprocess_entities(self, entities: list[str]) -> list[str]:
+        pass
+
+    @abstractmethod
+    def run(self, documents: Iterator[BasicDocument]) -> Iterator[BasicDocument]:
+        pass

r2r/core/pipelines/rag.py

@@ -26,7 +26,7 @@ class RAGPipeline(Pipeline):
 
     def __init__(
         self,
-        llm: "LLMProvider",
+        llm: LLMProvider,
         prompt_provider: PromptProvider,
         logging_connection: Optional[LoggingDatabaseConnection] = None,
         *args,

r2r/pipelines/basic/extraction.py

@@ -0,0 +1,53 @@
+from typing import Iterator
+
+from r2r.core import BasicDocument, EntityExtractionPipeline, GenerationConfig
+from r2r.pipelines import BasicPromptProvider
+
+class BasicEntityExtractionPipeline(EntityExtractionPipeline):
+    BASIC_SYSTEM_PROMPT = "You are a helpful assistant."
+    BASIC_TASK_PROMPT = """
+    ## Task:
+    Extract the named entities from the following text document, and return them in a comma-separated list.
+
+    ## Response:
+    """
+    def __init__(self, llm, logging_connection=None, *args, **kwargs):
+        super().__init__(prompt_provider=BasicPromptProvider(BasicEntityExtractionPipeline.BASIC_SYSTEM_PROMPT, BasicEntityExtractionPipeline.BASIC_TASK_PROMPT), logging_connection=logging_connection, **kwargs)
+        self.llm = llm
+
+    def preprocess_text(self, text: str) -> str:
+        # Optional - Implement text preprocessing logic here
+        return text
+
+    def extract_entities(self, text: str, generation_config: GenerationConfig) -> list[str]:
+        # entities = self.com
+        self._check_pipeline_initialized()
+        messages = [
+            {
+                "role": "system",
+                "content": self.prompt_provider.get_prompt("system_prompt"),
+
+            },
+            {
+                "role": "user",
+                "content": self.prompt_provider.get_prompt("task_prompt"),
+            },
+        ]
+        entities_list =  self.llm.get_completion(text, generation_config)
+        if not "," in entities_list:
+            entities = []
+        else:
+            entities = entities_list.split(",")
+        return entities
+
+    def postprocess_entities(self, entities: list[str]) -> list[str]:
+        # Implement entity postprocessing logic here
+        return [entity.upper() for entity in entities]
+
+    def run(self, documents: Iterator[BasicDocument]) -> Iterator[BasicDocument]:
+        for document in documents:
+            preprocessed_text = self.preprocess_text(document.text)
+            entities = self.extract_entities(preprocessed_text)
+            postprocessed_entities = self.postprocess_entities(entities)
+            document.metadata["entities"] = postprocessed_entities
+            yield document