diff --git a/doc_craft/core/constant.py b/doc_craft/core/constant.py new file mode 100644 index 0000000..e66ab57 --- /dev/null +++ b/doc_craft/core/constant.py @@ -0,0 +1 @@ +AI_MODEL = {"model_name": "gpt-3.5-turbo", "model_family": "openai"} diff --git a/doc_craft/utils/llm.py b/doc_craft/utils/llm.py index cf3c8c4..2af6802 100644 --- a/doc_craft/utils/llm.py +++ b/doc_craft/utils/llm.py @@ -3,13 +3,18 @@ from .env_apikey_handler import EnvAPIKeyHandler from .azurekeyvault_apikey_handler import AzureKeyVaultAPIKeyHandler import sys +from doc_craft.core.constant import AI_MODEL class LLM: """A class that represents a Language Model (LLM).""" # TODO: allow user to specify model_name and model_family - def __init__(self, model_name="gpt-3.5-turbo", model_family="openai"): + def __init__( + self, + model_name=AI_MODEL.get("model_name"), + model_family=AI_MODEL.get("model_family"), + ): """Initialize a Model object with the specified model name and model family. Args: diff --git a/setup.py b/setup.py index 8caa435..3ad6017 100644 --- a/setup.py +++ b/setup.py @@ -1,8 +1,8 @@ from setuptools import setup, find_packages setup( - name="DocCraft", - version="1.0.3", + name="doc-craft", + version="1.0.5", description="A tool to add docstrings to Python code using LLMs.", long_description=open("README.md").read(), long_description_content_type="text/markdown",