From 046236912faa5119965050c04ed8f1d9752f8a65 Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Thu, 12 Dec 2024 12:14:14 +0530 Subject: [PATCH 1/4] feat(agents-api): Add missing columns to indices Signed-off-by: Diwank Singh Tomer --- .../models/docs/search_docs_by_embedding.py | 9 ++++--- .../models/docs/search_docs_by_text.py | 9 ++++--- .../models/execution/update_execution.py | 2 +- ...grate_1733985509_add_columns_to_indices.py | 24 +++++++++++++++++++ 4 files changed, 37 insertions(+), 7 deletions(-) create mode 100644 agents-api/migrations/migrate_1733985509_add_columns_to_indices.py diff --git a/agents-api/agents_api/models/docs/search_docs_by_embedding.py b/agents-api/agents_api/models/docs/search_docs_by_embedding.py index 992e12f9d..a471c596c 100644 --- a/agents-api/agents_api/models/docs/search_docs_by_embedding.py +++ b/agents-api/agents_api/models/docs/search_docs_by_embedding.py @@ -91,7 +91,7 @@ def search_docs_by_embedding( snippet_counter[count(item)] := owners[owner_type, owner_id_str], owner_id = to_uuid(owner_id_str), - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id: item, @@ -142,13 +142,16 @@ def search_docs_by_embedding( input[owner_type, owner_id, query], # Restrict the search to all documents that match the owner - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id, - title, metadata, }}, + *docs {{ + doc_id, + title, + }}, # Search for snippets in the embedding space ~snippets:embedding_space {{ diff --git a/agents-api/agents_api/models/docs/search_docs_by_text.py b/agents-api/agents_api/models/docs/search_docs_by_text.py index ac1a9f54f..33b270465 100644 --- a/agents-api/agents_api/models/docs/search_docs_by_text.py +++ b/agents-api/agents_api/models/docs/search_docs_by_text.py @@ -90,7 +90,7 @@ def search_docs_by_text( candidate[doc_id] := input[owner_type, owner_id], - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id, @@ -145,13 +145,16 @@ def search_docs_by_text( metadata, ] := candidate[doc_id], - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id, - title, metadata, }}, + *docs {{ + doc_id, + title, + }}, search_result [ doc_id, snippet_data, diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py index f33368412..b485e34f8 100644 --- a/agents-api/agents_api/models/execution/update_execution.py +++ b/agents-api/agents_api/models/execution/update_execution.py @@ -79,7 +79,7 @@ def update_execution( validate_status_query = """ valid_status[count(status)] := - *executions { + *executions:execution_id_status_idx { status, execution_id: to_uuid($execution_id), task_id: to_uuid($task_id), diff --git a/agents-api/migrations/migrate_1733985509_add_columns_to_indices.py b/agents-api/migrations/migrate_1733985509_add_columns_to_indices.py new file mode 100644 index 000000000..5ae3b41a9 --- /dev/null +++ b/agents-api/migrations/migrate_1733985509_add_columns_to_indices.py @@ -0,0 +1,24 @@ +#/usr/bin/env python3 + +MIGRATION_ID = "add_columns_to_indices" +CREATED_AT = 1733985509.241258 + +def up(client): + client.run( + """ + { ::index drop executions:execution_id_status_idx } + { ::index create executions:execution_id_status_idx { task_id, execution_id, status } } + { ::index drop docs:owner_id_metadata_doc_id_idx } + { ::index create docs:owner_id_metadata_doc_id_idx { owner_type, owner_id, metadata, doc_id } } + """ + ) + +def down(client): + client.run( + """ + { ::index drop executions:execution_id_status_idx } + { ::index create executions:execution_id_status_idx { execution_id, status } } + { ::index drop docs:owner_id_metadata_doc_id_idx } + { ::index create docs:owner_id_metadata_doc_id_idx { owner_id, metadata, doc_id } } + """ + ) From 24e9afb34fac345349d2cf2333f5b4786b2cb8c0 Mon Sep 17 00:00:00 2001 From: creatorrr Date: Thu, 12 Dec 2024 06:45:26 +0000 Subject: [PATCH 2/4] refactor: Lint agents-api (CI) --- .../migrations/migrate_1733985509_add_columns_to_indices.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/agents-api/migrations/migrate_1733985509_add_columns_to_indices.py b/agents-api/migrations/migrate_1733985509_add_columns_to_indices.py index 5ae3b41a9..c461aa499 100644 --- a/agents-api/migrations/migrate_1733985509_add_columns_to_indices.py +++ b/agents-api/migrations/migrate_1733985509_add_columns_to_indices.py @@ -1,8 +1,9 @@ -#/usr/bin/env python3 +# /usr/bin/env python3 MIGRATION_ID = "add_columns_to_indices" CREATED_AT = 1733985509.241258 + def up(client): client.run( """ @@ -13,6 +14,7 @@ def up(client): """ ) + def down(client): client.run( """ From 3ec0a21020ced6670ebbb8b69d147c5a68b3246e Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Thu, 12 Dec 2024 13:56:10 +0300 Subject: [PATCH 3/4] feat: Add prometheus summary for cozo queries --- agents-api/agents_api/metrics/counters.py | 18 ++++++++++++------ .../agents_api/models/agent/create_agent.py | 4 ++-- .../models/agent/create_or_update_agent.py | 4 ++-- .../agents_api/models/agent/patch_agent.py | 4 ++-- .../agents_api/models/agent/update_agent.py | 4 ++-- .../agents_api/models/docs/create_doc.py | 4 ++-- .../agents_api/models/entry/create_entries.py | 4 ++-- .../models/execution/create_execution.py | 4 ++-- .../execution/create_execution_transition.py | 6 +++--- .../models/execution/create_temporal_lookup.py | 4 ++-- .../models/execution/update_execution.py | 4 ++-- .../agents_api/models/files/create_file.py | 4 ++-- .../models/session/create_or_update_session.py | 4 ++-- .../models/session/create_session.py | 4 ++-- .../models/session/update_session.py | 4 ++-- .../models/task/create_or_update_task.py | 4 ++-- .../agents_api/models/task/create_task.py | 4 ++-- .../agents_api/models/task/patch_task.py | 4 ++-- .../agents_api/models/task/update_task.py | 4 ++-- .../agents_api/models/tools/create_tools.py | 4 ++-- .../agents_api/models/tools/patch_tool.py | 4 ++-- .../agents_api/models/tools/update_tool.py | 4 ++-- .../models/user/create_or_update_user.py | 4 ++-- .../agents_api/models/user/create_user.py | 4 ++-- .../agents_api/models/user/patch_user.py | 4 ++-- .../agents_api/models/user/update_user.py | 4 ++-- 26 files changed, 63 insertions(+), 57 deletions(-) diff --git a/agents-api/agents_api/metrics/counters.py b/agents-api/agents_api/metrics/counters.py index f80236bf7..ec3f31436 100644 --- a/agents-api/agents_api/metrics/counters.py +++ b/agents-api/agents_api/metrics/counters.py @@ -2,28 +2,34 @@ from functools import wraps from typing import Awaitable, Callable, ParamSpec, TypeVar -from prometheus_client import Counter +from prometheus_client import Counter, Summary P = ParamSpec("P") T = TypeVar("T") -def increase_counter(metric_label: str, id_field_name: str = "developer_id"): +def query_metrics_update(metric_label: str, id_field_name: str = "developer_id"): def decor(func: Callable[P, T | Awaitable[T]]): - metric = Counter( - metric_label, + counter = Counter( + f"cozo_counter_{metric_label}", f"Number of {metric_label} calls", labelnames=(id_field_name,), ) + summary = Summary( + f"cozo_latency_{metric_label}", + metric_label, + ) + @summary.time() @wraps(func) def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: - metric.labels(kwargs.get(id_field_name, "not_set")).inc() + counter.labels(kwargs.get(id_field_name, "not_set")).inc() return func(*args, **kwargs) + @summary.time() @wraps(func) async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: - metric.labels(kwargs.get(id_field_name, "not_set")).inc() + counter.labels(kwargs.get(id_field_name, "not_set")).inc() return await func(*args, **kwargs) return async_wrapper if inspect.iscoroutinefunction(func) else wrapper diff --git a/agents-api/agents_api/models/agent/create_agent.py b/agents-api/agents_api/models/agent/create_agent.py index a9f0bfb8f..97d05459c 100644 --- a/agents-api/agents_api/models/agent/create_agent.py +++ b/agents-api/agents_api/models/agent/create_agent.py @@ -13,7 +13,7 @@ from ...autogen.openapi_model import Agent, CreateAgentRequest from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -58,7 +58,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("create_agent") +@query_metrics_update("create_agent") @beartype def create_agent( *, diff --git a/agents-api/agents_api/models/agent/create_or_update_agent.py b/agents-api/agents_api/models/agent/create_or_update_agent.py index 9a1feb717..2d01e1474 100644 --- a/agents-api/agents_api/models/agent/create_or_update_agent.py +++ b/agents-api/agents_api/models/agent/create_or_update_agent.py @@ -13,7 +13,7 @@ from ...autogen.openapi_model import Agent, CreateOrUpdateAgentRequest from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -49,7 +49,7 @@ Agent, one=True, transform=lambda d: {"id": UUID(d.pop("agent_id")), **d} ) @cozo_query -@increase_counter("create_or_update_agent") +@query_metrics_update("create_or_update_agent") @beartype def create_or_update_agent( *, diff --git a/agents-api/agents_api/models/agent/patch_agent.py b/agents-api/agents_api/models/agent/patch_agent.py index 99d4e3553..f57e15167 100644 --- a/agents-api/agents_api/models/agent/patch_agent.py +++ b/agents-api/agents_api/models/agent/patch_agent.py @@ -9,7 +9,7 @@ from ...autogen.openapi_model import PatchAgentRequest, ResourceUpdatedResponse from ...common.utils.cozo import cozo_process_mutate_data from ...common.utils.datetime import utcnow -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -37,7 +37,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("patch_agent") +@query_metrics_update("patch_agent") @beartype def patch_agent( *, diff --git a/agents-api/agents_api/models/agent/update_agent.py b/agents-api/agents_api/models/agent/update_agent.py index b36f687eb..7277bdbae 100644 --- a/agents-api/agents_api/models/agent/update_agent.py +++ b/agents-api/agents_api/models/agent/update_agent.py @@ -8,7 +8,7 @@ from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateAgentRequest from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -36,7 +36,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("update_agent") +@query_metrics_update("update_agent") @beartype def update_agent( *, diff --git a/agents-api/agents_api/models/docs/create_doc.py b/agents-api/agents_api/models/docs/create_doc.py index 3b9c8c9f7..3cd03d023 100644 --- a/agents-api/agents_api/models/docs/create_doc.py +++ b/agents-api/agents_api/models/docs/create_doc.py @@ -8,7 +8,7 @@ from ...autogen.openapi_model import CreateDocRequest, Doc from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -38,7 +38,7 @@ }, ) @cozo_query -@increase_counter("create_doc") +@query_metrics_update("create_doc") @beartype def create_doc( *, diff --git a/agents-api/agents_api/models/entry/create_entries.py b/agents-api/agents_api/models/entry/create_entries.py index a8671a6dd..724927b9c 100644 --- a/agents-api/agents_api/models/entry/create_entries.py +++ b/agents-api/agents_api/models/entry/create_entries.py @@ -10,7 +10,7 @@ from ...common.utils.cozo import cozo_process_mutate_data from ...common.utils.datetime import utcnow from ...common.utils.messages import content_to_json -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, mark_session_updated_query, @@ -41,7 +41,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("create_entries") +@query_metrics_update("create_entries") @beartype def create_entries( *, diff --git a/agents-api/agents_api/models/execution/create_execution.py b/agents-api/agents_api/models/execution/create_execution.py index 832532d6d..4c1b47bec 100644 --- a/agents-api/agents_api/models/execution/create_execution.py +++ b/agents-api/agents_api/models/execution/create_execution.py @@ -9,7 +9,7 @@ from ...autogen.openapi_model import CreateExecutionRequest, Execution from ...common.utils.cozo import cozo_process_mutate_data from ...common.utils.types import dict_like -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -38,7 +38,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("create_execution") +@query_metrics_update("create_execution") @beartype def create_execution( *, diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py index 59a63ed09..4acb11fef 100644 --- a/agents-api/agents_api/models/execution/create_execution_transition.py +++ b/agents-api/agents_api/models/execution/create_execution_transition.py @@ -12,7 +12,7 @@ ) from ...common.protocol.tasks import transition_to_execution_status, valid_transitions from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, cozo_query_async, @@ -224,7 +224,7 @@ def validate_transition_targets(data: CreateTransitionRequest) -> None: _kind="inserted", )( cozo_query( - increase_counter("create_execution_transition")( + query_metrics_update("create_execution_transition")( _create_execution_transition ) ) @@ -250,7 +250,7 @@ def validate_transition_targets(data: CreateTransitionRequest) -> None: _kind="inserted", )( cozo_query_async( - increase_counter("create_execution_transition_async")( + query_metrics_update("create_execution_transition_async")( _create_execution_transition ) ) diff --git a/agents-api/agents_api/models/execution/create_temporal_lookup.py b/agents-api/agents_api/models/execution/create_temporal_lookup.py index e47a505db..b2aa80989 100644 --- a/agents-api/agents_api/models/execution/create_temporal_lookup.py +++ b/agents-api/agents_api/models/execution/create_temporal_lookup.py @@ -8,7 +8,7 @@ from temporalio.client import WorkflowHandle from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -29,7 +29,7 @@ } ) @cozo_query -@increase_counter("create_temporal_lookup") +@query_metrics_update("create_temporal_lookup") @beartype def create_temporal_lookup( *, diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py index b485e34f8..cac545ab9 100644 --- a/agents-api/agents_api/models/execution/update_execution.py +++ b/agents-api/agents_api/models/execution/update_execution.py @@ -14,7 +14,7 @@ valid_previous_statuses as valid_previous_statuses_map, ) from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -43,7 +43,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("update_execution") +@query_metrics_update("update_execution") @beartype def update_execution( *, diff --git a/agents-api/agents_api/models/files/create_file.py b/agents-api/agents_api/models/files/create_file.py index 224597180..d995598fc 100644 --- a/agents-api/agents_api/models/files/create_file.py +++ b/agents-api/agents_api/models/files/create_file.py @@ -14,7 +14,7 @@ from pydantic import ValidationError from ...autogen.openapi_model import CreateFileRequest, File -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -63,7 +63,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("create_file") +@query_metrics_update("create_file") @beartype def create_file( *, diff --git a/agents-api/agents_api/models/session/create_or_update_session.py b/agents-api/agents_api/models/session/create_or_update_session.py index e34a63ca5..310bc8c9a 100644 --- a/agents-api/agents_api/models/session/create_or_update_session.py +++ b/agents-api/agents_api/models/session/create_or_update_session.py @@ -11,7 +11,7 @@ ResourceUpdatedResponse, ) from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -44,7 +44,7 @@ }, ) @cozo_query -@increase_counter("create_or_update_session") +@query_metrics_update("create_or_update_session") @beartype def create_or_update_session( *, diff --git a/agents-api/agents_api/models/session/create_session.py b/agents-api/agents_api/models/session/create_session.py index ce804399d..f20ee06a1 100644 --- a/agents-api/agents_api/models/session/create_session.py +++ b/agents-api/agents_api/models/session/create_session.py @@ -12,7 +12,7 @@ from pydantic import ValidationError from ...autogen.openapi_model import CreateSessionRequest, Session -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -45,7 +45,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("create_session") +@query_metrics_update("create_session") @beartype def create_session( *, diff --git a/agents-api/agents_api/models/session/update_session.py b/agents-api/agents_api/models/session/update_session.py index cc8b61f16..1765b811f 100644 --- a/agents-api/agents_api/models/session/update_session.py +++ b/agents-api/agents_api/models/session/update_session.py @@ -8,7 +8,7 @@ from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateSessionRequest from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -52,7 +52,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("update_session") +@query_metrics_update("update_session") @beartype def update_session( *, diff --git a/agents-api/agents_api/models/task/create_or_update_task.py b/agents-api/agents_api/models/task/create_or_update_task.py index 1f615a3ad..b4cf6e90e 100644 --- a/agents-api/agents_api/models/task/create_or_update_task.py +++ b/agents-api/agents_api/models/task/create_or_update_task.py @@ -18,7 +18,7 @@ from ...common.protocol.tasks import task_to_spec from ...common.utils.cozo import cozo_process_mutate_data from ...common.utils.datetime import utcnow -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -50,7 +50,7 @@ }, ) @cozo_query -@increase_counter("create_or_update_task") +@query_metrics_update("create_or_update_task") @beartype def create_or_update_task( *, diff --git a/agents-api/agents_api/models/task/create_task.py b/agents-api/agents_api/models/task/create_task.py index ab68a5b0c..b254d7572 100644 --- a/agents-api/agents_api/models/task/create_task.py +++ b/agents-api/agents_api/models/task/create_task.py @@ -17,7 +17,7 @@ ) from ...common.protocol.tasks import task_to_spec from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -49,7 +49,7 @@ }, ) @cozo_query -@increase_counter("create_task") +@query_metrics_update("create_task") @beartype def create_task( *, diff --git a/agents-api/agents_api/models/task/patch_task.py b/agents-api/agents_api/models/task/patch_task.py index 178b9daa3..80e136010 100644 --- a/agents-api/agents_api/models/task/patch_task.py +++ b/agents-api/agents_api/models/task/patch_task.py @@ -14,7 +14,7 @@ from ...autogen.openapi_model import PatchTaskRequest, ResourceUpdatedResponse, TaskSpec from ...common.protocol.tasks import task_to_spec from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -47,7 +47,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("patch_task") +@query_metrics_update("patch_task") @beartype def patch_task( *, diff --git a/agents-api/agents_api/models/task/update_task.py b/agents-api/agents_api/models/task/update_task.py index cd98d85d5..232f351ee 100644 --- a/agents-api/agents_api/models/task/update_task.py +++ b/agents-api/agents_api/models/task/update_task.py @@ -14,7 +14,7 @@ from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateTaskRequest from ...common.protocol.tasks import task_to_spec from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -46,7 +46,7 @@ }, ) @cozo_query -@increase_counter("update_task") +@query_metrics_update("update_task") @beartype def update_task( *, diff --git a/agents-api/agents_api/models/tools/create_tools.py b/agents-api/agents_api/models/tools/create_tools.py index 9b2be387a..d615ef877 100644 --- a/agents-api/agents_api/models/tools/create_tools.py +++ b/agents-api/agents_api/models/tools/create_tools.py @@ -9,7 +9,7 @@ from pydantic import ValidationError from ...autogen.openapi_model import CreateToolRequest, Tool -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -41,7 +41,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("create_tools") +@query_metrics_update("create_tools") @beartype def create_tools( *, diff --git a/agents-api/agents_api/models/tools/patch_tool.py b/agents-api/agents_api/models/tools/patch_tool.py index bc49b8121..45b655acc 100644 --- a/agents-api/agents_api/models/tools/patch_tool.py +++ b/agents-api/agents_api/models/tools/patch_tool.py @@ -8,7 +8,7 @@ from ...autogen.openapi_model import PatchToolRequest, ResourceUpdatedResponse from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -36,7 +36,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("patch_tool") +@query_metrics_update("patch_tool") @beartype def patch_tool( *, developer_id: UUID, agent_id: UUID, tool_id: UUID, data: PatchToolRequest diff --git a/agents-api/agents_api/models/tools/update_tool.py b/agents-api/agents_api/models/tools/update_tool.py index ef700a5f6..a36545000 100644 --- a/agents-api/agents_api/models/tools/update_tool.py +++ b/agents-api/agents_api/models/tools/update_tool.py @@ -11,7 +11,7 @@ UpdateToolRequest, ) from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -39,7 +39,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("update_tool") +@query_metrics_update("update_tool") @beartype def update_tool( *, diff --git a/agents-api/agents_api/models/user/create_or_update_user.py b/agents-api/agents_api/models/user/create_or_update_user.py index 3e9b1f3a6..7985bd92d 100644 --- a/agents-api/agents_api/models/user/create_or_update_user.py +++ b/agents-api/agents_api/models/user/create_or_update_user.py @@ -12,7 +12,7 @@ from pydantic import ValidationError from ...autogen.openapi_model import CreateOrUpdateUserRequest, User -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -46,7 +46,7 @@ ) @wrap_in_class(User, one=True, transform=lambda d: {"id": UUID(d.pop("user_id")), **d}) @cozo_query -@increase_counter("create_or_update_user") +@query_metrics_update("create_or_update_user") @beartype def create_or_update_user( *, diff --git a/agents-api/agents_api/models/user/create_user.py b/agents-api/agents_api/models/user/create_user.py index ba96bd2b5..2b203e59b 100644 --- a/agents-api/agents_api/models/user/create_user.py +++ b/agents-api/agents_api/models/user/create_user.py @@ -12,7 +12,7 @@ from pydantic import ValidationError from ...autogen.openapi_model import CreateUserRequest, User -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -57,7 +57,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("create_user") +@query_metrics_update("create_user") @beartype def create_user( *, diff --git a/agents-api/agents_api/models/user/patch_user.py b/agents-api/agents_api/models/user/patch_user.py index bd3fc0246..ee91bfbdc 100644 --- a/agents-api/agents_api/models/user/patch_user.py +++ b/agents-api/agents_api/models/user/patch_user.py @@ -11,7 +11,7 @@ from ...autogen.openapi_model import PatchUserRequest, ResourceUpdatedResponse from ...common.utils.cozo import cozo_process_mutate_data from ...common.utils.datetime import utcnow -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -51,7 +51,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("patch_user") +@query_metrics_update("patch_user") @beartype def patch_user( *, diff --git a/agents-api/agents_api/models/user/update_user.py b/agents-api/agents_api/models/user/update_user.py index 68e6e6c25..26d3e87a6 100644 --- a/agents-api/agents_api/models/user/update_user.py +++ b/agents-api/agents_api/models/user/update_user.py @@ -8,7 +8,7 @@ from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateUserRequest from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter +from ...metrics.counters import query_metrics_update from ..utils import ( cozo_query, partialclass, @@ -48,7 +48,7 @@ _kind="inserted", ) @cozo_query -@increase_counter("update_user") +@query_metrics_update("update_user") @beartype def update_user( *, developer_id: UUID, user_id: UUID, data: UpdateUserRequest From f47cb924821267e9a9cb24147813292dfae70466 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Thu, 12 Dec 2024 16:31:14 +0300 Subject: [PATCH 4/4] fix: Fix and refactor metrics --- agents-api/agents_api/metrics/counters.py | 100 ++++++++++++++++------ 1 file changed, 75 insertions(+), 25 deletions(-) diff --git a/agents-api/agents_api/metrics/counters.py b/agents-api/agents_api/metrics/counters.py index ec3f31436..c5455248b 100644 --- a/agents-api/agents_api/metrics/counters.py +++ b/agents-api/agents_api/metrics/counters.py @@ -1,37 +1,87 @@ import inspect +import time from functools import wraps from typing import Awaitable, Callable, ParamSpec, TypeVar -from prometheus_client import Counter, Summary +from prometheus_client import Counter, Histogram, Summary +from prometheus_client.utils import INF P = ParamSpec("P") T = TypeVar("T") -def query_metrics_update(metric_label: str, id_field_name: str = "developer_id"): +id_field_name = "developer_id" +labelnames = ("developer_id", "query_name") +buckets = ( + 0.005, + 0.01, + 0.025, + 0.05, + 0.075, + 0.1, + 0.25, + 0.5, + 0.75, + 1.0, + 2.5, + 5.0, + 7.5, + 10.0, + 15.0, + 20.0, + 25.0, + 30.0, + INF, +) +counter = Counter( + f"db_query_counter", + f"Number of db calls", + labelnames=labelnames, +) +summary = Summary( + f"db_query_latency_summary", + "Database query latency summary", + labelnames=labelnames, +) +hist = Histogram( + f"db_query_latency_hist", + "Database query latency histogram", + labelnames=labelnames, + buckets=buckets, +) + + +def query_metrics_update(metric_label: str): def decor(func: Callable[P, T | Awaitable[T]]): - counter = Counter( - f"cozo_counter_{metric_label}", - f"Number of {metric_label} calls", - labelnames=(id_field_name,), - ) - summary = Summary( - f"cozo_latency_{metric_label}", - metric_label, - ) - - @summary.time() - @wraps(func) - def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: - counter.labels(kwargs.get(id_field_name, "not_set")).inc() - return func(*args, **kwargs) - - @summary.time() - @wraps(func) - async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: - counter.labels(kwargs.get(id_field_name, "not_set")).inc() - return await func(*args, **kwargs) - - return async_wrapper if inspect.iscoroutinefunction(func) else wrapper + if inspect.iscoroutinefunction(func): + + @wraps(func) + async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: + fld_id = kwargs.get(id_field_name, "not_set") + labels_dict = {id_field_name: fld_id, "query_name": metric_label} + counter.labels(**labels_dict).inc() + start_time = time.time() + result = await func(*args, **kwargs) + end_time = time.time() - start_time + summary.labels(**labels_dict).observe(end_time) + hist.labels(**labels_dict).observe(end_time) + return result + + return async_wrapper + else: + + @wraps(func) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: + fld_id = kwargs.get(id_field_name, "not_set") + labels_dict = {id_field_name: fld_id, "query_name": metric_label} + counter.labels(**labels_dict).inc() + start_time = time.time() + result = func(*args, **kwargs) + end_time = time.time() - start_time + summary.labels(**labels_dict).observe(end_time) + hist.labels(**labels_dict).observe(end_time) + return result + + return wrapper return decor