Merge pull request #96 from mindsdb/agents-base

Add Agents Support to SDK

Author: tmichaeldb

examples/using_agents.py

@@ -0,0 +1,34 @@
+import mindsdb_sdk
+
+con = mindsdb_sdk.connect()
+
+# We currently support Langchain as a backend.
+print('Creating underlying langchain model for the agent to use...')
+try:
+    langchain_engine = con.ml_engines.get('langchain')
+except Exception:
+    # Create the engine if it doesn't exist.
+    langchain_engine = con.ml_engines.create('langchain', handler='langchain')
+
+# Actually create the underlying model the agent will use.
+langchain_model = con.models.create(
+    'agent_model',
+    predict='answer',
+    engine='langchain',
+    prompt_template='You are a spicy, cheeky assistant. Add some personality and flare when responding to the user question: {{question}}',
+    model_name='gpt-4-0125-preview' # This is the underlying LLM. Can use OpenAI, Claude, local Ollama, etc
+    # Can optionally set LLM args here. For example:
+    # temperature=0.0,
+    # max_tokens=1000,
+    # top_p=1.0,
+    # top_k=0,
+    # ...
+)
+print('Agent ready to use.')
+
+# Now create an agent that will use the model we just created.
+agent = con.agents.create('new_agent', langchain_model)
+print('Ask a question: ')
+question = input()
+answer = agent.completion([{'question': question, 'answer': None}])
+print(answer.content)

mindsdb_sdk/agents.py

@@ -0,0 +1,223 @@
+from requests.exceptions import HTTPError
+from typing import List, Union
+import datetime
+
+from mindsdb_sdk.models import Model
+from mindsdb_sdk.skills import Skill, Skills
+from mindsdb_sdk.utils.objects_collection import CollectionBase
+
+
+class AgentCompletion:
+    """Represents a full MindsDB agent completion"""
+    def __init__(self, content: str):
+        self.content = content
+
+    def __repr__(self):
+        return self.content
+
+
+class Agent:
+    """Represents a MindsDB agent.
+    
+    Working with agents:
+
+    Get an agent by name:
+
+    >>> agent = agents.get('my_agent')
+
+    Query an agent:
+
+    >>> completion = agent.completion([{'question': 'What is your name?', 'answer': None}])
+    >>> print(completion.content)
+
+    List all agents:
+
+    >>> agents = agents.list()
+
+    Create a new agent:
+
+    >>> model = models.get('my_model') # Or use models.create(...)
+    >>> # Connect your agent to a MindsDB table.
+    >>> text_to_sql_skill = skills.create('text_to_sql', 'sql', { 'tables': ['my_table'], 'database': 'my_database' })
+    >>> agent = agents.create('my_agent', model, [text_to_sql_skill])
+
+    Update an agent:
+
+    >>> new_model = models.get('new_model')
+    >>> agent.model_name = new_model.name
+    >>> new_skill = skills.create('new_skill', 'sql', { 'tables': ['new_table'], 'database': 'new_database' })
+    >>> updated_agent.skills.append(new_skill)
+    >>> updated_agent = agents.update('my_agent', agent)
+
+    Delete an agent by name:
+
+    >>> agents.delete('my_agent')
+    """
+    def __init__(
+            self,
+            name: str,
+            model_name: str,
+            skills: List[Skill],
+            params: dict,
+            created_at: datetime.datetime,
+            updated_at: datetime.datetime,
+            collection: CollectionBase = None
+            ):
+        self.name = name
+        self.model_name = model_name
+        self.skills = skills
+        self.params = params
+        self.created_at = created_at
+        self.updated_at = updated_at
+        self.collection = collection
+
+    def completion(self, messages: List[dict]) -> AgentCompletion:
+        return self.collection.completion(self.name, messages)
+
+    def __repr__(self):
+        return f'{self.__class__.__name__}(name: {self.name})'
+
+    def __eq__(self, other):
+        if self.name != other.name:
+            return False
+        if self.model_name != other.model_name:
+            return False
+        if self.skills != other.skills:
+            return False
+        if self.params != other.params:
+            return False
+        if self.created_at != other.created_at:
+            return False
+        return self.updated_at == other.updated_at
+
+    @classmethod
+    def from_json(cls, json: dict, collection: CollectionBase):
+        return cls(
+            json['name'],
+            json['model_name'],
+            [Skill.from_json(skill) for skill in json['skills']],
+            json['params'],
+            json['created_at'],
+            json['updated_at'],
+            collection
+        )
+
+
+class Agents(CollectionBase):
+    """Collection for agents"""
+    def __init__(self, api, project: str, skills: Skills = None):
+        self.api = api
+        self.project = project
+        self.skills = skills or Skills(self.api, project)
+
+    def list(self) -> List[Agent]:
+        """
+        List available agents.
+
+        :return: list of agents
+        """
+        data = self.api.agents(self.project)
+        return [Agent.from_json(agent, self) for agent in data]
+
+    def get(self, name: str) -> Agent:
+        """
+        Gets an agent by name.
+
+        :param name: Name of the agent
+
+        :return: agent with given name
+        """
+        data = self.api.agent(self.project, name)
+        return Agent.from_json(data, self)
+
+    def completion(self, name: str, messages: List[dict]) -> AgentCompletion:
+        """
+        Queries the agent for a completion.
+
+        :param name: Name of the agent
+        :param messages: List of messages to be sent to the agent
+
+        :return: completion from querying the agent
+        """
+        data = self.api.agent_completion(self.project, name, messages)
+        return AgentCompletion(data['message']['content'])
+
+    def create(
+            self,
+            name: str,
+            model: Model,
+            skills: List[Union[Skill, str]] = None,
+            params: dict = None) -> Agent:
+        """
+        Create new agent and return it
+
+        :param name: Name of the agent to be created
+        :param model: Model to be used by the agent
+        :param skills: List of skills to be used by the agent. Currently only 'sql' is supported.
+        :param params: Parameters for the agent
+
+        :return: created agent object
+        """
+        skills = skills or []
+        skill_names = []
+        for skill in skills:
+            if isinstance(skill, str):
+                # Check if skill exists.
+                _ = self.skills.get(skill)
+                skill_names.append(skill)
+                continue
+            # Create the skill if it doesn't exist.
+            _ = self.skills.create(skill.name, skill.type, skill.params)
+            skill_names.append(skill.name)
+
+        data = self.api.create_agent(self.project, name, model.name, skill_names, params)
+        return Agent.from_json(data, self)
+
+    def update(self, name: str, updated_agent: Agent):
+        """
+        Update an agent by name.
+
+        :param name: Name of the agent to be updated
+        :param updated_agent: Agent with updated fields
+
+        :return: updated agent object
+        """
+        updated_skills = set()
+        for skill in updated_agent.skills:
+            if isinstance(skill, str):
+                # Skill must exist.
+                _ = self.skills.get(skill)
+                updated_skills.add(skill)
+                continue
+            try:
+                # Create the skill if it doesn't exist.
+                _ = self.skills.get(skill.name)
+            except HTTPError as e:
+                if e.response.status_code != 404:
+                    raise e
+                # Doesn't exist
+                _ = self.skills.create(skill.name, skill.type, skill.params)
+            updated_skills.add(skill.name)
+
+        existing_agent = self.api.agent(self.project, name)
+        existing_skills = set([s.name for s in existing_agent['skills']])
+        skills_to_add = updated_skills.difference(existing_skills)
+        skills_to_remove = existing_skills.difference(updated_skills)
+        data = self.api.update_agent(
+            self.project,
+            name,
+            updated_agent.name,
+            updated_agent.model_name,
+            list(skills_to_add),
+            list(skills_to_remove),
+            updated_agent.params
+        )
+        return Agent.from_json(data, self)
+
+    def drop(self, name: str):
+        """
+        Drop an agent by name.
+
+        :param name: Name of the agent to be dropped
+        """
+        _ = self.api.delete_agent(self.project, name)

mindsdb_sdk/connect.py

@@ -2,6 +2,9 @@
 
 from mindsdb_sdk.connectors.rest_api import RestAPI
 
+DEFAULT_LOCAL_API_URL = 'http://127.0.0.1:47334'
+DEFAULT_CLOUD_API_URL = 'https://cloud.mindsdb.com'
+
 
 def connect(url: str = None, login: str = None, password: str = None, is_managed: bool = False, headers=None) -> Server:
     """
@@ -37,11 +40,11 @@ def connect(url: str = None, login: str = None, password: str = None, is_managed
     if url is None:
         if login is not None:
             # default is cloud
-            url = 'https://cloud.mindsdb.com'
+            url = DEFAULT_CLOUD_API_URL
         else:
             # is local
-            url = 'http://127.0.0.1:47334'
+            url = DEFAULT_LOCAL_API_URL
 
     api = RestAPI(url, login, password, is_managed, headers=headers)
 
-    return Server(api)
+    return Server(api)