Add add_datasource tool

docs/api_integrations.md

@@ -13,7 +13,7 @@ Most of the time, you will only need to specify `instructions` for a new API to
 Below is the default.
 
 ```yaml
-definitions_root: ./api_definitions
+definitions_root: ./datasources
 drafter:
     model: models/gemini-1.5-pro-001
     ttl_seconds: 1800

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
 dependencies = [
   "beaker-kernel~=1.9.4",
   "archytas~=1.3.18",
-  "adhoc-api~=2.3.0",
+  "adhoc-api~=2.4.0",
   "requests",
   "PyYAML",
   "idc-index",

src/biome/agent.py

@@ -23,6 +23,8 @@
 
 JSON_OUTPUT = False
 
+DATASOURCES_FOLDER = "datasources"
+
 class MessageLogger():
     def __init__(self, agent_log_function, print_logger):
         self.agent_log = agent_log_function
@@ -54,12 +56,15 @@ def decorator(func):
 class BiomeAgent(BaseAgent):
     """
     You are the Biome Agent, a chat assistant that helps users with biomedical research tasks.
-    """
 
+    A 'datasource' is defined as an API or dataset or general collection of knowledge that you have access to.
+
+    An API should be considered a type of datasource.
+    """
     def __init__(self, context: BaseContext = None, tools: list = None, **kwargs):
         self.root_folder = Path(__file__).resolve().parent
 
-        api_def_dir = os.path.join(self.root_folder, 'api_definitions')
+        api_def_dir = os.path.join(self.root_folder, DATASOURCES_FOLDER)
         data_dir = (self.root_folder / ".." / "..").resolve() / "data"
         print(f"data_dir: {data_dir}")
 
@@ -117,7 +122,7 @@ def __init__(self, context: BaseContext = None, tools: list = None, **kwargs):
                 logger=self.logger,
             )
         except ValueError as e:
-            self.add_context(f"The APIs failed to load for this reason: {str(e)}. Please inform the user immediately.")
+            self.add_context(f"The datasources failed to load for this reason: {str(e)}. Please inform the user immediately.")
             self.api = None
 
         self.api_list = [spec['name'] for spec in specs]
@@ -199,13 +204,13 @@ async def drs_uri_info(self, uris: List[str]) -> List[dict]:
     @tool()
     async def add_example(self, api: str, code: str, query: str, notes: str = None) -> str:
         """
-        Add a successful code example to the API's examples.yaml documentation file.
-        This tool should be used after successfully completing a task with an API to capture the working code for future reference.
+        Add a successful code example to the datasource's examples.yaml documentation file.
+        This tool should be used after successfully completing a task with an datasource to capture the working code for future reference.
 
-        The API names must match one of the names in the agent's API list.
+        The API names must match one of the names in the agent's datasource list.
 
         Args:
-            api (str): The name of the API the example is for
+            api (str): The name of the datasource the example is for
             code (str): The working, successful code to add as an example
             query (str): A brief description of what the example demonstrates
             notes (str, optional): Additional notes about the example, such as implementation details
@@ -219,7 +224,7 @@ async def add_example(self, api: str, code: str, query: str, notes: str = None)
         try:
             api_folder = self.api_directories[api]
             # Construct path to examples.yaml file
-            examples_path = os.path.join(self.root_folder, "api_definitions", api_folder, "documentation", "examples.yaml")
+            examples_path = os.path.join(self.root_folder, DATASOURCES_FOLDER, api_folder, "documentation", "examples.yaml")
             os.makedirs(os.path.dirname(examples_path), exist_ok=True)
 
             # Create new example entry as a dictionary
@@ -270,12 +275,84 @@ def represent_str_as_block(dumper, data):
 
 
     @tool()
-    async def get_available_apis(self) -> list:
+    async def add_datasource(self,
+                             datasource: str,
+                             description: str,
+                             base_url: str,
+                             schema_location: str) -> str:
+        """
+        Adds a datasource to the list of supported datasources usable within Biome.
+        This will be added to the API and data source list.
+
+        Args:
+            datasource (str): The name of the target data source or API that will be added.
+            description (str): A plain text description of what the data source is based on your knowledge of what the user is asking for, combined with their description if their description is relevant, or, if you do not know about the target data source. If the user does not provide any information, rely on what you know. Target a paragraph in length.
+            schema_location (str): A URL or local filepath to fetch an OpenAPI schema from. If the user does not provide one, ask them for the URL or local filepath to the schema.
+            base_url (str): The base URL for the datasource that will be used for making OpenAPI calls. If the user does not provide one, ask them for the base URL of the API.
+        Returns:
+            str: Message indicating success or failure of adding the datasource.
+        """
+        datasources_path = os.path.join(self.root_folder, DATASOURCES_FOLDER)
+        datasource_name = datasource.lower().replace(' ', '_')
+        datasource_folder = os.path.join(datasources_path, datasource_name)
+        if (os.path.isdir(datasource_folder)):
+            return f"Failed to add datasource: {datasource}: {DATASOURCES_FOLDER}/{datasource_folder} already exists."
+
+        logger.info(f'adding datasource {datasource} to {datasource_folder}')
+        try:
+            if schema_location.startswith('http'):
+                response = requests.get(schema_location)
+                if response.status_code != 200:
+                    return f'Failed to get OpenAPI schema: {response.status_code}'
+                schema = response.content.decode("utf-8")
+            else: 
+                with open(schema_location, 'r') as f:
+                    schema = f.read()
+
+            documentation_folder = os.path.join(datasource_folder, 'documentation')
+            os.makedirs(documentation_folder, exist_ok=True)
+            with open(os.path.join(documentation_folder, 'schema.yaml'), 'w') as f:
+                f.write(str(schema))
+            with open(os.path.join(documentation_folder, 'examples.yaml'), 'a'):
+                pass
+
+        except Exception as e:
+            return f'Failed to get OpenAPI schema: {e}'
+
+        formatted_description = description.replace('\n', ' ')
+        template = f"""
+name: {datasource}
+description: {formatted_description}
+examples: !load_yaml documentation/examples.yaml
+cache_key: "api_assistant_{datasource_name}"
+raw_documentation: !load_txt documentation/schema.yaml
+
+documentation: !fill |
+    The base URL for the service is `{base_url}`
+    Below is the OpenAPI schema for the desired service.
+
+    {{raw_documentation}}
+"""
+        with open(os.path.join(datasource_folder, 'api.yaml'), 'w') as f:
+            f.write(template)
+        try:
+            datasource_yaml = Path(os.path.join(datasource_folder, 'api.yaml'))
+            datasource_spec = load_yaml_api(datasource_yaml)
+            self.api_specs.append(datasource_spec)
+            self.api.add_api(datasource_spec)
+
+        except Exception as e:
+            return f"Failed to add datasource: {e}"
+        logger.info(f'Successfully added {datasource} to {datasource_folder} and it is now available for use.')
+        return f'Successfully added {datasource} and it is now available for use.'
+
+    @tool()
+    async def get_available_datasources(self) -> list:
         """
-        Get list of APIs that the agent is designed to interact with.
+        Get list of datasources that the agent is designed to interact with.
 
         Returns:
-            list: The list of available APIs.
+            list: The list of available datasources.
         """
         return self.api_list
 

src/biome/api_agent.yaml

@@ -1,22 +1,22 @@
-definitions_root: ./api_definitions
+definitions_root: ./datasources
 drafter:
     model: models/gemini-1.5-pro-001
     ttl_seconds: 1800
 finalizer:
     model: gpt-4o
-default_cache_body: !fill | 
-    You will be given the entire API documentation. 
-    When you write code against this API, you should avail yourself of the appropriate query parameters, 
+default_cache_body: !fill |
+    You will be given the entire API documentation.
+    When you write code against this API, you should avail yourself of the appropriate query parameters,
     your understanding of the response model, and be cognizant that not all data is public and thus may require a token, etc.
     Unless you receive a 403 forbidden, assume the endpoints are unauthenticated.
     If the user says the API does not require authentication, OMIT code about tokens and token handling and token headers.
     When you are downloading, communicate via stdout that something has been downloaded.
     When you are doing complex things try to break them down step by step and implement appropriate exception handling.
     You must NEVER print out the entire result of a workload or API search result. Always slice it and show just a subset to the user. You can save it as a variable
     for later use, etc.
-    
+
     {instructions}
-    
+
     Here is the documentation.
-    
+
     {documentation}