Created Primary Documentation

src/vanna/ZhipuAI/ZhipuAI_Chat.py.md

@@ -0,0 +1,153 @@
+## 🤖 ZhipuAI_Chat 🤖
+
+### Table of Contents
+
+- [Introduction](#introduction)
+- [Dependencies](#dependencies)
+- [Class Definition](#class-definition)
+  - [Attributes](#attributes)
+  - [Methods](#methods)
+- [Usage](#usage)
+- [Example](#example)
+
+### Introduction
+
+The `ZhipuAI_Chat` class is an easy-to-use wrapper for the ZhipuAI Chat API, which allows you to generate text and code from natural language prompts. It provides a simplified interface for sending prompts to the API and receiving the results.
+
+### Dependencies
+
+This library depends on the following:
+
+- `pandas`
+- `zhipuai`
+
+### Class Definition
+
+#### Attributes
+
+- `api_key`: Your ZhipuAI API key.
+- `model`: The ZhipuAI model to use. The default is "glm-4".
+
+#### Methods
+
+- `get_sql_prompt`: Generates a prompt for the ZhipuAI Chat API to generate SQL code from a natural language prompt.
+- `get_followup_questions_prompt`: Generates a prompt for the ZhipuAI Chat API to generate followup questions from a previous question and SQL code.
+- `generate_question`: Generates a natural language question from a SQL query.
+- `generate_plotly_code`: Generates Python Plotly code to chart the results of a pandas DataFrame.
+- `submit_prompt`: Sends a prompt to the ZhipuAI Chat API and returns the response.
+
+### Usage
+
+To use the `ZhipuAI_Chat` class, simply initialize it with your API key. You can then use the `get_sql_prompt`, `get_followup_questions_prompt`, `generate_question`, `generate_plotly_code`, and `submit_prompt` methods to interact with the ZhipuAI Chat API.
+
+### Example
+
+The following code shows how to use the `ZhipuAI_Chat` class to generate SQL code from a natural language prompt:
+
+```
+from zhipuai_chat import ZhipuAI_Chat
+
+# Initialize the ZhipuAI_Chat class with your API key
+chat = ZhipuAI_Chat(api_key="YOUR_API_KEY")
+
+# Generate a prompt for the ZhipuAI Chat API to generate SQL code
+prompt = chat.get_sql_prompt(
+    "What is the total sales for each product category?",
+    question_sql_list=[],
+    ddl_list=[],
+    doc_list=[],
+)
+
+# Send the prompt to the ZhipuAI Chat API and receive the response
+response = chat.submit_prompt(prompt)
+
+# The response will be a SQL query that answers the natural language prompt
+print(response)
+```
+
+### Example Usage - Generate SQL
+
+```python
+from zhipuai_chat import ZhipuAI_Chat
+
+# Initialize the ZhipuAI_Chat class with your API key
+chat = ZhipuAI_Chat(api_key="YOUR_API_KEY")
+
+# Generate a prompt for the ZhipuAI Chat API to generate SQL code
+prompt = chat.get_sql_prompt(
+    "What is the total sales for each product category?",
+    question_sql_list=[],
+    ddl_list=[],
+    doc_list=[],
+)
+
+# Send the prompt to the ZhipuAI Chat API and receive the response
+response = chat.submit_prompt(prompt)
+
+# The response will be a SQL query that answers the natural language prompt
+sql_query = response
+```
+
+### Example Usage - Generate Follow-up Question
+
+```python
+from zhipuai_chat import ZhipuAI_Chat
+
+# Initialize the ZhipuAI_Chat class with your API key
+chat = ZhipuAI_Chat(api_key="YOUR_API_KEY")
+
+# Generate a prompt for the ZhipuAI Chat API to generate SQL code
+prompt = chat.get_followup_questions_prompt(
+    "What is the total sales for each product category?",
+    df=pd.DataFrame(),
+    question_sql_list=[],
+    ddl_list=[],
+    doc_list=[],
+)
+
+# Send the prompt to the ZhipuAI Chat API and receive the response
+response = chat.submit_prompt(prompt)
+
+# The response will be a list of follow-up questions
+followup_questions = response
+```
+
+### Example Usage - Generate Question from SQL
+
+```python
+from zhipuai_chat import ZhipuAI_Chat
+
+# Initialize the ZhipuAI_Chat class with your API key
+chat = ZhipuAI_Chat(api_key="YOUR_API_KEY")
+
+# Generate a prompt for the ZhipuAI Chat API to generate SQL code
+prompt = chat.generate_question(
+    "SELECT SUM(sales) FROM products GROUP BY product_category"
+)
+
+# Send the prompt to the ZhipuAI Chat API and receive the response
+response = chat.submit_prompt(prompt)
+
+# The response will be a question in natural language
+question = response
+```
+
+### Example Usage - Generate Plotly Code
+
+```python
+from zhipuai_chat import ZhipuAI_Chat
+
+# Initialize the ZhipuAI_Chat class with your API key
+chat = ZhipuAI_Chat(api_key="YOUR_API_KEY")
+
+# Generate a prompt for the ZhipuAI Chat API to generate SQL code
+prompt = chat.generate_plotly_code(
+    df=pd.DataFrame(),
+)
+
+# Send the prompt to the ZhipuAI Chat API and receive the response
+response = chat.submit_prompt(prompt)
+
+# The response will be Python Plotly code
+plotly_code = response
+```

src/vanna/ZhipuAI/ZhipuAI_embeddings.py.md

@@ -0,0 +1,139 @@
+## 🔗 Table of Contents
+
+1. Overview
+2. Usage
+3. Example
+
+---
+
+## 📚 Overview
+
+### Description
+The `ZhipuAI_Embeddings` class is a `VannaBase` class that allows users to generate embeddings from ZhipuAI. It is designed to work with `chromadb` and can be used to create embedding functions for `chromadb` vector stores.
+
+### 🛠️ Requirements
+- Python 3.6 or higher
+- `zhipuai` package
+- `chromadb` package
+
+---
+
+## 🧙‍♂️ Usage
+
+### Initialization
+```
+from typing import List
+from zhipuai import ZhipuAI
+from chromadb import Documents, EmbeddingFunction, Embeddings
+from ..base import VannaBase
+
+class ZhipuAI_Embeddings(VannaBase):
+    """
+    [future functionality] This function is used to generate embeddings from ZhipuAI.
+
+    Args:
+        VannaBase (_type_): _description_
+    """
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+        if "api_key" not in config:
+            raise Exception("Missing api_key in config")
+        self.api_key = config["api_key"]
+        self.client = ZhipuAI(api_key=self.api_key)
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        
+        embedding = self.client.embeddings.create(
+            model="embedding-2",
+            input=data,
+        )
+
+        return embedding.data[0].embedding
+    
+
+
+class ZhipuAIEmbeddingFunction(EmbeddingFunction[Documents]):
+    """
+    A embeddingFunction that uses ZhipuAI to generate embeddings which can use in chromadb.
+    usage: 
+    class MyVanna(ChromaDB_VectorStore, ZhipuAI_Chat):
+        def __init__(self, config=None):
+            ChromaDB_VectorStore.__init__(self, config=config)
+            ZhipuAI_Chat.__init__(self, config=config)
+    
+    config={'api_key': 'xxx'}
+    zhipu_embedding_function = ZhipuAIEmbeddingFunction(config=config)
+    config = {"api_key": "xxx", "model": "glm-4","path":"xy","embedding_function":zhipu_embedding_function}
+    
+    vn = MyVanna(config)
+    
+    """
+    def __init__(self, config=None):
+        if config is None or "api_key" not in config:
+            raise ValueError("Missing 'api_key' in config")
+        
+        self.api_key = config["api_key"]
+        self.model_name = config.get("model_name", "embedding-2")
+        
+        try:
+            self.client = ZhipuAI(api_key=self.api_key)
+        except Exception as e:
+            raise ValueError(f"Error initializing ZhipuAI client: {e}")
+
+    def __call__(self, input: Documents) -> Embeddings:
+        # Replace newlines, which can negatively affect performance.
+        input = [t.replace("\n", " ") for t in input]
+        all_embeddings = []
+        print(f"Generating embeddings for {len(input)} documents")
+
+        # Iterating over each document for individual API calls
+        for document in input:
+            try:
+                response = self.client.embeddings.create(
+                    model=self.model_name,
+                    input=document
+                )
+                # print(response)
+                embedding = response.data[0].embedding
+                all_embeddings.append(embedding)
+                # print(f"Cost required: {response.usage.total_tokens}")
+            except Exception as e:
+                raise ValueError(f"Error generating embedding for document: {e}")
+
+        return all_embeddings
+```
+
+**Initialization**
+  - `__init__(self, config=None)`: Initializes the ZhipuAIEmbeddingFunction object. Requires an `api_key` in the `config` dictionary.
+  - `model_name`: The name of the ZhipuAI model to use for generating embeddings. Defaults to "embedding-2".
+
+**Embedding Generation**
+  - `__call__(self, input: Documents) -> Embeddings`: Generates embeddings for a list of documents. The input documents should be a list of strings. The function returns a list of embeddings, where each embedding is a list of floats.
+
+---
+
+## 🏃 Example
+
+```python
+# Import the necessary libraries
+from zhipuai_embeddings import ZhipuAIEmbeddingFunction
+
+# Initialize the ZhipuAIEmbeddingFunction object
+embedding_function = ZhipuAIEmbeddingFunction(config={
+    "api_key": "YOUR_API_KEY",
+    "model_name": "YOUR_MODEL_NAME"
+})
+
+# Generate embeddings for a list of documents
+documents = ["document 1", "document 2", "document 3"]
+embeddings = embedding_function(documents)
+
+# The `embeddings` variable now contains a list of embeddings for the input documents.
+```
+
+---
+
+## 💡 Additional Notes
+
+- The cost of using the ZhipuAIEmbeddingFunction depends on the number of tokens in the input documents and the model used for generating embeddings.
+- The ZhipuAIEmbeddingFunction can be used with any `chromadb` vector store.

src/vanna/ZhipuAI/__init__.py.md

@@ -0,0 +1,59 @@
+**ZhipuAI Language Models**
+
+  **Table of Contents**
+
+  - [ZhipuAI_Chat](#zhipuai_chat)
+  - [ZhipuAI_Embeddings](#zhipuai_embeddings)
+      - [ZhipuAIEmbeddingFunction](#zhipuaiembeddingfunction)
+
+
+**ZhipuAI_Chat**
+
+The `ZhipuAI_Chat` class provides an interface to interact with the ZhipuAI chat assistant API. This class makes it easy to send messages to the API and receive responses.
+
+```python
+from .ZhipuAI_Chat import ZhipuAI_Chat
+
+# Create a new ZhipuAI_Chat object
+chat = ZhipuAI_Chat()
+
+# Send a message to the API
+response = chat.send_message("Hello, ZhipuAI!")
+
+# Print the response from the API
+print(response)
+```
+
+**ZhipuAI_Embeddings**
+
+The `ZhipuAI_Embeddings` class provides an interface to interact with the ZhipuAI embeddings API. This class makes it easy to generate embeddings for text data.
+
+```python
+from .ZhipuAI_Embeddings import ZhipuAI_Embeddings, ZhipuAIEmbeddingFunction
+
+# Create a new ZhipuAI_Embeddings object
+embeddings = ZhipuAI_Embeddings()
+
+# Generate embeddings for a piece of text
+embeddings = embeddings.generate_embeddings("Hello, world!")
+
+# Print the embeddings
+print(embeddings)
+```
+
+**ZhipuAIEmbeddingFunction**
+
+The `ZhipuAIEmbeddingFunction` class is a helper class that can be used to generate embeddings for text data using a specific embedding function.
+
+```python
+from .ZhipuAI_Embeddings import ZhipuAIEmbeddingFunction
+
+# Create a new ZhipuAIEmbeddingFunction object
+embedding_function = ZhipuAIEmbeddingFunction("average")
+
+# Generate embeddings for a piece of text
+embeddings = embedding_function.generate_embeddings("Hello, world!")
+
+# Print the embeddings
+print(embeddings)
+```