From 2c720f83a857f3fe102d42d9193e48c8b46da55c Mon Sep 17 00:00:00 2001 From: Raju Komati Date: Mon, 24 Apr 2023 14:34:14 +0530 Subject: [PATCH] added validation for model parameter --- quora/__init__.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/quora/__init__.py b/quora/__init__.py index 88c9d0c..4eeeb69 100644 --- a/quora/__init__.py +++ b/quora/__init__.py @@ -380,8 +380,13 @@ class Completion: class Poe: def __init__(self, model: str = "ChatGPT"): - self.cookie = self.__load_cookie() + # validating the model + if model and model not in MODELS: + raise RuntimeError( + "Sorry, the model you provided does not exist. Please check and try again." + ) self.model = MODELS[model] + self.cookie = self.__load_cookie() self.client = PoeClient(self.cookie) def __load_cookie(self) -> str: @@ -443,7 +448,11 @@ class Poe: return cookie def chat(self, message: str, model: Optional[str] = None) -> str: - model = MODELS[model] or self.model + if model and model not in MODELS: + raise RuntimeError( + "Sorry, the model you provided does not exist. Please check and try again." + ) + model = MODELS[model] if model else self.model response = None for chunk in self.client.send_message(model, message): response = chunk["text"]