fix: update all enum classes to support python 3.11+ users

lanarky/adapters/langchain/callbacks.py

@@ -1,4 +1,3 @@
-from enum import Enum
 from typing import Any, Optional
 
 from fastapi.websockets import WebSocket
@@ -12,10 +11,10 @@
 from starlette.types import Message, Send
 
 from lanarky.events import Events, ServerSentEvent, ensure_bytes
-from lanarky.utils import model_dump_json
+from lanarky.utils import StrEnum, model_dump_json
 
 
-class LangchainEvents(str, Enum):
+class LangchainEvents(StrEnum):
     SOURCE_DOCUMENTS = "source_documents"
 
 
@@ -78,7 +77,7 @@ def _construct_message(self, data: str, event: Optional[str] = None) -> Message:
         }
 
 
-class TokenStreamMode(str, Enum):
+class TokenStreamMode(StrEnum):
     TEXT = "text"
     JSON = "json"
 

lanarky/adapters/langchain/responses.py

@@ -1,5 +1,4 @@
 import asyncio
-from enum import Enum
 from functools import partial
 from typing import Any
 
@@ -11,9 +10,10 @@
 from lanarky.logging import logger
 from lanarky.responses import HTTPStatusDetail
 from lanarky.responses import StreamingResponse as _StreamingResponse
+from lanarky.utils import StrEnum
 
 
-class ChainRunMode(str, Enum):
+class ChainRunMode(StrEnum):
     """Enum for LangChain run modes."""
 
     ASYNC = "async"

lanarky/events.py

@@ -1,10 +1,10 @@
-from enum import Enum
-
 from sse_starlette.sse import ServerSentEvent as ServerSentEvent
 from sse_starlette.sse import ensure_bytes as ensure_bytes
 
+from lanarky.utils import StrEnum
+
 
-class Events(str, Enum):
+class Events(StrEnum):
     COMPLETION = "completion"
     ERROR = "error"
     END = "end"

lanarky/responses.py

@@ -1,4 +1,3 @@
-from enum import Enum
 from typing import Any
 
 from fastapi import status
@@ -7,9 +6,10 @@
 
 from lanarky.events import Events, ServerSentEvent, ensure_bytes
 from lanarky.logging import logger
+from lanarky.utils import StrEnum
 
 
-class HTTPStatusDetail(str, Enum):
+class HTTPStatusDetail(StrEnum):
     INTERNAL_SERVER_ERROR = "Internal Server Error"
 
 

lanarky/utils.py

@@ -3,6 +3,15 @@
 import pydantic
 from pydantic.fields import FieldInfo
 
+try:
+    from enum import StrEnum  # type: ignore
+except ImportError:
+    from enum import Enum
+
+    class StrEnum(str, Enum):
+        ...
+
+
 PYDANTIC_V2 = pydantic.VERSION.startswith("2.")
 
 

lanarky/websockets.py

@@ -1,13 +1,13 @@
 from contextlib import asynccontextmanager
-from enum import Enum
 from typing import Generator
 
 from fastapi.websockets import WebSocket, WebSocketDisconnect
 
 from lanarky.logging import logger
+from lanarky.utils import StrEnum
 
 
-class DataMode(str, Enum):
+class DataMode(StrEnum):
     JSON = "json"
     TEXT = "text"
     BYTES = "bytes"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "lanarky"
-version = "0.8.2"
+version = "0.8.3"
 description = "The web framework for building LLM microservices"
 authors = ["Ajinkya Indulkar <ajndkr@gmail.com>"]
 readme = "README.pypi.md"