Polished agents and skills collection

mindsdb_sdk/agents.py

@@ -5,16 +5,15 @@
 import datetime
 import json
 
-from mindsdb_sdk.databases import Databases
-from mindsdb_sdk.knowledge_bases import KnowledgeBase, KnowledgeBases
-from mindsdb_sdk.ml_engines import MLEngines
-from mindsdb_sdk.models import Model, Models
-from mindsdb_sdk.skills import Skill, Skills
+from mindsdb_sdk.knowledge_bases import KnowledgeBase
+from mindsdb_sdk.models import Model
+from mindsdb_sdk.skills import Skill
 from mindsdb_sdk.utils.objects_collection import CollectionBase
 
 _DEFAULT_LLM_MODEL = 'gpt-4o'
 _DEFAULT_LLM_PROMPT = 'Answer the user"s question in a helpful way: {{question}}'
 
+
 class AgentCompletion:
     """
     Represents a full MindsDB agent completion response.
@@ -25,6 +24,7 @@ class AgentCompletion:
 
 
     """
+
     def __init__(self, content: str, context: List[dict] = None):
         self.content = content
         self.context = context
@@ -76,6 +76,7 @@ class Agent:
 
     >>> agents.drop('my_agent')
     """
+
     def __init__(
             self,
             name: str,
@@ -84,9 +85,9 @@ def __init__(
             params: dict,
             created_at: datetime.datetime,
             updated_at: datetime.datetime,
-            provider: str= None,
+            provider: str = None,
             collection: CollectionBase = None
-            ):
+    ):
         self.name = name
         self.model_name = model_name
         self.provider = provider
@@ -178,22 +179,25 @@ def from_json(cls, json: dict, collection: CollectionBase):
 
 class Agents(CollectionBase):
     """Collection for agents"""
-    def __init__(self, api, project: str, knowledge_bases: KnowledgeBases, databases: Databases, models: Models, ml_engines: MLEngines, skills: Skills = None):
+
+    def __init__(self, project, api):
         self.api = api
         self.project = project
-        self.skills = skills or Skills(self.api, project)
-        self.databases = databases
-        self.knowledge_bases = knowledge_bases
-        self.ml_engines = ml_engines
-        self.models = models
+
+        self.knowledge_bases = project.knowledge_bases
+        self.models = project.models
+        self.skills = project.skills
+
+        self.databases = project.server.databases
+        self.ml_engines = project.server.ml_engines
 
     def list(self) -> List[Agent]:
         """
         List available agents.
 
         :return: list of agents
         """
-        data = self.api.agents(self.project)
+        data = self.api.agents(self.project.name)
         return [Agent.from_json(agent, self) for agent in data]
 
     def get(self, name: str) -> Agent:
@@ -204,7 +208,7 @@ def get(self, name: str) -> Agent:
 
         :return: agent with given name
         """
-        data = self.api.agent(self.project, name)
+        data = self.api.agent(self.project.name, name)
         return Agent.from_json(data, self)
 
     def completion(self, name: str, messages: List[dict]) -> AgentCompletion:
@@ -216,7 +220,7 @@ def completion(self, name: str, messages: List[dict]) -> AgentCompletion:
 
         :return: completion from querying the agent
         """
-        data = self.api.agent_completion(self.project, name, messages)
+        data = self.api.agent_completion(self.project.name, name, messages)
         if 'context' in data['message']:
             return AgentCompletion(data['message']['content'], data['message'].get('context'))
 
@@ -231,7 +235,7 @@ def completion_stream(self, name, messages: List[dict]) -> Iterable[object]:
 
         :return: iterable of completion chunks from querying the agent.
         """
-        return self.api.agent_completion_stream(self.project, name, messages)
+        return self.api.agent_completion_stream(self.project.name, name, messages)
 
     def _create_default_knowledge_base(self, agent: Agent, name: str) -> KnowledgeBase:
         # Make sure default ML engine for embeddings exists.
@@ -244,7 +248,7 @@ def _create_default_knowledge_base(self, agent: Agent, name: str) -> KnowledgeBa
             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}
+            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)
@@ -298,7 +302,6 @@ def add_files(self, name: str, file_paths: List[str], description: str, knowledg
         agent.skills.append(file_retrieval_skill)
         self.update(agent.name, agent)
 
-
     def add_file(self, name: str, file_path: str, description: str, knowledge_base: str = None):
         """
         Add a file to the agent for retrieval.
@@ -418,6 +421,7 @@ def create(
         for skill in skills:
             if isinstance(skill, str):
                 # Check if skill exists.
+                # TODO what this line does?
                 _ = self.skills.get(skill)
                 skill_names.append(skill)
                 continue
@@ -438,7 +442,7 @@ def create(
             model = model.name
             provider = 'mindsdb'
 
-        data = self.api.create_agent(self.project, name, model, provider, skill_names, params)
+        data = self.api.create_agent(self.project.name, name, model, provider, skill_names, params)
         return Agent.from_json(data, self)
 
     def update(self, name: str, updated_agent: Agent):
@@ -467,12 +471,12 @@ def update(self, name: str, updated_agent: Agent):
                 _ = self.skills.create(skill.name, skill.type, skill.params)
             updated_skills.add(skill.name)
 
-        existing_agent = self.api.agent(self.project, name)
+        existing_agent = self.api.agent(self.project.name, 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,
+            self.project.name,
             name,
             updated_agent.name,
             updated_agent.model_name,
@@ -488,4 +492,4 @@ def drop(self, name: str):
 
         :param name: Name of the agent to be dropped
         """
-        _ = self.api.delete_agent(self.project, name)
+        _ = self.api.delete_agent(self.project.name, name)

mindsdb_sdk/projects.py

@@ -48,9 +48,10 @@ class Project:
 
     """
 
-    def __init__(self, api, name, agents: Agents = None, skills: Skills = None, knowledge_bases: KnowledgeBases = None, databases: Databases = None, ml_engines: MLEngines = None):
+    def __init__(self, server, api, name):
         self.name = name
         self.api = api
+        self.server = server
 
         self.models = Models(self, api)
 
@@ -76,11 +77,10 @@ def __init__(self, api, name, agents: Agents = None, skills: Skills = None, know
         self.create_job = self.jobs.create
         self.drop_job = self.jobs.drop
 
-        self.databases = databases or Databases(api)
-        self.knowledge_bases = knowledge_bases or KnowledgeBases(self, api)
+        self.knowledge_bases = KnowledgeBases(self, api)
 
-        self.skills = skills or Skills(api, name)
-        self.agents = agents or Agents(api, name, self.knowledge_bases, self.databases, self.models, ml_engines, self.skills)
+        self.skills = Skills(self, api)
+        self.agents = Agents(self, api)
 
     def __repr__(self):
         return f'{self.__class__.__name__}({self.name})'
@@ -106,7 +106,6 @@ def drop_model_version(self, name: str, version: int):
         self.query(ast_query.to_string()).fetch()
 
 
-
 class Projects(CollectionBase):
     """
     Projects
@@ -133,8 +132,9 @@ class Projects(CollectionBase):
 
     """
 
-    def __init__(self, api):
+    def __init__(self, server, api):
         self.api = api
+        self.server = server
 
     def _list_projects(self):
         data = self.api.sql_query("select NAME from information_schema.databases where TYPE='project'")
@@ -147,7 +147,7 @@ def list(self) -> List[Project]:
         :return: list of Project objects
         """
         # select * from information_schema.databases where TYPE='project'
-        return [Project(self.api, name) for name in self._list_projects()]
+        return [Project(self.server, self.api, name) for name in self._list_projects()]
 
     def get(self, name: str = 'mindsdb') -> Project:
         """
@@ -158,7 +158,7 @@ def get(self, name: str = 'mindsdb') -> Project:
         """
         if name not in self._list_projects():
             raise AttributeError("Project doesn't exist")
-        return Project(self.api, name)
+        return Project(self.server, self.api, name)
 
     def create(self, name: str) -> Project:
         """
@@ -175,7 +175,7 @@ def create(self, name: str) -> Project:
         )
 
         self.api.sql_query(ast_query.to_string())
-        return Project(self.api, name)
+        return Project(self.server, self.api, name)
 
     def drop(self, name: str):
         """

mindsdb_sdk/server.py

@@ -27,27 +27,24 @@ class Server(Project):
 
     def __init__(self, api, skills: Skills = None, agents: Agents = None):
         # server is also mindsdb project
-        project_name = 'mindsdb'
         self.databases = Databases(api)
         self.ml_engines = MLEngines(api)
-        super().__init__(api, project_name, skills=skills, agents=agents, databases=self.databases, ml_engines=self.ml_engines)
+        super().__init__(self, api, 'mindsdb')
 
-        self.projects = Projects(api)
+        self.projects = Projects(self, api)
 
         # old api
         self.get_project = self.projects.get
         self.list_projects = self.projects.list
         self.create_project = self.projects.create
         self.drop_project = self.projects.drop
 
-
         # old api
         self.get_database = self.databases.get
         self.list_databases = self.databases.list
         self.create_database = self.databases.create
         self.drop_database = self.databases.drop
 
-
         self.ml_handlers = Handlers(self.api, 'ml')
         self.data_handlers = Handlers(self.api, 'data')
 

mindsdb_sdk/skills.py

@@ -82,7 +82,7 @@ def __init__(self, name: str, knowledge_base: str, description: str):
 
 class Skills(CollectionBase):
     """Collection for skills"""
-    def __init__(self, api, project: str):
+    def __init__(self, project, api):
         self.api = api
         self.project = project
 
@@ -92,7 +92,7 @@ def list(self) -> List[Skill]:
 
         :return: list of skills
         """
-        data = self.api.skills(self.project)
+        data = self.api.skills(self.project.name)
         return [Skill.from_json(skill) for skill in data]
 
     def get(self, name: str) -> Skill:
@@ -103,7 +103,7 @@ def get(self, name: str) -> Skill:
 
         :return: skill with the given name
         """
-        data = self.api.skill(self.project, name)
+        data = self.api.skill(self.project.name, name)
         return Skill.from_json(data)
 
     def create(self, name: str, type: str, params: dict = None) -> Skill:
@@ -116,7 +116,7 @@ def create(self, name: str, type: str, params: dict = None) -> Skill:
 
         :return: created skill object
         """
-        _ = self.api.create_skill(self.project, name, type, params)
+        _ = self.api.create_skill(self.project.name, name, type, params)
         if type == 'sql':
             return SQLSkill(name, params['tables'], params['database'], params['description'])
         return Skill(name, type, params)
@@ -130,7 +130,7 @@ def update(self, name: str, updated_skill: Skill) -> Skill:
 
         :return: updated skillobject
         """
-        data = self.api.update_skill(self.project, name, updated_skill.name, updated_skill.type, updated_skill.params)
+        data = self.api.update_skill(self.project.name, name, updated_skill.name, updated_skill.type, updated_skill.params)
         return Skill.from_json(data)
 
     def drop(self, name: str):
@@ -139,4 +139,4 @@ def drop(self, name: str):
 
         :param name: Name of the skill to be dropped
         """
-        _ = self.api.delete_skill(self.project, name)
+        _ = self.api.delete_skill(self.project.name, name)