Merge pull request #102 from mindsdb/agents-retrieval

Add File Retrieval to Agents

Author: tmichaeldb

mindsdb_sdk/agents.py

@@ -1,7 +1,11 @@
 from requests.exceptions import HTTPError
 from typing import List, Union
+from uuid import uuid4
 import datetime
+import pandas as pd
 
+from mindsdb_sdk.databases import Databases
+from mindsdb_sdk.knowledge_bases import KnowledgeBases
 from mindsdb_sdk.models import Model
 from mindsdb_sdk.skills import Skill, Skills
 from mindsdb_sdk.utils.objects_collection import CollectionBase
@@ -74,6 +78,14 @@ def __init__(
     def completion(self, messages: List[dict]) -> AgentCompletion:
         return self.collection.completion(self.name, messages)
 
+    def add_file(self, file_path: str, description: str, knowledge_base: str = None):
+        """
+        Add a file to the agent for retrieval.
+
+        :param file_path: Path to the file to be added.
+        """
+        self.collection.add_file(self.name, file_path, description, knowledge_base)
+
     def __repr__(self):
         return f'{self.__class__.__name__}(name: {self.name})'
 
@@ -105,10 +117,12 @@ def from_json(cls, json: dict, collection: CollectionBase):
 
 class Agents(CollectionBase):
     """Collection for agents"""
-    def __init__(self, api, project: str, skills: Skills = None):
+    def __init__(self, api, project: str, knowledge_bases: KnowledgeBases, databases: Databases, skills: Skills = None):
         self.api = api
         self.project = project
         self.skills = skills or Skills(self.api, project)
+        self.databases = databases
+        self.knowledge_bases = knowledge_bases
 
     def list(self) -> List[Agent]:
         """
@@ -142,6 +156,55 @@ def completion(self, name: str, messages: List[dict]) -> AgentCompletion:
         data = self.api.agent_completion(self.project, name, messages)
         return AgentCompletion(data['message']['content'])
 
+    def add_file(self, name: str, file_path: str, description: str, knowledge_base: str = None):
+        """
+        Add a file to the agent for retrieval.
+
+        :param name: Name of the agent
+        :param file_path: Path to the file to be added, or name of existing file.
+        :param description: Description of the file. Used by agent to know when to do retrieval
+        :param knowledge_base: Name of an existing knowledge base to be used. Will create a default knowledge base if not given.
+        """
+        filename = file_path.split('/')[-1]
+        filename_no_extension = filename.split('.')[0]
+        try:
+            _ = self.api.get_file_metadata(filename_no_extension)
+        except HTTPError as e:
+            if e.response.status_code >= 400 and e.response.status_code != 404:
+                raise e
+            # Upload file if it doesn't exist.
+            with open(file_path, 'rb') as file:
+                content = file.read()
+                df = pd.DataFrame.from_records([{'content': content}])
+                self.api.upload_file(filename_no_extension, df)
+
+        # Insert uploaded file into new knowledge base.
+        if knowledge_base is not None:
+            kb = self.knowledge_bases.get(knowledge_base)
+        else:
+            kb_name = f'{name}_{filename_no_extension}_kb'
+            try:
+                kb = self.knowledge_bases.get(kb_name)
+            except AttributeError as e:
+                # Create KB if it doesn't exist.
+                kb = self.knowledge_bases.create(kb_name)
+                # Wait for underlying embedding model to finish training.
+                kb.model.wait_complete()
+
+        # Insert the entire file.
+        kb.insert(self.databases.files.tables.get(filename_no_extension))
+
+        # Make sure skill name is unique.
+        skill_name = f'{filename_no_extension}_retrieval_skill_{uuid4()}'
+        retrieval_params = {
+            'source': kb.name,
+            'description': description,
+        }
+        file_retrieval_skill = self.skills.create(skill_name, 'knowledge_base', retrieval_params)
+        agent = self.get(name)
+        agent.skills.append(file_retrieval_skill)
+        self.update(agent.name, agent)
+
     def create(
             self,
             name: str,
@@ -200,7 +263,7 @@ def update(self, name: str, updated_agent: Agent):
             updated_skills.add(skill.name)
 
         existing_agent = self.api.agent(self.project, name)
-        existing_skills = set([s.name for s in existing_agent['skills']])
+        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(

mindsdb_sdk/connectors/rest_api.py

@@ -146,6 +146,19 @@ def upload_file(self, name: str, df: pd.DataFrame):
         )
         _raise_for_status(r)
 
+    @_try_relogin
+    def get_file_metadata(self, name: str) -> dict:
+        # No endpoint currently to get single file.
+        url = self.url + f'/api/files/'
+        r = self.session.get(url)
+        _raise_for_status(r)
+        all_file_metadata = r.json()
+        for metadata in all_file_metadata:
+            if metadata.get('name', None) == name:
+                return metadata
+        r.status_code = 404
+        raise requests.HTTPError(f'Not found: No file named {name} found', response=r)
+
     @_try_relogin
     def upload_byom(self, name: str, code: str, requirements: str):
 

mindsdb_sdk/projects.py

@@ -7,6 +7,7 @@
 from mindsdb_sdk.utils.sql import dict_to_binary_op
 
 from mindsdb_sdk.agents import Agents
+from mindsdb_sdk.databases import Databases
 from mindsdb_sdk.skills import Skills
 from mindsdb_sdk.utils.objects_collection import CollectionBase
 
@@ -48,7 +49,7 @@ class Project:
 
     """
 
-    def __init__(self, api, name):
+    def __init__(self, api, name, agents: Agents = None, skills: Skills = None, knowledge_bases: KnowledgeBases = None, databases: Databases = None):
         self.name = name
         self.api = api
 
@@ -76,9 +77,11 @@ def __init__(self, api, name):
         self.create_job = self.jobs.create
         self.drop_job = self.jobs.drop
 
-        self.skills = Skills(api, name)
-        self.agents = Agents(api, name, self.skills)
-        self.knowledge_bases = KnowledgeBases(self, api)
+        self.databases = databases or Databases(api)
+        self.knowledge_bases = knowledge_bases or KnowledgeBases(self, api)
+
+        self.skills = skills or Skills(api, name)
+        self.agents = agents or Agents(api, name, self.knowledge_bases, self.databases, self.skills)
 
     def __repr__(self):
         return f'{self.__class__.__name__}({self.name})'

mindsdb_sdk/server.py

@@ -25,10 +25,11 @@ class Server(Project):
 
     """
 
-    def __init__(self, api):
+    def __init__(self, api, skills: Skills = None, agents: Agents = None):
         # server is also mindsdb project
         project_name = 'mindsdb'
-        super().__init__(api, project_name)
+        self.databases = Databases(api)
+        super().__init__(api, project_name, skills=skills, agents=agents, databases=self.databases)
 
         self.projects = Projects(api)
 
@@ -38,7 +39,6 @@ def __init__(self, api):
         self.create_project = self.projects.create
         self.drop_project = self.projects.drop
 
-        self.databases = Databases(api)
 
         # old api
         self.get_database = self.databases.get

mindsdb_sdk/skills.py

@@ -52,6 +52,8 @@ def __repr__(self):
     def from_json(cls, json: dict):
         if json['type'] == 'sql':
             return SQLSkill(json['name'], json['params']['tables'], json['params']['database'])
+        if json['type'] == 'retrieval':
+            return RetrievalSkill(json['name'], json['params']['knowledge_base'], json['params']['description'])
         return Skill(json['name'], json['type'], json['params'])
 
 
@@ -64,6 +66,15 @@ def __init__(self, name: str, tables: List[str], database: str):
         }
         super().__init__(name, 'sql', params)
 
+class RetrievalSkill(Skill):
+    """Represents a MindsDB skill for agents to interact with MindsDB data sources"""
+    def __init__(self, name: str, knowledge_base: str, description: str):
+        params = {
+            'knowledge_base': knowledge_base,
+            'description': description
+        }
+        super().__init__(name, 'knowledge_base', params)
+
 
 class Skills(CollectionBase):
     """Collection for skills"""