fix: Use ruamel.yaml instead of pyyaml

Signed-off-by: Diwank Singh Tomer <[email protected]>

agents-api/agents_api/activities/excecute_api_call.py

@@ -1,15 +1,13 @@
 import base64
-from typing import Annotated, Any, Optional, TypedDict, Union
+from typing import Any, Optional, TypedDict, Union
 
 import httpx
 from beartype import beartype
-from pydantic import Field
 from temporalio import activity
 
 from ..autogen.openapi_model import ApiCallDef
 
 # from ..clients import integrations
-from ..common.protocol.tasks import StepContext
 from ..env import testing
 
 # from ..models.tools import get_tool_args_from_metadata

agents-api/agents_api/activities/execute_system.py

@@ -19,9 +19,7 @@
 from ..models.agent.get_agent import get_agent as get_agent_query
 from ..models.agent.list_agents import list_agents as list_agents_query
 from ..models.agent.update_agent import update_agent as update_agent_query
-from ..models.docs.create_doc import create_doc as create_doc_query
 from ..models.docs.delete_doc import delete_doc as delete_doc_query
-from ..models.docs.get_doc import get_doc as get_doc_query
 from ..models.docs.list_docs import list_docs as list_docs_query
 from ..models.session.create_session import create_session as create_session_query
 from ..models.session.delete_session import delete_session as delete_session_query
@@ -121,7 +119,7 @@ async def execute_system(
                     )
 
             # NO SUBRESOURCE
-            elif system.subresource == None:
+            elif system.subresource is None:
                 if system.operation == "list":
                     return list_agents_query(**arguments)
                 elif system.operation == "get":
@@ -190,7 +188,7 @@ async def execute_system(
                     )
 
             # NO SUBRESOURCE
-            elif system.subresource == None:
+            elif system.subresource is None:
                 if system.operation == "list":
                     return list_users_query(**arguments)
                 elif system.operation == "get":

agents-api/agents_api/activities/task_steps/get_value_step.py

@@ -11,7 +11,7 @@
 async def get_value_step(
     context: StepContext,
 ) -> StepOutcome:
-    key: str = context.current_step.get
+    key: str = context.current_step.get  # noqa: F841
     raise NotImplementedError("Not implemented yet")
 
 

agents-api/agents_api/activities/utils.py

@@ -9,20 +9,18 @@
 import string
 import time
 import urllib.parse
-from typing import Any, Callable, ParamSpec, Type, TypeVar, cast
+from typing import Any, Callable, ParamSpec, TypeVar
 
 import re2
-import yaml
 import zoneinfo
 from beartype import beartype
 from simpleeval import EvalWithCompoundTypes, SimpleEval
-from yaml import CSafeDumper, CSafeLoader
-
-T = TypeVar("T")
 
+from ..common.utils import yaml
 
-P = ParamSpec("P")
+T = TypeVar("T")
 R = TypeVar("R")
+P = ParamSpec("P")
 
 
 # TODO: We need to make sure that we dont expose any security issues
@@ -51,7 +49,7 @@
     "zip": zip,
     "search_regex": lambda pattern, string: re2.search(pattern, string),
     "load_json": json.loads,
-    "load_yaml": lambda string: yaml.load(string, Loader=CSafeLoader),
+    "load_yaml": yaml.load,
     "match_regex": lambda pattern, string: bool(re2.fullmatch(pattern, string)),
 }
 
@@ -74,8 +72,8 @@ class stdlib_json:
 
 
 class stdlib_yaml:
-    load = lambda string: yaml.load(string, Loader=CSafeLoader)  # noqa: E731
-    dump = lambda value: yaml.dump(value, Dumper=CSafeDumper)  # noqa: E731
+    load = yaml.load
+    dump = yaml.dump
 
 
 class stdlib_time:

agents-api/agents_api/common/protocol/tasks.py

@@ -1,4 +1,4 @@
-from typing import Annotated, Any, Type
+from typing import Annotated, Any
 from uuid import UUID
 
 from pydantic import BaseModel, Field, computed_field

agents-api/agents_api/common/utils/template.py

@@ -4,12 +4,13 @@
 
 import arrow
 import re2
-import yaml
 from beartype import beartype
 from jinja2.sandbox import ImmutableSandboxedEnvironment
 from jinja2schema import infer, to_json_schema
 from jsonschema import validate
 
+from . import yaml
+
 __all__: List[str] = [
     "render_template",
 ]

agents-api/agents_api/common/utils/yaml.py

@@ -0,0 +1,18 @@
+from io import StringIO
+from typing import Any
+
+from ruamel.yaml import YAML
+
+yaml = YAML(typ="safe", pure=True)  # pure is needed for yaml 1.2 support
+yaml.version = (1, 2)
+
+
+def dump(value: Any) -> str:
+    stream = StringIO()
+    yaml.dump(value, stream)
+
+    return stream.getvalue()
+
+
+def load(string: str) -> Any:
+    return yaml.load(string)

agents-api/agents_api/routers/internal/__init__.py

@@ -1 +1 @@
-from .router import router
+from .router import router as router

agents-api/agents_api/routers/internal/router.py

@@ -1,4 +1,4 @@
-from fastapi import APIRouter, Request, Response
+from fastapi import APIRouter, Request
 from google.protobuf import json_format
 from temporalio.api.common.v1 import Payloads
 

agents-api/agents_api/routers/tasks/router.py

@@ -1,9 +1,10 @@
 from typing import Callable
 
-import yaml
 from fastapi import APIRouter, Request, Response
 from fastapi.routing import APIRoute
 
+from ...common.utils import yaml
+
 
 class YamlRequest(Request):
     async def body(self) -> bytes:

agents-api/agents_api/routers/tasks/update_execution.py

@@ -34,7 +34,7 @@ async def update_execution(
                     *get_temporal_workflow_data(execution_id=execution_id)
                 )
                 await wf_handle.cancel()
-            except Exception as e:
+            except Exception:
                 raise HTTPException(status_code=500, detail="Failed to stop execution")
 
         case ResumeExecutionRequest():
@@ -59,7 +59,7 @@ async def update_execution(
                 )
             try:
                 await act_handle.complete(data.input)
-            except Exception as e:
+            except Exception:
                 raise HTTPException(
                     status_code=500, detail="Failed to resume execution"
                 )