Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 13 additions & 45 deletions src/main/brmcpserver/tools.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
import time
import logging

from pydantic import Field, BeforeValidator
from pydantic import Field
from typing_extensions import Annotated
from datetime import datetime, timezone
from typing import List, Optional, Dict
from models import Dataset, LogEvent, Datapoint, Timeseries

Expand All @@ -21,6 +20,7 @@ def __init__(self, bronto_client: BrontoClient):
def register(self, mcp):
mcp.tool(
name='search_logs',
title='Execute Event Query',
description="""Searches log data. This tool returns a list of log events and their attributes
The prompt should be a question or statement that you want for log data to be searched,
such as "Can you please search some log data from datasets related to the Bronto ingestion system?".
Expand All @@ -29,7 +29,8 @@ def register(self, mcp):
)(self.search_logs)

mcp.tool(
name='compute_metrics',
name='timeseries',
title='Execute Aggregate or Time-Series Query',
description="""Computes metric data from log data. This tool returns a list of data points for each key in the group_by_keys
list. Each list represents the value of the computed metrics for a subset of the provided time range.

Expand All @@ -38,24 +39,17 @@ def register(self, mcp):
time per path for the last hour?". The answer would then return the AVG(response_time) metric, grouped by URL path,
and split into a list of data points, one per every 5 minutes of the provided time range.
"""
)(self.compute_metrics)

mcp.tool(
name='get_timestamp_as_unix_epoch',
description="""Provides a unix timestamp (in milliseconds) since epoch representation of the input time. This tool
takes 1 string as parameters, representing a time in the following format '%Y-%m-%d %H:%M:%S'. For instance with
input_time='2025-05-01 00:00:00' then this tool returns 1746054000000. And with input_time='2025-05-01 01:00:00',
then this tool returns 1746057600000
"""
)(self.get_timestamp_as_unix_epoch)
)(self.timeseries)

mcp.tool(
name='get_datasets',
title='Retrieve a List of Logs',
description='Fetches all dataset details'
)(self.get_datasets)

mcp.tool(
name='get_datasets_by_name',
title='Find Datasets by Name and Collection',
description="""Fetches details about a Bronto dataset. A dataset is uniquely identify by its name and its
collection name. In other words, several datasets with the same name can be associated with different collections.
However only one dataset with a given name can be associated to a given collection.
Expand All @@ -64,6 +58,7 @@ def register(self, mcp):

mcp.tool(
name='get_keys',
title='Get Top Keys for a Specific Log ID',
description="""Fetches all keys present in a dataset, which is represented by a log ID.
This tool takes a log ID as parameter. A log ID is a string representing a UUID. A log ID maps to a dataset and
collection name. So given a dataset and collection name, it is possible to retrieve its log ID by using another tool
Expand All @@ -74,6 +69,7 @@ def register(self, mcp):

mcp.tool(
name='get_all_datasets_keys',
title='Retrieve Top Keys for Logs',
description="""Fetches all keys present in all datasets.
This tool returns a list of strings. Each string provides the name of a key present in the provided
dataset. This tool is useful in cases such as:
Expand All @@ -84,16 +80,12 @@ def register(self, mcp):

mcp.tool(
name='get_key_values',
title='Retrieve Values for a Dataset Key',
description="""Fetches the values of the provided key and dataset ID.
This tool returns a list of strings. Each string provides the value of the key provided as input, for
the dataset provided as input."""
)(self.get_key_values)

mcp.tool(
name='get_current_time',
description="This tool provides the current time in the YYYY-MM-DD HH:mm:ss format"
)(self.get_current_time)

def search_logs(
self,
timerange_start: Annotated[Optional[int], Field(
Expand All @@ -118,8 +110,8 @@ def search_logs(
exposed by this MCP server. In any case, following SQL syntax,
- key names should be double-quoted
- key value should be single-quoted if they are expected to be strings of characters
- key value should not be quoted if they are expected to be numbers."""
)] = '',
- key value should not be quoted if they are expected to be numbers.""")
] = '',
) -> Annotated[
List[LogEvent],
Field(description='A list of log events and their attributes. Attributes are key-value pairs associated with '
Expand All @@ -129,7 +121,7 @@ def search_logs(
log_events = self.bronto_client.search(timerange_start, timerange_end, log_ids, search_filter, _select=['*', '@raw'])
return log_events

def compute_metrics(
def timeseries(
self,
timerange_start: Annotated[int, Field(
description='Unix timestamp in millisecond representing the start of a time range, e.g. 1756063146000',
Expand Down Expand Up @@ -187,26 +179,6 @@ def compute_metrics(
result[group_serie['name']] = timeseries
return result

@staticmethod
def _validate_input_time(input_time: str) -> str:
try:
datetime.strptime(input_time, '%Y-%m-%d %H:%M:%S')
return input_time
except ValueError as e:
raise e

@staticmethod
def get_timestamp_as_unix_epoch(
input_time: Annotated[
str,
BeforeValidator(_validate_input_time),
Field(description='Time represented in the "%Y-%m-%d %H:%M:%S" format. Timezone is assumed to be UTC')]
) -> Annotated[
int,
Field(description='A unix timestamp (in milliseconds) since epoch, representing the `input_time` parameter')
]:
return int(datetime.strptime(input_time, '%Y-%m-%d %H:%M:%S').replace(tzinfo=timezone.utc).timestamp()) * 1000

def get_datasets(self) -> Annotated[
List[Dataset],
Field(description='''List of datasets. Each dataset object contains
Expand Down Expand Up @@ -275,7 +247,3 @@ def get_key_values(
keys_and_values = datasets_top_keys_and_values.get(log_id, {})
key_and_values = keys_and_values.get(key, {})
return key_and_values.get('values', {}).get(key, [])

@staticmethod
def get_current_time() -> Annotated[str, Field(description='the current time in the YYYY-MM-DD HH:mm:ss format')]:
return datetime.strftime(datetime.now(), '%Y-%m-%d %H:%M:%S')
34 changes: 7 additions & 27 deletions src/test/test_bronto_tools.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
import pytest
import time
from datetime import datetime
from unittest.mock import Mock

from tools import BrontoTools
from models import Datapoint, Timeseries, DatasetKey, LogEvent
from models import Datapoint, Timeseries, LogEvent
from clients import BrontoClient


Expand All @@ -19,25 +18,6 @@ def bronto_tools(mock_bronto_client):
return BrontoTools(mock_bronto_client)


def test_get_current_time():
current_time = BrontoTools.get_current_time()
assert isinstance(current_time, str)
datetime.strptime(current_time, "%Y-%m-%d %H:%M:%S")


def test_get_timestamp_as_unix_epoch():
test_time = "2025-05-01 00:00:00"
timestamp = BrontoTools.get_timestamp_as_unix_epoch(test_time)
assert isinstance(timestamp, int)
assert timestamp == 1746057600000


def test_get_timestamp_as_unix_epoch_wrong_format():
test_time = "2025/05/01 00:00:00"
with pytest.raises(ValueError):
BrontoTools.get_timestamp_as_unix_epoch(test_time)


def test_get_datasets(bronto_tools, mock_bronto_client):
mock_datasets = [
{
Expand Down Expand Up @@ -116,8 +96,8 @@ def test_search_logs(bronto_tools, mock_bronto_client):
assert log_event2 in log_events


def test_compute_metrics_no_group(bronto_tools, mock_bronto_client):
timestamp = BrontoTools.get_timestamp_as_unix_epoch("2023-01-01 00:00:00")
def test_timeseries_no_group(bronto_tools, mock_bronto_client):
timestamp = 1672531200000
mock_response = {
"totals": {
"count": 100,
Expand All @@ -128,7 +108,7 @@ def test_compute_metrics_no_group(bronto_tools, mock_bronto_client):
}
mock_bronto_client.search_post.return_value = mock_response

metrics = bronto_tools.compute_metrics(
metrics = bronto_tools.timeseries(
log_ids=["test_log_id"],
metric_functions=["SUM"],
timerange_start=int(time.time()) * 1000,
Expand All @@ -147,8 +127,8 @@ def test_compute_metrics_no_group(bronto_tools, mock_bronto_client):
)


def test_compute_metrics_single_group(bronto_tools, mock_bronto_client):
timestamp = BrontoTools.get_timestamp_as_unix_epoch("2023-01-01 00:00:00")
def test_timeseries_single_group(bronto_tools, mock_bronto_client):
timestamp = 1672531200000
mock_response = {
"groups_series": [
{
Expand All @@ -167,7 +147,7 @@ def test_compute_metrics_single_group(bronto_tools, mock_bronto_client):
}
mock_bronto_client.search_post.return_value = mock_response

metrics = bronto_tools.compute_metrics(
metrics = bronto_tools.timeseries(
log_ids=["test_log_id"],
metric_functions=["SUM"],
timerange_start=int(time.time()) * 1000,
Expand Down
Loading