Add serialisation props to Fireworks and ChatFireworks (langchain-ai#…

…12255)
EricLiclair · Oct 25, 2023 · c25b174 · c25b174
1 parent fd5f549
commit c25b174
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/libs/langchain/langchain/chat_models/fireworks.py b/libs/langchain/langchain/chat_models/fireworks.py
@@ -90,6 +90,14 @@ class ChatFireworks(BaseChatModel):
     fireworks_api_key: Optional[str] = None
     max_retries: int = 20
 
+    @property
+    def lc_secrets(self) -> Dict[str, str]:
+        return {"fireworks_api_key": "FIREWORKS_API_KEY"}
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        return True
+
     @root_validator()
     def validate_environment(cls, values: Dict) -> Dict:
         """Validate that api key in environment."""

diff --git a/libs/langchain/langchain/llms/fireworks.py b/libs/langchain/langchain/llms/fireworks.py
@@ -37,6 +37,14 @@ class Fireworks(LLM):
     fireworks_api_key: Optional[str] = None
     max_retries: int = 20
 
+    @property
+    def lc_secrets(self) -> Dict[str, str]:
+        return {"fireworks_api_key": "FIREWORKS_API_KEY"}
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        return True
+
     @root_validator()
     def validate_environment(cls, values: Dict) -> Dict:
         """Validate that api key in environment."""