chatbot

README.md

@@ -95,6 +95,17 @@ Simply read a model and it's context to explain what the model actually does, an
 dbtai explain <model_name>
 ```
 
+### Chat
+You can open an interactive chat with a dbt model:
+
+```bash
+dbtai chat <model_name>
+```
+
+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!
 

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})

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"""

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
+        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

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}
 """