feat: Upgrade ZenML

llm_engineering/application/dataset/generation.py

@@ -38,7 +38,7 @@ class DatasetGenerator:
     @classmethod
     def get_system_prompt(cls) -> Prompt:
         return Prompt(
-            template=PromptTemplate.from_template(cls.system_prompt_template),
+            template=cls.system_prompt_template,
             input_variables={},
             content=cls.system_prompt_template,
         )
@@ -93,7 +93,7 @@ def get_prompt(cls, documents: list[CleanedDocument]) -> GenerateDatasetSamplesP
             prompt = cls.tokenizer.decode(prompt_tokens)
 
         prompt = GenerateDatasetSamplesPrompt(
-            template=prompt_template,
+            template=prompt_template.template,
             input_variables=input_variables,
             content=prompt,
             num_tokens=len(prompt_tokens),

llm_engineering/domain/base/nosql.py

@@ -44,7 +44,7 @@ def to_mongo(self: T, **kwargs) -> dict:
         exclude_unset = kwargs.pop("exclude_unset", False)
         by_alias = kwargs.pop("by_alias", True)
 
-        parsed = self.dict(exclude_unset=exclude_unset, by_alias=by_alias, **kwargs)
+        parsed = self.model_dump(exclude_unset=exclude_unset, by_alias=by_alias, **kwargs)
 
         if "_id" not in parsed and "id" in parsed:
             parsed["_id"] = str(parsed.pop("id"))
@@ -55,8 +55,8 @@ def to_mongo(self: T, **kwargs) -> dict:
 
         return parsed
 
-    def dict(self: T, **kwargs) -> dict:
-        dict_ = super().dict(**kwargs)
+    def model_dump(self: T, **kwargs) -> dict:
+        dict_ = super().model_dump(**kwargs)
 
         for key, value in dict_.items():
             if isinstance(value, uuid.UUID):

llm_engineering/domain/base/vector.py

@@ -48,7 +48,7 @@ def to_point(self: T, **kwargs) -> PointStruct:
         exclude_unset = kwargs.pop("exclude_unset", False)
         by_alias = kwargs.pop("by_alias", True)
 
-        payload = self.dict(exclude_unset=exclude_unset, by_alias=by_alias, **kwargs)
+        payload = self.model_dump(exclude_unset=exclude_unset, by_alias=by_alias, **kwargs)
 
         _id = str(payload.pop("id"))
         vector = payload.pop("embedding", {})
@@ -57,8 +57,8 @@ def to_point(self: T, **kwargs) -> PointStruct:
 
         return PointStruct(id=_id, vector=vector, payload=payload)
 
-    def dict(self: T, **kwargs) -> dict:
-        dict_ = super().dict(**kwargs)
+    def model_dump(self: T, **kwargs) -> dict:
+        dict_ = super().model_dump(**kwargs)
 
         for key, value in dict_.items():
             if isinstance(value, UUID):

llm_engineering/domain/prompt.py

@@ -1,18 +1,17 @@
-from langchain_core.prompts import PromptTemplate
-
 from llm_engineering.domain.base import VectorBaseDocument
 from llm_engineering.domain.cleaned_documents import CleanedDocument
 from llm_engineering.domain.types import DataCategory
 
 
 class Prompt(VectorBaseDocument):
-    template: PromptTemplate
+    template: str
     input_variables: dict
     content: str
     num_tokens: int | None = None
 
     class Config:
         category = DataCategory.PROMPT
+        arbitrary_types_allowed = True
 
 
 class GenerateDatasetSamplesPrompt(Prompt):

llm_engineering/settings.py

@@ -1,4 +1,4 @@
-from pydantic import BaseSettings
+from pydantic_settings import BaseSettings
 
 
 class Settings(BaseSettings):

pyproject.toml

@@ -8,9 +8,10 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "~3.11"
-zenml = { version = "0.58.1", extras = ["server"] }
+zenml = { version = "0.63.0", extras = ["server"] }
 pymongo = "^4.6.2"
-pydantic = "^1.9.0"
+pydantic = "^2.7.0"
+pydantic-settings = "^2.2.1"
 click = "^8.0.1"
 loguru = "^0.7.2"
 rich = "^13.7.1"
@@ -25,11 +26,11 @@ jmespath = "^1.0.1"
 # Feature engineering
 qdrant-client = "^1.8.0"
 unstructured = "^0.12.0"
-langchain = "^0.1.13"
+langchain = "^0.2.11"
 sentence-transformers = "^2.6.1"
 
 # RAG
-langchain-openai = "^0.1.3"
+langchain-openai = "^0.1.19"
 jinja2 = "^3.1.4"
 tiktoken = "^0.7.0"
 fake-useragent = "^1.5.1"