fix comment

Chainlit · Oct 2, 2024 · a854002 · a854002
1 parent 586bed0
commit a854002
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/literalai/observability/generation.py b/literalai/observability/generation.py
@@ -13,10 +13,19 @@
 
 
 @unique
-class GenerationType(Enum):
+class GenerationType(str, Enum):
     CHAT = "CHAT"
     COMPLETION = "COMPLETION"
 
+    def __str__(self):
+        return self.value
+
+    def __repr__(self):
+        return f"GenerationType.{self.name}"
+
+    def to_json(self):
+        return self.value
+
 
 class GenerationMessage(TypedDict, total=False):
     uuid: Optional[str]

diff --git a/literalai/prompt_engineering/prompt.py b/literalai/prompt_engineering/prompt.py
@@ -79,12 +79,11 @@ class Prompt(Utils):
     variables_default_values: Optional[Dict[str, Any]]
 
     def to_dict(self) -> PromptDict:
-        # Convert the Prompt instance to a dictionary matching the PromptDict structure
         return {
             "id": self.id,
             "createdAt": self.created_at,
             "updatedAt": self.updated_at,
-            "type": self.type.value if isinstance(self.type, GenerationType) else self.type,
+            "type": self.type,
             "name": self.name,
             "version": self.version,
             "url": self.url,