Fixing lint error

Author: ritaagarwala-sf

src/datacustomcode/llm_gateway/base.py

@@ -16,9 +16,13 @@
 
 from abc import abstractmethod
 
+from datacustomcode.llm_gateway.types.generate_text_request import GenerateTextRequest
+from datacustomcode.llm_gateway.types.generate_text_response import GenerateTextResponse
+
+
 class LLMGateway:
     def __init__(self):
         pass
 
     @abstractmethod
-    def generate_text(self, GenerateTextRequest) -> GenerateTextResponse: ...
+    def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse: ...

src/datacustomcode/llm_gateway/default.py

@@ -16,22 +16,18 @@
 from datacustomcode.llm_gateway.base import LLMGateway
 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,
+)
 
-from datacustomcode.llm_gateway.types.generate_text_response_builder import GenerateTextResponseBuilder
 
 class DefaultLLMGateway(LLMGateway):
-    def generate_text(
-            self,
-            request: GenerateTextRequest
-    ) -> GenerateTextResponse:
-
+    def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse:
 
         response_data = {
-            'version': 'v1',
-            'status_code': 200,
-            'data' : {
-                'generation': {'generatedText': 'Hello World'}
-            }
+            "version": "v1",
+            "status_code": 200,
+            "data": {"generation": {"generatedText": "Hello World"}},
         }
 
         return GenerateTextResponseBuilder.build(response_data)

src/datacustomcode/llm_gateway/types/generate_text_request.py

@@ -13,20 +13,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Dict, Any, Literal
-from pydantic import BaseModel, Field, ConfigDict
+from typing import (
+    Any,
+    Dict,
+    Literal,
+    Optional,
+)
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+)
 from pydantic.alias_generators import to_camel
 
 
 class GenerateTextRequest(BaseModel):
 
     model_config = ConfigDict(
         alias_generator=to_camel,
-        populate_by_name=True  # Allows both snake_case and camelCase input
+        populate_by_name=True,  # Allows both snake_case and camelCase input
     )
 
-    version: Literal["v1"] = Field(default="v1", description="API version, must be 'v1'")
+    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")
+    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

@@ -13,15 +13,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import (
+    Any,
+    Dict,
+    Optional,
+)
+
 from datacustomcode.llm_gateway.types.generate_text_request import GenerateTextRequest
 
 
 class GenerateTextRequestBuilder:
     def __init__(self):
         self._prompt = ""
         self._model_name = ""
-        self._localization = None
-        self._tags = None
+        self._localization: Optional[Dict[str, Any]] = None
+        self._tags: Optional[Dict[str, Any]] = None
 
     def set_prompt(self, prompt: str):
         self._prompt = prompt
@@ -31,7 +37,11 @@ def set_model(self, model_name: str):
         self._model_name = model_name
         return self
 
-    def set_localization(self, localization: dict = None, locale: str = None):
+    def set_localization(
+        self,
+        localization: Optional[Dict[str, Any]] = None,
+        locale: Optional[str] = None,
+    ):
         """
         Set localization either from a dict or a simple locale string.
 
@@ -52,7 +62,7 @@ def set_localization(self, localization: dict = None, locale: str = None):
 
         return self
 
-    def set_tags(self, tags: dict):
+    def set_tags(self, tags: Dict[str, Any]):
         self._tags = tags
         return self
 
@@ -62,7 +72,7 @@ def build(self) -> GenerateTextRequest:
             prompt=self._prompt,
             model_name=self._model_name,
             localization=self._localization,
-            tags=self._tags
+            tags=self._tags,
         )
 
         return request

src/datacustomcode/llm_gateway/types/generate_text_response.py

@@ -13,10 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Dict, Any
+from typing import (
+    Any,
+    Dict,
+    Optional,
+)
 
 from pydantic import BaseModel, Field
 
+
 class GenerateTextResponse(BaseModel):
     """Response from LLM text generation"""
 
@@ -38,12 +43,16 @@ def is_error(self) -> bool:
     def text(self) -> str:
         """Generated text (convenience property)."""
         if self.is_success and self.data:
-            return self.data.get('generation', {}).get('generatedText', '')
-        return ''
+            generation = self.data.get("generation", {})
+            if isinstance(generation, dict):
+                text = generation.get("generatedText", "")
+                return str(text) if text else ""
+        return ""
 
     @property
     def error_code(self) -> str:
         """Generated text (convenience property)."""
         if self.is_error and self.data:
-            return self.data.get('errorCode', self.status_code)
-        return ''
+            error_code = self.data.get("errorCode", str(self.status_code))
+            return str(error_code)
+        return ""

src/datacustomcode/llm_gateway/types/generate_text_response_builder.py

@@ -13,13 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Any
+from typing import Any, Dict
+
 from datacustomcode.llm_gateway.types.generate_text_response import GenerateTextResponse
 
 
 class GenerateTextResponseBuilder:
     def __init__(self):
-        self._version = "v1" # Hardcoded default for your SDK
+        self._version = "v1"  # Hardcoded default for your SDK
         self._status_code = None
         self._data = None
 
@@ -33,4 +34,4 @@ def set_data(self, data: dict):
 
     @staticmethod
     def build(response_dict: Dict[str, Any]) -> GenerateTextResponse:
-        return GenerateTextResponse.model_validate(response_dict)
+        return GenerateTextResponse.model_validate(response_dict)

src/datacustomcode/proxy/__init__.py

@@ -11,4 +11,4 @@
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
-# limitations under the License.
+# limitations under the License.

src/datacustomcode/proxy/base.py

@@ -21,4 +21,4 @@
 
 class BaseProxyAccessLayer(ABC, UserExtendableNamedConfigMixin):
     def __init__(self):
-        pass
+        pass

src/datacustomcode/proxy/client/LocalProxyClientProvider.py

@@ -31,4 +31,4 @@ def call_llm_gateway(self, llmModelId: str, prompt: str, maxTokens: int) -> str:
     def llm_gateway_generate_text(
         self, template, values, llmModelId: str, maxTokens: int
     ):
-        return f"Using Generate Text with {llmModelId} and maxTokens: {maxTokens}"
+        return f"Using Generate Text with {llmModelId} and maxTokens: {maxTokens}"

src/datacustomcode/proxy/client/__init__.py

@@ -11,4 +11,4 @@
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
-# limitations under the License.
+# limitations under the License.