cleanup

Signed-off-by: Mandana Vaziri <[email protected]>

Author: vazirim

src/pdl/pdl_context.py

@@ -1,16 +1,19 @@
-from typing import Any
 from enum import StrEnum
+from typing import Any
+
 from .pdl_lazy import (
     PdlDict,
     PdlList,
-    PdlLazy,
 )
 
+
 class SerializeMode(StrEnum):
     LITELLM = "litellm"
     GRANITEIO = "graniteio"
 
+
 class PDLContext():
+
     def serialize(self, mode: SerializeMode) -> list[dict[str, Any]]:
         return []
 
@@ -30,14 +33,12 @@ def serialize(self, mode: SerializeMode) -> list[dict[str, Any]]:
         return [result]
 
 
-
 class IndependentContext(PDLContext):
     context: PdlList[PDLContext]
 
     def __init__(self, context: PdlList[PDLContext]):
         self.context = context
 
-
     def serialize(self, mode: SerializeMode) -> list[dict[str, Any]]:
         result = self.context.result()
         contexts = [m.serialize(mode) for m in result]
@@ -57,9 +58,9 @@ def serialize(self, mode: SerializeMode) -> list[dict[str, Any]]:
         result = self.context.result()
         contexts = [m.serialize(mode) for m in result]
         return [x for xs in contexts for x in xs]
-    
 
-def deserialize(context: list[dict[str, Any]]) -> DependentContext: # Only support dependent for now 
+
+def deserialize(context: list[dict[str, Any]]) -> DependentContext:  # Only support dependent for now
     ret: DependentContext = DependentContext(PdlList([]))
     for message in context:
         if isinstance(message, dict):
@@ -71,4 +72,4 @@ def deserialize(context: list[dict[str, Any]]) -> DependentContext: # Only suppo
         else:
             ret = DependentContext(PdlList([ret, message]))
 
-    return ret
+    return ret

tests/test_context.py

@@ -1,4 +1,9 @@
-from pdl.pdl_context import DependentContext, BaseMessage, IndependentContext, SerializeMode
+from pdl.pdl_context import (
+    BaseMessage,
+    DependentContext,
+    IndependentContext,
+    SerializeMode,
+)
 from pdl.pdl_lazy import PdlList
 
 a = BaseMessage({"role": "user", "content": "hello"})
@@ -17,6 +22,7 @@
 p = DependentContext(PdlList([d1, d2]))
 p1 = DependentContext(PdlList([d1, i]))
 
+
 def test_p():
     assert p.serialize(SerializeMode.LITELLM) == [{"role": "user", "content": "hello"},
                                                   {"role": "user", "content": "bye"},
@@ -27,11 +33,11 @@ def test_p():
 
 def test_p1():
     assert p1.serialize(SerializeMode.LITELLM) == [{"role": "user", "content": "hello"},
-                                                  {"role": "user", "content": "bye"},
-                                                  {"role": "user", "content": "hello2"},
-                                                  {"role": "user", "content": "bye2"}]
-    
+                                                   {"role": "user", "content": "bye"},
+                                                   {"role": "user", "content": "hello2"},
+                                                   {"role": "user", "content": "bye2"}]
+
     assert p1.serialize(SerializeMode.GRANITEIO) == [{"role": "user", "content": "hello"},
-                                                  {"role": "user", "content": "bye"},
-                                                  {"independent": [{"role": "user", "content": "hello2"},
-                                                  {"role": "user", "content": "bye2"}]}]
+                                                     {"role": "user", "content": "bye"},
+                                                     {"independent": [{"role": "user", "content": "hello2"},
+                                                                      {"role": "user", "content": "bye2"}]}]