diff --git a/README.md b/README.md index 929a0ae..631bd1b 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,17 @@ Simply read a model and it's context to explain what the model actually does, an dbtai explain ``` +### Chat +You can open an interactive chat with a dbt model: + +```bash +dbtai chat +``` + +This will open a CLI chat, letting you ask questions and get answers interactively, keeping the chat history. + +Save the chat history to file by typing `\save` inside the chat. You can still continue the chat after saving. + ## That's all, folks! diff --git a/dbtai/chatbot.py b/dbtai/chatbot.py new file mode 100644 index 0000000..0eed6a6 --- /dev/null +++ b/dbtai/chatbot.py @@ -0,0 +1,73 @@ +import openai +import os +import appdirs +import yaml +import click +import datetime + +class ModelChatBot: + def __init__( + self, + model_name, + system_prompt + ): + + self.config = self._load_config() + self.model_name = model_name + self.chat_history = [ + {"role": "system", "content": system_prompt} + ] + + self.client = openai.OpenAI(api_key=self.config["api_key"] or os.getenv("OPENAI_API_KEY")) + + def _load_config(self): + """Convenience function to load the user config from the config file.""" + configdir = appdirs.user_data_dir("dbtai", "dbtai") + + if os.path.exists(os.path.join(configdir, "config.yaml")): + with open(os.path.join(configdir, "config.yaml"), "r") as f: + return yaml.load(f, Loader=yaml.FullLoader) + + return {'language': 'english', "backend": "OpenAI"} + + def chat_completion(self, messages): + """Convenience method to call the chat completion endpoint. + + Args: + messages (list): A list of messages to send to the chat API + response_format_type (str, optional): The response format. Defaults to "json_object". + + Returns: + openai.ChatCompletion: The response from the chat API + """ + + return self.client.chat.completions.create( + model="gpt-3.5-turbo", + messages=messages + ) + + def run(self): + print(f""" +Hi! I'm here to chat about the dbt model {self.model_name}. What's on your mind? +Type 'quit' or hit Ctrl-C to exit) + """ + ) + while True: + user_input = input(">>> ") + if user_input.lower() in ['quit', 'exit']: + print("Goodbye!") + break + + if user_input == "\save": + with open('chat_history.txt', 'a') as f: + f.write(f"Chat history for the dbt model: {self.model_name}, on {datetime.datetime.now().isoformat()}\n\n") + for item in self.chat_history: + f.write("%s\n" % item) + click.echo(click.style(response.choices[0].message.content, fg='blue')) + print("Chat history saved to chat_history.txt") + continue + + self.chat_history.append({"role": "user", "content": user_input}) + response = self.chat_completion(self.chat_history) + click.echo(click.style(response.choices[0].message.content, fg='blue')) + self.chat_history.append({"role": "system", "content": response.choices[0].message.content}) diff --git a/dbtai/cli.py b/dbtai/cli.py index 7d2bd12..49bf48e 100644 --- a/dbtai/cli.py +++ b/dbtai/cli.py @@ -6,6 +6,7 @@ import yaml from dbtai.templates.prompts import languages, GENERATE_MODEL from dbtai.manifest import Manifest +from dbtai.chatbot import ModelChatBot APPNAME = "dbtai" APPAUTHOR = "dbtai" @@ -183,6 +184,22 @@ def explain(model): result = manifest.explain(model) click.echo(result) + +@dbtai.command(help="Chat with a dbt model") +@click.argument("model", required=True) +def chat(model): + manifest = Manifest() + chatbot_prompt = manifest.generate_chatbot_prompt(model) + chatbot = ModelChatBot( + model_name=model, + system_prompt=chatbot_prompt + ) + chatbot.run() + + + + + @dbtai.command(help="Show logo") def hello(): greeting = r""" diff --git a/dbtai/manifest.py b/dbtai/manifest.py index 370683b..078534d 100644 --- a/dbtai/manifest.py +++ b/dbtai/manifest.py @@ -148,6 +148,26 @@ def compile_upstream_description_markdown(self, model_name): return '\n\n'.join(all_models) + def get_model_description(self, model_name): + """Get the description of a model. + + Args: + model_name (str): The name of the model. + + Returns: + str: The description of the model. + """ + model = self.get_model_from_name(model_name) + + top_level_description = model.get('description') or "(no description)" + column_descriptions = '\n'.join( + [f'* {column}: {content.get("description") or "(no description)"}' + for column, content in model['columns'].items()] + ) or '(no columns defined)' + + return f'{top_level_description}\nColumns:\n{column_descriptions}' + + def create_documentation_instructions(self, model_name): """Create a markdown prompt with instructions for documenting the model. @@ -384,4 +404,16 @@ def explain(self, model_name): response_format_type="text" ) - return response.choices[0].message.content \ No newline at end of file + return response.choices[0].message.content + + def generate_chatbot_prompt(self, model): + model_docs = self.get_model_description(model) + upstream_docs = self.compile_upstream_description_markdown(model) + model_code = self.get_model_from_name(model)['raw_code'] + prompt = languages[self.config['language']]['chatbot_prompt'].format( + model_name=model, + raw_code=model_code, + upstream_models = upstream_docs, + model_description=model_docs + ) + return prompt \ No newline at end of file diff --git a/dbtai/templates/prompts.py b/dbtai/templates/prompts.py index 31d4654..7276a12 100644 --- a/dbtai/templates/prompts.py +++ b/dbtai/templates/prompts.py @@ -45,6 +45,20 @@ The upstream models referenced in the code are: {upstream_models} +The documentation of the model itself is: +{model_description} +""", + +"chatbot_prompt": """ +You are a friendly chatbot, ready to answer questions about the following dbt model, named {model_name}. The code is: + +``` +{raw_code} +``` + +The upstream models referenced in the code are: +{upstream_models} + The documentation of the model itself is: {model_description} """ @@ -89,6 +103,19 @@ De oppstrømsmodellene som refereres i koden er: {upstream_models} +Dokumentasjonen av selve modellen er: +{model_description} +""", + +"chatbot_prompt": """ +Du er en vennlig chatbot, klar til å svare på spørsmål om følgende dbt-modell, med navn {model_name}. Koden er: +``` +{raw_code} +``` + +De oppstrømsmodellene som refereres i koden er: +{upstream_models} + Dokumentasjonen av selve modellen er: {model_description} """ @@ -134,6 +161,19 @@ 代码中引用的上游模型有： {upstream_models} +模型本身的文档说明是： +{model_description} +""", + +"chatbot_prompt": """ +你是一个友好的聊天机器人，随时准备回答有关以下名为{model_name}的dbt模型的问题。代码如下： +``` +{raw_code} +``` + +代码中引用的上游模型有： +{upstream_models} + 模型本身的文档说明是： {model_description} """ @@ -180,6 +220,19 @@ La documentación del propio modelo es: {model_description} +""", + +"chatbot_prompt": """ +Eres un chatbot amigable, listo para responder preguntas sobre el siguiente modelo dbt, llamado {model_name}. El código es: +``` +{raw_code} +``` + +Los modelos previos referenciados en el código son: +{upstream_models} + +La documentación del propio modelo es: +{model_description} """ }, "french": { @@ -223,6 +276,19 @@ Les modèles en amont référencés dans le code sont : {upstream_models} +La documentation du modèle lui-même est : +{model_description} +""", + +"chatbot_prompt": """ +Vous êtes un chatbot amical, prêt à répondre aux questions concernant le modèle dbt suivant, nommé {model_name}. Le code est : +``` +{raw_code} +``` + +Les modèles en amont référencés dans le code sont : +{upstream_models} + La documentation du modèle lui-même est : {model_description} """ @@ -268,6 +334,19 @@ Die im Code referenzierten vorgelagerten Modelle sind: {upstream_models} +Die Dokumentation des Modells selbst ist: +{model_description} +""", + +"chatbot_prompt": """ +Du bist ein freundlicher Chatbot, bereit Fragen zum folgenden dbt-Modell, genannt {model_name}, zu beantworten. Der Code lautet: +``` +{raw_code} +``` + +Die im Code referenzierten vorgelagerten Modelle sind: +{upstream_models} + Die Dokumentation des Modells selbst ist: {model_description} """