Fix data class caches

Author: ajhai

llmstack/data/apis.py

@@ -10,7 +10,7 @@
 from rest_framework.response import Response as DRFResponse
 
 from llmstack.base.models import VectorstoreEmbeddingEndpoint
-from llmstack.data.schemas import DataDocument
+from llmstack.data.sources.base import DataDocument
 from llmstack.data.yaml_loader import (
     get_data_pipeline_template_by_slug,
     get_data_pipelines_from_contrib,

llmstack/data/destinations/__init__.py

@@ -1,5 +1,19 @@
+from functools import cache
+
 from llmstack.data.destinations.stores.singlestore import SingleStore
+from llmstack.data.destinations.vector_stores.chromadb import ChromaDB
 from llmstack.data.destinations.vector_stores.pinecone import Pinecone
+from llmstack.data.destinations.vector_stores.qdrant import Qdrant
+from llmstack.data.destinations.vector_stores.vector_store import PromptlyVectorStore
 from llmstack.data.destinations.vector_stores.weaviate import Weaviate
 
+
+@cache
+def get_destination_cls(slug, provider_slug):
+    for cls in [ChromaDB, Weaviate, SingleStore, Pinecone, Qdrant, PromptlyVectorStore]:
+        if cls.slug() == slug and cls.provider_slug() == provider_slug:
+            return cls
+    return None
+
+
 __all__ = ["SingleStore", "Pinecone", "Weaviate"]

llmstack/data/destinations/base.py

@@ -4,7 +4,7 @@
     CustomGenerateJsonSchema,
     get_ui_schema_from_json_schema,
 )
-from llmstack.data.schemas import DataDocument
+from llmstack.data.sources.base import DataDocument
 
 
 class BaseDestination(BaseModel):

llmstack/data/destinations/utils.py

@@ -22,7 +22,7 @@
 from weaviate.connect.helpers import connect_to_custom, connect_to_wcs
 
 from llmstack.data.destinations.base import BaseDestination
-from llmstack.data.schemas import DataDocument
+from llmstack.data.sources.base import DataDocument
 from llmstack.processors.providers.weaviate import (
     APIKey,
     WeaviateCloudInstance,

llmstack/data/destinations/vector_stores/weaviate.py

@@ -9,6 +9,8 @@
 
 from llmstack.assets.models import Assets
 from llmstack.base.models import Profile
+from llmstack.data.pipeline import DataIngestionPipeline, DataQueryPipeline
+from llmstack.data.schemas import PipelineBlock
 from llmstack.events.apis import EventsViewSet
 
 logger = logging.getLogger(__name__)
@@ -122,8 +124,6 @@ def type_slug(self):
 
     @property
     def pipeline_obj(self):
-        from llmstack.data.schemas import PipelineBlock
-
         if self.config.get("pipeline"):
             return PipelineBlock(**self.config.get("pipeline"))
 
@@ -147,13 +147,9 @@ def pipeline(self):
         return self.config.get("pipeline", {})
 
     def create_data_ingestion_pipeline(self):
-        from llmstack.data.pipeline import DataIngestionPipeline
-
         return DataIngestionPipeline(self)
 
     def create_data_query_pipeline(self):
-        from llmstack.data.pipeline import DataQueryPipeline
-
         return DataQueryPipeline(self)
 
 

llmstack/data/models.py

@@ -5,8 +5,7 @@
 from llama_index.core.schema import Document as LlamaDocument
 
 from llmstack.common.blocks.data.store.vectorstore import Document
-from llmstack.data.models import DataSource
-from llmstack.data.schemas import DataDocument
+from llmstack.data.sources.base import DataDocument
 
 logger = logging.getLogger(__name__)
 
@@ -18,17 +17,15 @@ class LlamaDocumentShim(LlamaDocument):
 
 
 class DataIngestionPipeline:
-    def __init__(self, datasource: DataSource):
+    def __init__(self, datasource):
         self.datasource = datasource
         self._source_cls = self.datasource.pipeline_obj.source_cls
         self._destination_cls = self.datasource.pipeline_obj.destination_cls
-        logger.debug("Initializing DataIngestionPipeline")
 
         self._destination = None
         self._transformations = self.datasource.pipeline_obj.transformation_objs
         embedding_cls = self.datasource.pipeline_obj.embedding_cls
         if embedding_cls:
-            logger.debug("Initializing DataIngestionPipeline Transformation")
             embedding_additional_kwargs = {
                 **self.datasource.pipeline_obj.embedding.data.get("additional_kwargs", {}),
                 **{"datasource": datasource},
@@ -41,29 +38,21 @@ def __init__(self, datasource: DataSource):
                     }
                 )
             )
-            logger.debug("Finished Initializing DataIngestionPipeline Transformation")
 
         if self._destination_cls:
-            logger.debug("Initializing DataIngestionPipeline Destination")
             self._destination = self._destination_cls(**self.datasource.pipeline_obj.destination_data)
             self._destination.initialize_client(datasource=self.datasource, create_collection=True)
-            logger.debug("Finished Initializing DataIngestionPipeline Destination")
 
     def process(self, document: DataDocument) -> DataDocument:
-        logger.debug(f"Processing document: {document.name}")
         document = self._source_cls.process_document(document)
-        logger.debug(f"Creating IngestionPipeline for document: {document.name}")
         ingestion_pipeline = IngestionPipeline(transformations=self._transformations)
         ldoc = LlamaDocumentShim(**document.model_dump())
         ldoc.metadata = {**ldoc.metadata, **document.metadata}
-        logger.debug(f"Running IngestionPipeline for document: {document.name}")
         document.nodes = ingestion_pipeline.run(documents=[ldoc])
-        logger.debug(f"Finished running IngestionPipeline for document: {document.name}")
         document.node_ids = list(map(lambda x: x.id_, document.nodes))
+
         if self._destination:
-            logger.debug(f"Adding document: {document.name} to destination")
             self._destination.add(document=document)
-            logger.debug(f"Finished adding document: {document.name} to destination")
 
         return document
 
@@ -80,55 +69,44 @@ def delete_all_entries(self) -> None:
 
 
 class DataQueryPipeline:
-    def __init__(self, datasource: DataSource):
+    def __init__(self, datasource):
         self.datasource = datasource
         self._destination_cls = self.datasource.pipeline_obj.destination_cls
         self._destination = None
         self._embedding_generator = None
-        logger.debug("Initializing DataQueryPipeline")
 
         if self._destination_cls:
-            logger.debug("Initializing DataQueryPipeline Destination")
             self._destination = self._destination_cls(**self.datasource.pipeline_obj.destination_data)
             self._destination.initialize_client(datasource=self.datasource, create_collection=False)
-            logger.debug("Finished Initializing DataQueryPipeline Destination")
 
         if self.datasource.pipeline_obj.embedding:
-            logger.debug("Initializing DataQueryPipeline Embedding")
             embedding_data = self.datasource.pipeline_obj.embedding.data
             embedding_data["additional_kwargs"] = {
                 **embedding_data.get("additional_kwargs", {}),
                 **{"datasource": self.datasource},
             }
             self._embedding_generator = self.datasource.pipeline_obj.embedding_cls(**embedding_data)
-            logger.debug("Finished Initializing DataQueryPipeline Embedding")
 
     def search(self, query: str, use_hybrid_search=True, **kwargs) -> List[dict]:
         content_key = self.datasource.destination_text_content_key
         query_embedding = None
 
-        logger.debug(f"Initializing Search for query: {query}")
-
         if kwargs.get("search_filters", None):
             raise NotImplementedError("Search filters are not supported for this data source.")
 
         documents = []
 
         if self._embedding_generator:
-            logger.debug("Generating embedding for query")
             query_embedding = self._embedding_generator.get_embedding(query)
-            logger.debug("Finished generating embedding for query")
 
         if self._destination:
-            logger.debug(f"Searching for query: {query} in destination")
             query_result = self._destination.search(
                 query=query,
                 use_hybrid_search=use_hybrid_search,
                 query_embedding=query_embedding,
                 datasource_uuid=str(self.datasource.uuid),
                 **kwargs,
             )
-            logger.debug(f"Received results for query: {query} from destination")
             documents = list(
                 map(
                     lambda x: Document(page_content_key=content_key, page_content=x.text, metadata=x.metadata),

llmstack/data/pipeline.py

@@ -1,7 +1,10 @@
-import uuid
-from typing import Any, List, Optional
+from typing import List, Optional
 
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, PrivateAttr
+
+from llmstack.data.destinations import get_destination_cls
+from llmstack.data.sources import get_source_cls
+from llmstack.data.transformations import get_transformer_cls
 
 
 class BaseProcessorBlock(BaseModel):
@@ -36,26 +39,20 @@ def default_dict(self):
 
 class PipelineSource(BaseProcessorBlock):
     def __init__(self, **data):
-        from llmstack.data.sources.utils import get_source_cls
-
         super().__init__(**data)
 
         self._processor_cls = get_source_cls(slug=self.slug, provider_slug=self.provider_slug)
 
 
 class PipelineDestination(BaseProcessorBlock):
     def __init__(self, **data):
-        from llmstack.data.destinations.utils import get_destination_cls
-
         super().__init__(**data)
 
         self._processor_cls = get_destination_cls(slug=self.slug, provider_slug=self.provider_slug)
 
 
 class PipelineTransformation(BaseProcessorBlock):
     def __init__(self, **data):
-        from llmstack.data.transformations.utils import get_transformer_cls
-
         super().__init__(**data)
 
         self._processor_cls = get_transformer_cls(slug=self.slug, provider_slug=self.provider_slug)
@@ -66,8 +63,6 @@ def get_default_data(self, **kwargs):
 
 class PipelineEmbedding(BaseProcessorBlock):
     def __init__(self, **data):
-        from llmstack.data.transformations.utils import get_transformer_cls
-
         super().__init__(**data)
 
         self._processor_cls = get_transformer_cls(slug=self.slug, provider_slug=self.provider_slug)
@@ -128,20 +123,3 @@ def default_dict(self):
                 "destination": self.pipeline.destination.default_dict() if self.pipeline.destination else None,
             },
         }
-
-
-class DataDocument(BaseModel):
-    id_: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Unique ID of the document.")
-    name: Optional[str] = None
-    text: Optional[str] = None
-    text_objref: Optional[str] = None
-    content: Optional[str] = None
-    mimetype: str = Field(default="text/plain", description="MIME type of the content.")
-    metadata: Optional[dict] = None
-    extra_info: Optional[dict] = {}
-    nodes: Optional[List[Any]] = None
-    embeddings: Optional[List[float]] = None
-    processing_errors: Optional[List[str]] = None
-    datasource_uuid: Optional[str] = None
-    request_data: Optional[dict] = {}
-    node_ids: Optional[List[str]] = []

llmstack/data/schemas.py

@@ -1,5 +1,18 @@
+from functools import cache
+
+from llmstack.data.sources.files.csv import CSVFileSchema
 from llmstack.data.sources.files.file import FileSchema
+from llmstack.data.sources.files.pdf import PdfSchema
 from llmstack.data.sources.text.text_data import TextSchema
 from llmstack.data.sources.website.url import URLSchema
 
 __all__ = ["FileSchema", "TextSchema", "URLSchema"]
+
+
+@cache
+def get_source_cls(slug, provider_slug):
+    for cls in [CSVFileSchema, FileSchema, PdfSchema, URLSchema, TextSchema]:
+        if cls.slug() == slug and cls.provider_slug() == provider_slug:
+            return cls
+
+    return None

llmstack/data/sources/__init__.py

@@ -1,12 +1,29 @@
-from typing import List
+import uuid
+from typing import Any, List, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from llmstack.common.blocks.base.schema import (
     CustomGenerateJsonSchema,
     get_ui_schema_from_json_schema,
 )
-from llmstack.data.schemas import DataDocument
+
+
+class DataDocument(BaseModel):
+    id_: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Unique ID of the document.")
+    name: Optional[str] = None
+    text: Optional[str] = None
+    text_objref: Optional[str] = None
+    content: Optional[str] = None
+    mimetype: str = Field(default="text/plain", description="MIME type of the content.")
+    metadata: Optional[dict] = None
+    extra_info: Optional[dict] = {}
+    nodes: Optional[List[Any]] = None
+    embeddings: Optional[List[float]] = None
+    processing_errors: Optional[List[str]] = None
+    datasource_uuid: Optional[str] = None
+    request_data: Optional[dict] = {}
+    node_ids: Optional[List[str]] = []
 
 
 class BaseSource(BaseModel):