Merge pull request #132 from mindsdb/mind-data-source-config

Update create_mind to Use New DataSourceConfig

Author: tmichaeldb

examples/using_database_mind_text2sql.py

@@ -1,7 +1,7 @@
 from uuid import uuid4
 
 from openai import OpenAI
-from mindsdb_sdk.utils.mind import create_mind
+from mindsdb_sdk.utils.mind import create_mind, DatabaseConfig
 import os
 
 
@@ -21,22 +21,28 @@
     base_url=base_url
 )
 
-# create a database mind
-mind = create_mind(
-    name = f'my_house_data_mind_{uuid4().hex}',
-    description= 'House Sales',
-    base_url= base_url,
-    api_key= MINDSDB_API_KEY,
-    model= model_name,
-    data_source_type='postgres',
-    data_source_connection_args={
+# Create a Database Config.
+pg_config = DatabaseConfig(
+    description='House Sales',
+    type='postgres',
+    connection_args={
         'user': 'demo_user',
         'password': 'demo_password',
         'host': 'samples.mindsdb.com',
         'port': '5432',
         'database': 'demo',
         'schema': 'demo_data'
-    }
+    },
+    tables=['house_sales']
+)
+
+# create a database mind
+mind = create_mind(
+    base_url= base_url,
+    api_key= MINDSDB_API_KEY,
+    name = f'my_house_data_mind_{uuid4().hex}',
+    data_source_configs=[pg_config],
+    model= model_name
 )
 
 # Actually pass in our tool to get a SQL completion.

mindsdb_sdk/utils/mind.py

@@ -1,4 +1,6 @@
-from typing import Optional
+from pydantic import BaseModel, Field
+from typing import List, Optional
+from uuid import uuid4
 
 import requests
 from logging import getLogger
@@ -16,27 +18,48 @@ def __init__(self, name):
         self.name = name
 
 
+class DataSourceConfig(BaseModel):
+    """
+    Represents a data source that can be made available to a Mind.
+    """
+    id: str = Field(default_factory=lambda: uuid4().hex)
+
+    # Description for underlying agent to know, based on context, whether to access this data source.
+    description: str
+
+
+class DatabaseConfig(DataSourceConfig):
+    """
+    Represents a database that can be made available to a Mind.
+    """
+
+    # Integration name (e.g. postgres)
+    type: str
+
+    # Args for connecting to database.
+    connection_args: dict
+
+    # Tables to make available to the Mind (defaults to ALL).
+    tables: List[str] = []
+
+
 # Create mind entity util function
 def create_mind(
         base_url: str,
         api_key: str,
         name: str,
-        description: str,
-        data_source_type: str,
-        data_source_connection_args: dict,
+        data_source_configs: List[DataSourceConfig] = None,
         model: Optional[str] = None,
 ) -> Mind:
     """
     Create a mind entity in LiteLLM proxy.
 
     Args:
-        base_url: MindsDB base URL
-        api_key: MindsDB API key
-        name: Mind name
-        description: Mind description
+        base_url (str): MindsDB base URL
+        api_key (str): MindsDB API key
+        name (str): Mind name
+        data_source_configs (List[DataSourceConfig]): Data sources to make available to the mind
         model: Model orchestrating the AI reasoning loop
-        data_source_type: Data source type
-        data_source_connection_args: Data source connection arguments
 
     Returns:
         Mind: Mind entity
@@ -46,10 +69,8 @@ def create_mind(
     headers = {"Authorization": f"Bearer {api_key}"}
     payload = {
         "name": name,
-        "description": description,
-        "model": model,
-        "data_source_type": data_source_type,
-        "data_source_connection_args": data_source_connection_args
+        "data_source_configs": [d.model_dump() for d in data_source_configs],
+        "model": model
     }
     try:
         response = requests.post(url, json=payload, headers=headers)