combine datatypes.py and endpoints.py into api.py

llama_toolchain/safety/api/__init__.py

@@ -4,5 +4,4 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from .datatypes import *  # noqa
-from .endpoints import *  # noqa
+from .api import *  # noqa: F401 F403

llama_toolchain/safety/api/api.py

@@ -5,13 +5,12 @@
 # the root directory of this source tree.
 
 from enum import Enum
-from typing import Dict, Optional, Union
-
-from llama_models.llama3.api.datatypes import ToolParamDefinition
-from llama_models.schema_utils import json_schema_type
+from typing import Dict, List, Optional, Protocol, Union
 
+from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, validator
 
+from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_toolchain.common.deployment_types import RestAPIExecutionConfig
 
 
@@ -70,3 +69,22 @@ class ShieldResponse(BaseModel):
             except ValueError:
                 return v
         return v
+
+
+@json_schema_type
+class RunShieldRequest(BaseModel):
+    messages: List[Message]
+    shields: List[ShieldDefinition]
+
+
+@json_schema_type
+class RunShieldResponse(BaseModel):
+    responses: List[ShieldResponse]
+
+
+class Safety(Protocol):
+    @webmethod(route="/safety/run_shields")
+    async def run_shields(
+        self,
+        request: RunShieldRequest,
+    ) -> RunShieldResponse: ...

llama_toolchain/safety/api/endpoints.py

@@ -1,32 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from .datatypes import *  # noqa: F403
-from typing import List, Protocol
-
-from llama_models.llama3.api.datatypes import Message
-
-# this dependency is annoying and we need a forked up version anyway
-from llama_models.schema_utils import webmethod
-
-
-@json_schema_type
-class RunShieldRequest(BaseModel):
-    messages: List[Message]
-    shields: List[ShieldDefinition]
-
-
-@json_schema_type
-class RunShieldResponse(BaseModel):
-    responses: List[ShieldResponse]
-
-
-class Safety(Protocol):
-    @webmethod(route="/safety/run_shields")
-    async def run_shields(
-        self,
-        request: RunShieldRequest,
-    ) -> RunShieldResponse: ...