Merge pull request #134 from mindsdb/staging

Release 3.0.0

examples/using_agents_with_retrieval.py

@@ -1,12 +1,19 @@
 import mindsdb_sdk
+from uuid import uuid4
+import os
 
 con = mindsdb_sdk.connect()
 
+open_ai_key = os.getenv('OPENAI_API_KEY')
+model_name = 'gpt-4'
+
 # Now create an agent that will use the model we just created.
-agent = con.agents.get('agent_with_retrieval')
-agent.add_file('./data/tokaido-rulebook.pdf', 'rule book for the board game takaido')
+agent = con.agents.create(name=f'mindsdb_retrieval_agent_{model_name}_{uuid4().hex}',
+                          model='gpt-4')
+
+agent.add_file('./data/tokaido-rulebook.pdf', 'rule book for the board game Tokaido')
+
 
-print('Ask a question: ')
-question = input()
+question = "what are the rules for the game takaido?"
 answer = agent.completion([{'question': question, 'answer': None}])
 print(answer.content)

examples/using_agents_with_text2sql.py

@@ -9,8 +9,8 @@
 
 # Now create an agent that will use the model we just created.
 agent = con.agents.create(name=f'mindsdb_sql_agent_{model_name}_{uuid4().hex}',
-                          model={'model_name': model_name, 'openai_api_key': open_ai_key},
-                          params={'openai_api_key': open_ai_key})
+                          model='gpt-4')
+
 
 # Set up a Postgres data source with our new agent.
 data_source = 'postgres'

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/agents.py

@@ -1,5 +1,5 @@
 from requests.exceptions import HTTPError
-from typing import List, Union
+from typing import Iterable, List, Union
 from urllib.parse import urlparse
 from uuid import uuid4
 import datetime
@@ -13,6 +13,7 @@
 from mindsdb_sdk.skills import Skill, Skills
 from mindsdb_sdk.utils.objects_collection import CollectionBase
 
+_DEFAULT_LLM_MODEL = 'gpt-4o'
 
 class AgentCompletion:
     """Represents a full MindsDB agent completion"""
@@ -37,6 +38,12 @@ class Agent:
     >>> completion = agent.completion([{'question': 'What is your name?', 'answer': None}])
     >>> print(completion.content)
 
+    Query an agent with streaming:
+
+    >>> completion = agent.completion_stream([{'question': 'What is your name?', 'answer': None}])
+    >>> for chunk in completion:
+            print(chunk.choices[0].delta.content)
+
     List all agents:
 
     >>> agents = agents.list()
@@ -68,10 +75,12 @@ def __init__(
             params: dict,
             created_at: datetime.datetime,
             updated_at: datetime.datetime,
+            provider: str= None,
             collection: CollectionBase = None
             ):
         self.name = name
         self.model_name = model_name
+        self.provider = provider
         self.skills = skills
         self.params = params
         self.created_at = created_at
@@ -81,6 +90,9 @@ def __init__(
     def completion(self, messages: List[dict]) -> AgentCompletion:
         return self.collection.completion(self.name, messages)
 
+    def completion_stream(self, messages: List[dict]) -> Iterable[object]:
+        return self.collection.completion_stream(self.name, messages)
+
     def add_files(self, file_paths: List[str], description: str, knowledge_base: str = None):
         """
         Add a list of files to the agent for retrieval.
@@ -131,6 +143,8 @@ def __eq__(self, other):
             return False
         if self.model_name != other.model_name:
             return False
+        if self.provider != other.provider:
+            return False
         if self.skills != other.skills:
             return False
         if self.params != other.params:
@@ -148,6 +162,7 @@ def from_json(cls, json: dict, collection: CollectionBase):
             json['params'],
             json['created_at'],
             json['updated_at'],
+            json['provider'],
             collection
         )
 
@@ -195,18 +210,32 @@ def completion(self, name: str, messages: List[dict]) -> AgentCompletion:
         data = self.api.agent_completion(self.project, name, messages)
         return AgentCompletion(data['message']['content'])
 
+    def completion_stream(self, name, messages: List[dict]) -> Iterable[object]:
+        """
+        Queries the agent for a completion and streams the response as an iterable object.
+
+        :param name: Name of the agent
+        :param messageS: List of messages to be sent to the agent
+
+        :return: iterable of completion chunks from querying the agent.
+        """
+        return self.api.agent_completion_stream(self.project, name, messages)
+
     def _create_default_knowledge_base(self, agent: Agent, name: str) -> KnowledgeBase:
         # Make sure default ML engine for embeddings exists.
         try:
             _ = self.ml_engines.get('langchain_embedding')
         except AttributeError:
             _ = self.ml_engines.create('langchain_embedding', 'langchain_embedding')
         # Include API keys in embeddings.
-        agent_model = self.models.get(agent.model_name)
-        training_options = json.loads(agent_model.data.get('training_options', '{}'))
-        training_options_using = training_options.get('using', {})
-        api_key_params = {k:v for k, v in training_options_using.items() if 'api_key' in k}
-        kb = self.knowledge_bases.create(name, params=api_key_params)
+        if agent.provider == "mindsdb":
+            agent_model = self.models.get(agent.model_name)
+            training_options = json.loads(agent_model.data.get('training_options', '{}'))
+            training_options_using = training_options.get('using', {})
+            api_key_params = {k:v for k, v in training_options_using.items() if 'api_key' in k}
+            kb = self.knowledge_bases.create(name, params=api_key_params)
+        else:
+            kb = self.knowledge_bases.create(name)
         # Wait for underlying embedding model to finish training.
         kb.model.wait_complete()
         return kb
@@ -344,6 +373,13 @@ def add_database(self, name: str, database: str, tables: List[str], description:
         }
         database_sql_skill = self.skills.create(skill_name, 'sql', sql_params)
         agent = self.get(name)
+
+        if not agent.params:
+            agent.params = {}
+        if 'prompt_template' not in agent.params:
+            # Set default prompt template. This is for langchain agent check.
+            agent.params['prompt_template'] = 'using mindsdb sqltoolbox'
+
         agent.skills.append(database_sql_skill)
         self.update(agent.name, agent)
 
@@ -354,37 +390,41 @@ def _create_ml_engine_if_not_exists(self, name: str = 'langchain'):
             # Create the engine if it doesn't exist.
             _ = self.ml_engines.create('langchain', handler='langchain')
 
-    def _create_model_if_not_exists(self, name: str, model: Union[Model, dict]) -> Model:
+    def _create_model_if_not_exists(self, name: str, model: Union[Model, dict, str]) -> str:
         # Create langchain engine if it doesn't exist.
         self._create_ml_engine_if_not_exists()
         # Create a default model if it doesn't exist.
         default_model_params = {
             'predict': 'answer',
-            'mode': 'retrieval',
             'engine': 'langchain',
             'prompt_template': 'Answer the user"s question in a helpful way: {{question}}',
             # Use GPT-4 by default.
             'provider': 'openai',
             'model_name': 'gpt-4'
         }
-        if model is None:
-            return self.models.create(
-                f'{name}_default_model',
-                **default_model_params
-            )
+
         if isinstance(model, dict):
             default_model_params.update(model)
             # Create model with passed in params.
             return self.models.create(
                 f'{name}_default_model',
                 **default_model_params
-            )
+            ).name
+
+        if model is None:
+            # Create model with default params.
+            return _DEFAULT_LLM_MODEL
+
+        if isinstance(model, Model):
+            return model.name
+
         return model
 
     def create(
             self,
             name: str,
-            model: Union[Model, dict] = None,
+            model: Union[Model, dict, str] = None,
+            provider: str = None,
             skills: List[Union[Skill, str]] = None,
             params: dict = None) -> Agent:
         """
@@ -409,9 +449,9 @@ def create(
             _ = self.skills.create(skill.name, skill.type, skill.params)
             skill_names.append(skill.name)
 
-        # Create a default model if it doesn't exist.
         model = self._create_model_if_not_exists(name, model)
-        data = self.api.create_agent(self.project, name, model.name, skill_names, params)
+
+        data = self.api.create_agent(self.project, name, model, provider, skill_names, params)
         return Agent.from_json(data, self)
 
     def update(self, name: str, updated_agent: Agent):

mindsdb_sdk/connectors/rest_api.py

@@ -1,11 +1,13 @@
 from functools import wraps
 from typing import List, Union
 import io
+import json
 
 import requests
 import pandas as pd
 
 from mindsdb_sdk import __about__
+from sseclient import SSEClient
 
 
 def _try_relogin(fnc):
@@ -261,14 +263,24 @@ def agent_completion(self, project: str, name: str, messages: List[dict]):
         return r.json()
 
     @_try_relogin
-    def create_agent(self, project: str, name: str, model: str, skills: List[str] = None, params: dict = None):
+    def agent_completion_stream(self, project: str, name: str, messages: List[dict]):
+        url = self.url + f'/api/projects/{project}/agents/{name}/completions/stream'
+        stream = requests.post(url, json={'messages': messages}, stream=True)
+        client = SSEClient(stream)
+        for chunk in client.events():
+            # Stream objects loaded from SSE events 'data' param.
+            yield json.loads(chunk.data)
+
+    @_try_relogin
+    def create_agent(self, project: str, name: str, model: str = None, provider: str = None, skills: List[str] = None, params: dict = None):
         url = self.url + f'/api/projects/{project}/agents'
         r = self.session.post(
             url,
             json={
                 'agent': {
                     'name': name,
                     'model_name': model,
+                    'provider': provider,
                     'skills': skills,
                     'params': params
                 }

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)

requirements.txt

@@ -1,6 +1,7 @@
 requests
 pandas >= 1.3.5
-mindsdb-sql >= 0.13.0, < 0.14.0
+mindsdb-sql >= 0.17.0, < 1.0.0
 docstring-parser >= 0.7.3
 tenacity >= 8.0.1
 openai >= 1.15.0
+sseclient-py >= 1.8.0