From 089d38ec9b13ce97bf99493bf4a5cf89b3fefd0c Mon Sep 17 00:00:00 2001 From: Vitaly Neyman Date: Sun, 14 Apr 2024 14:35:31 +0300 Subject: [PATCH] Small refactor: improve ChatModelInfo computing and reporting short_doc as a property, in langchain_integration --- ps_fuzz/cli.py | 4 +--- ps_fuzz/langchain_integration.py | 4 ++++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/ps_fuzz/cli.py b/ps_fuzz/cli.py index 4897c05..e172579 100644 --- a/ps_fuzz/cli.py +++ b/ps_fuzz/cli.py @@ -53,9 +53,7 @@ def main(): if args.list_providers: print("Available providers:") for provider_name, provider_info in get_langchain_chat_models_info().items(): - doc = provider_info.doc - short_doc = doc[:doc.find('\n')] - print(f" {BRIGHT}{provider_name}{RESET}: {short_doc}") + print(f" {BRIGHT}{provider_name}{RESET}: {provider_info.short_doc}") sys.exit(0) if args.system_prompt_file is None: diff --git a/ps_fuzz/langchain_integration.py b/ps_fuzz/langchain_integration.py index c16d4a2..78e021e 100644 --- a/ps_fuzz/langchain_integration.py +++ b/ps_fuzz/langchain_integration.py @@ -76,6 +76,10 @@ def __str__(self): s += f" {param_name}: {param}\n" return s + @property + def short_doc(self): + return self.doc[:self.doc.find('\n')] + def get_langchain_chat_models_info() -> Dict[str, Dict[str, Any]]: """ Introspects a langchain library, extracting information about supported chat models and required/optional parameters