Using pydantic for llm_gateway models

Author: ritaagarwala-sf

pyproject.toml

@@ -99,7 +99,7 @@ click = "^8.1.8"
 loguru = "^0.7.3"
 numpy = "*"
 pandas = "*"
-pydantic = "^1.8.2 || ^2.0.0"
+pydantic = "2.13.1"
 pyspark = "3.5.1"
 python = ">=3.10,<3.12"
 pyyaml = "^6.0"

src/datacustomcode/llm_gateway/default.py

@@ -17,6 +17,7 @@
 from datacustomcode.llm_gateway.types.generate_text_request import GenerateTextRequest
 from datacustomcode.llm_gateway.types.generate_text_response import GenerateTextResponse
 
+from datacustomcode.llm_gateway.types.generate_text_response_builder import GenerateTextResponseBuilder
 
 class DefaultLLMGateway(LLMGateway):
     def generate_text(
@@ -26,7 +27,11 @@ def generate_text(
 
 
         response_data = {
-            'generation' : {'generatedText' : "I am dreaming!!"},
+            'version': 'v1',
+            'status_code': 200,
+            'data' : {
+                'generation': {'generatedText': 'Hello World'}
+            }
         }
 
-        return GenerateTextResponse(200, {"data": response_data})
+        return GenerateTextResponseBuilder.build(response_data)

src/datacustomcode/llm_gateway/types/generate_text_request.py

@@ -1,17 +1,12 @@
-from dataclasses import dataclass
-
-import betterproto
-
-from .google import protobuf
-
-
-@dataclass
-class GenerateTextRequest(betterproto.Message):
-    version: str = betterproto.string_field(1)
-    model_name: str = betterproto.string_field(2)
-    prompt: str = betterproto.string_field(3)
-    localization: protobuf.Struct = betterproto.message_field(4)
-    tags: protobuf.Struct = betterproto.message_field(5)
+from typing import Optional, Dict, Any, Literal
+from pydantic import BaseModel, Field
 
 
+class GenerateTextRequest(BaseModel):
+    """Request for LLM text generation"""
 
+    version: Literal["v1"] = Field(default="v1", description="API version, must be 'v1'")
+    model_name: str = Field(..., min_length=1, description="Name of the model to use")
+    prompt: str = Field(..., min_length=1, max_length=1000, description="Input prompt")
+    localization: Optional[Dict[str, Any]] = Field(default=None, description="Localization settings")
+    tags: Optional[Dict[str, Any]] = Field(default=None, description="Additional tags")

src/datacustomcode/llm_gateway/types/generate_text_request_builder.py

@@ -1,36 +1,12 @@
-from datacustomcode.validator.base import Validator
 from datacustomcode.llm_gateway.types.generate_text_request import GenerateTextRequest
 
 
-class GenerateTextRequestValidator:
-    @staticmethod
-    def create_validator() -> Validator:
-        """Create a validator with all CEL rules for GenerateTextRequest"""
-        validator = Validator()
-
-        # Rule 1: version == "v1" (CEL: const)
-        validator.add_rule(
-            id="request.version_v1",
-            message="Platform currently only supports version 'v1'",
-            expression=lambda this: this.version == "v1"
-        )
-
-        # Rule 2: modelName.size() >= 1 (CEL: min_len)
-        validator.add_rule(
-            id="request.model_name_required",
-            message="modelName must not be empty (min_len: 1)",
-            expression=lambda this: len(this.model_name) >= 1
-        )
-
-        return validator
-
 class GenerateTextRequestBuilder:
     def __init__(self):
-        self._validator = GenerateTextRequestValidator.create_validator()
-        self._version = "v1" # Hardcoded default for your SDK
         self._prompt = ""
         self._model_name = ""
-
+        self._localization = None
+        self._tags = None
 
     def set_prompt(self, prompt: str):
         self._prompt = prompt
@@ -40,18 +16,39 @@ def set_model(self, model_name: str):
         self._model_name = model_name
         return self
 
+    def set_localization(self, localization: dict = None, locale: str = None):
+        """
+        Set localization either from a dict or a simple locale string.
+
+        Args:
+          localization: Full localization dict (if provided, locale is ignored)
+          locale: Simple locale string for defaultLocale only
+
+        Returns:
+            self for method chaining
+        """
+
+        if localization is not None:
+            self._localization = localization
+        elif locale is not None:
+            self._localization = {"defaultLocale": locale}
+        else:
+            raise ValueError("Must provide either localization or locale")
+
+        self._localization = localization
+        return self
+
+    def set_tags(self, tags: dict):
+        self._tags = tags
+        return self
+
     def build(self) -> GenerateTextRequest:
 
         request = GenerateTextRequest(
-            version=self._version,
             prompt=self._prompt,
-            model_name=self._model_name
+            model_name=self._model_name,
+            localization=self._localization,
+            tags=self._tags
         )
 
-        # 2. Run the Protovalidate check
-        # This reads the 'max_len: 1000' rule from the .proto metadata
-        violations = self._validator.validate(request)
-        if violations:
-            raise ValueError(f"Validation Error: {violations}")
-
-        return request
+        return request

src/datacustomcode/llm_gateway/types/generate_text_response.py

@@ -1,14 +1,34 @@
+from typing import Optional, Dict, Any
 
-from dataclasses import dataclass
+from pydantic import BaseModel, Field
 
-import betterproto
+class GenerateTextResponse(BaseModel):
+    """Response from LLM text generation"""
 
-from .google import protobuf
+    version: str = Field(default="v1", description="API version")
+    status_code: int = Field(..., description="HTTP status code", ge=0)
+    data: Optional[Dict[str, Any]] = Field(default=None, description="Response data")
 
+    @property
+    def is_success(self) -> bool:
+        """Check if request succeeded."""
+        return self.status_code == 200
 
+    @property
+    def is_error(self) -> bool:
+        """Check if request failed."""
+        return not self.is_success
 
-@dataclass
-class GenerateTextResponse(betterproto.Message):
-    version: str = betterproto.string_field(1)
-    status_code: int = betterproto.uint32_field(2)
-    data: protobuf.Struct = betterproto.message_field(3)
+    @property
+    def text(self) -> str:
+        """Generated text (convenience property)."""
+        if self.is_success:
+            return self.data.get('generation', {}).get('generatedText', '')
+        return ''
+
+    @property
+    def error_code(self) -> str:
+        """Generated text (convenience property)."""
+        if self.is_error:
+            return self.data.get('errorCode', self.status_code)
+        return ''

src/datacustomcode/llm_gateway/types/generate_text_response_builder.py

@@ -1,55 +1,21 @@
-from .google import protobuf
+from typing import Dict, Any
+from datacustomcode.llm_gateway.types.generate_text_response import GenerateTextResponse
 
 
 class GenerateTextResponseBuilder:
     def __init__(self):
-        self._validator = Validator()
-        self._rules = validate_pb2.MessageConstraints()
-
-        # Rule 1: Prompt Length
-        prompt_rule = self.rules.cel.add()
-        prompt_rule.id = "request.prompt_limit"
-        prompt_rule.message = "Prompt must be 1-1000 characters."
-        prompt_rule.expression = "this.prompt.size() > 0 && this.prompt.size() <= 1000"
-
-        # Rule 3: ModelName Constraint
-        model_name_rule = self.rules.cel.add()
-        version_rule.id = "request.version_v1"
-        version_rule.message = "Platform currently only supports version 'v1'."
-        version_rule.expression = "this.version == 'v1'"
-
         self._version = "v1" # Hardcoded default for your SDK
-        self._prompt = ""
-        self._model_name = ""
-
+        self._status_code = None
+        self._data = None
 
-    def validate(self, request: GenerateTextRequest):
-        violations = self.validator.validate(request, constraints=self.rules)
-        if violations:
-            # protovalidate returns a structured 'Violations' object
-            error_msg = "; ".join([v.message for v in violations.violations])
-            raise ValueError(f"GenerateTextRequest Validation Failed: {error_msg}")
-
-    def set_prompt(self, prompt: str):
-        self._prompt = prompt
+    def set_status_code(self, status_code: int):
+        self._status_code = status_code
         return self
 
-    def set_model(self, model_name: str):
-        self._model_name = model_name
+    def set_data(self, data: dict):
+        self._data = data
         return self
 
-    def build(self) -> GenerateTextRequest:
-
-        request = GenerateTextRequest(
-            version=self._version,
-            prompt=self._prompt,
-            model_name=self._model_name
-        )
-
-        # 2. Run the Protovalidate check
-        # This reads the 'max_len: 1000' rule from the .proto metadata
-        violations = _validator.validate(request)
-        if violations:
-            raise ValueError(f"Validation Error: {violations}")
-
-        return request
+    @staticmethod
+    def build(response_dict: Dict[str, Any]) -> GenerateTextResponse:
+        return GenerateTextResponse.model_validate(response_dict)

src/datacustomcode/templates/function/payload/entrypoint.py

@@ -48,17 +48,13 @@ def function(request: dict, runTime: Runtime) -> dict:
 
 
     builder = GenerateTextRequestBuilder()
-    request = builder.set_prompt("Hello").set_model("gpt-4").build()
-    response = runTime.llm_gateway.generate_text(request)
+    llm_request = builder.set_prompt("Hello").set_model("").build()
+    llm_response = runTime.llm_gateway.generate_text(llm_request)
 
-    if response.is_success:
-        print(response.text)
+    if llm_response.is_success:
+        print(llm_response.text)
     else:
-        print(response.error_code)
-
-    file_path = runTime.file.find_file_path("data.csv")
-    content = open(file_path, 'r').read()
-    logger.info(content)
+        print(llm_response.error_code)
 
     for item in items:
         # Item is DocElement as dict
@@ -126,7 +122,7 @@ def function(request: dict, runTime: Runtime) -> dict:
     }
 
     # Run the function
-    result = function(test_request)
+    result = function(test_request, Runtime())
 
     # Print the results in a more readable format
     print("\nChunking Results:")