diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py index ebfdf7b4a..555b73785 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/magics.py @@ -16,6 +16,8 @@ from langchain.chains import LLMChain from langchain.schema import HumanMessage from langchain_core.messages import AIMessage +from py_markdown_table.markdown_table import markdown_table +from py_markdown_table.utils import find_longest_contiguous_strings from ._version import __version__ from .parsers import ( @@ -195,7 +197,7 @@ def _ai_bulleted_list_models_for_provider(self, provider_id, Provider): return output def _ai_inline_list_models_for_provider(self, provider_id, Provider): - output = "" + return output # Is the required environment variable set? def _ai_env_status_for_provider_markdown(self, provider_id): + result = {"env_var": "Not applicable", "emoji": "NA"} na_message = "Not applicable. | " + NA_MESSAGE if ( provider_id not in self.providers or self.providers[provider_id].auth_strategy == None ): - return na_message # No emoji + return result # No emoji not_set_title = ENV_NOT_SET set_title = ENV_SET @@ -227,6 +232,7 @@ def _ai_env_status_for_provider_markdown(self, provider_id): var_name = auth_strategy.name env_var_display = f"`{var_name}`" env_status_ok = var_name in os.environ + result["env_var"] = var_name elif auth_strategy.type == "multienv": # Check multiple environment variables var_names = self.providers[provider_id].auth_strategy.names @@ -235,16 +241,19 @@ def _ai_env_status_for_provider_markdown(self, provider_id): env_status_ok = all(var_name in os.environ for var_name in var_names) not_set_title = MULTIENV_NOT_SET set_title = MULTIENV_SET + result["env_var"] = env_var_display else: # No environment variables - return na_message + return result output = f"{env_var_display} | " if env_status_ok: output += f'' + result["emoji"] = "✅" else: output += f'' + result["emoji"] = "❌" - return output + return result def _ai_env_status_for_provider_text(self, provider_id): # only handle providers with "env" or "multienv" auth strategy @@ -355,22 +364,28 @@ def _ai_list_command_markdown(self, single_provider=None): "| Provider | Environment variable | Set? | Models |\n" + "|----------|----------------------|------|--------|\n" ) + + results = [] if single_provider is not None and single_provider not in self.providers: return f"There is no model provider with ID `{single_provider}`." for provider_id, Provider in self.providers.items(): if single_provider is not None and provider_id != single_provider: continue - - output += ( - f"| `{provider_id}` | " - + self._ai_env_status_for_provider_markdown(provider_id) - + " | " - + self._ai_inline_list_models_for_provider(provider_id, Provider) - + " |\n" + env_var_result = self._ai_env_status_for_provider_markdown(provider_id) + results.append( + { + "Provider": provider_id, + "Environment variable": env_var_result["env_var"], + "Set?": env_var_result["emoji"], + "Models": self._ai_inline_list_models_for_provider( + provider_id, Provider + ), + } ) # Also list aliases. + aliasOutput = [] if single_provider is None and len(self.custom_model_registry) > 0: output += ( "\nAliases and custom commands:\n\n" @@ -378,6 +393,7 @@ def _ai_list_command_markdown(self, single_provider=None): + "|------|--------|\n" ) for key, value in self.custom_model_registry.items(): + aliasOutput.append({"Name": key, "Target": value}) output += f"| `{key}` | " if isinstance(value, str): output += f"`{value}`" @@ -386,7 +402,30 @@ def _ai_list_command_markdown(self, single_provider=None): output += " |\n" - return output + # markdown = markdown_table(results).get_markdown() + longest = find_longest_contiguous_strings(results, True) + print(longest) + markdown = ( + markdown_table(results) + .set_params( + padding_width=2, + padding_weight="centerleft", + quote=False, + multiline={ + "Provider": 35, + "Environment variable": 24, + "Set?": 8, + "Models": 119, + }, + ) + .get_markdown() + ) + # print(markdown) + # markdown1 = markdown_table(aliasOutput).get_markdown() + # print(markdown1) + # print(results) + # print(aliasOutput) + return markdown def _ai_list_command_text(self, single_provider=None): output = "" diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index 66299a7a5..41b5ff2a1 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/packages/jupyter-ai-magics/pyproject.toml @@ -30,6 +30,7 @@ dependencies = [ "typing_extensions>=4.5.0", "click~=8.0", "jsonpath-ng>=1.5.3,<2", + "py_markdown_table>=1.3.0" ] [project.optional-dependencies]