Add queries and tokens metrics

agents-api/agents_api/metrics/counters.py

@@ -0,0 +1,25 @@
+from functools import wraps
+from typing import Callable, ParamSpec, TypeVar
+
+from prometheus_client import Counter
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def increase_counter(metric_label: str, id_field_name: str = "developer_id"):
+    def decor(func: Callable[P, T] | None = None):
+        metric = Counter(
+            metric_label,
+            f"Number of {metric_label} calls",
+            labelnames=(id_field_name,),
+        )
+
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            metric.labels(kwargs.get(id_field_name, "not_set")).inc()
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decor

agents-api/agents_api/models/agent/create_agent.py

@@ -13,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -57,6 +58,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("create_agent")
 @beartype
 def create_agent(
     *,

agents-api/agents_api/models/agent/create_or_update_agent.py

@@ -13,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -48,6 +49,7 @@
     Agent, one=True, transform=lambda d: {"id": UUID(d.pop("agent_id")), **d}
 )
 @cozo_query
+@increase_counter("create_or_update_agent")
 @beartype
 def create_or_update_agent(
     *,

agents-api/agents_api/models/agent/patch_agent.py

@@ -9,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -36,6 +37,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("patch_agent")
 @beartype
 def patch_agent(
     *,

agents-api/agents_api/models/agent/update_agent.py

@@ -8,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -35,6 +36,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("update_agent")
 @beartype
 def update_agent(
     *,

agents-api/agents_api/models/docs/create_doc.py

@@ -8,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -37,6 +38,7 @@
     },
 )
 @cozo_query
+@increase_counter("create_doc")
 @beartype
 def create_doc(
     *,

agents-api/agents_api/models/entry/create_entries.py

@@ -10,6 +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 ..utils import (
     cozo_query,
     mark_session_updated_query,
@@ -40,6 +41,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("create_entries")
 @beartype
 def create_entries(
     *,

agents-api/agents_api/models/execution/create_execution.py

@@ -9,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -36,6 +37,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("create_execution")
 @beartype
 def create_execution(
     *,

agents-api/agents_api/models/execution/create_execution_transition.py

@@ -12,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -72,6 +73,7 @@ def validate_transition_targets(data: CreateTransitionRequest) -> None:
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("create_execution_transition")
 @beartype
 def create_execution_transition(
     *,

agents-api/agents_api/models/execution/create_temporal_lookup.py

@@ -8,6 +8,7 @@
 from temporalio.client import WorkflowHandle
 
 from ...common.utils.cozo import cozo_process_mutate_data
+from ...metrics.counters import increase_counter
 from ..utils import (
     cozo_query,
     partialclass,
@@ -28,6 +29,7 @@
     }
 )
 @cozo_query
+@increase_counter("create_temporal_lookup")
 @beartype
 def create_temporal_lookup(
     *,

agents-api/agents_api/models/execution/update_execution.py

@@ -14,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -41,6 +42,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("update_execution")
 @beartype
 def update_execution(
     *,

agents-api/agents_api/models/session/create_or_update_session.py

@@ -11,6 +11,7 @@
     ResourceUpdatedResponse,
 )
 from ...common.utils.cozo import cozo_process_mutate_data
+from ...metrics.counters import increase_counter
 from ..utils import (
     cozo_query,
     partialclass,
@@ -43,6 +44,7 @@
     },
 )
 @cozo_query
+@increase_counter("create_or_update_session")
 @beartype
 def create_or_update_session(
     *,

agents-api/agents_api/models/session/create_session.py

@@ -12,6 +12,7 @@
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import CreateSessionRequest, Session
+from ...metrics.counters import increase_counter
 from ..utils import (
     cozo_query,
     partialclass,
@@ -44,6 +45,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("create_session")
 @beartype
 def create_session(
     *,

agents-api/agents_api/models/session/update_session.py

@@ -8,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -51,6 +52,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("update_session")
 @beartype
 def update_session(
     *,

agents-api/agents_api/models/task/create_or_update_task.py

@@ -18,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -49,6 +50,7 @@
     },
 )
 @cozo_query
+@increase_counter("create_or_update_task")
 @beartype
 def create_or_update_task(
     *,

agents-api/agents_api/models/task/create_task.py

@@ -16,6 +16,7 @@
 )
 from ...common.protocol.tasks import spec_to_task, task_to_spec
 from ...common.utils.cozo import cozo_process_mutate_data
+from ...metrics.counters import increase_counter
 from ..utils import (
     cozo_query,
     partialclass,
@@ -38,6 +39,7 @@
 )
 @wrap_in_class(spec_to_task, one=True, _kind="inserted")
 @cozo_query
+@increase_counter("create_task")
 @beartype
 def create_task(
     *,

agents-api/agents_api/models/task/patch_task.py

@@ -14,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -46,6 +47,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("patch_task")
 @beartype
 def patch_task(
     *,

agents-api/agents_api/models/task/update_task.py

@@ -14,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -45,6 +46,7 @@
     },
 )
 @cozo_query
+@increase_counter("update_task")
 @beartype
 def update_task(
     *,

agents-api/agents_api/models/tools/create_tools.py

@@ -9,6 +9,7 @@
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import CreateToolRequest, Tool
+from ...metrics.counters import increase_counter
 from ..utils import (
     cozo_query,
     partialclass,
@@ -39,6 +40,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("create_tools")
 @beartype
 def create_tools(
     *,

agents-api/agents_api/models/tools/patch_tool.py

@@ -8,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -35,6 +36,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("patch_tool")
 @beartype
 def patch_tool(
     *, developer_id: UUID, agent_id: UUID, tool_id: UUID, data: PatchToolRequest

agents-api/agents_api/models/tools/update_tool.py

@@ -11,6 +11,7 @@
     UpdateToolRequest,
 )
 from ...common.utils.cozo import cozo_process_mutate_data
+from ...metrics.counters import increase_counter
 from ..utils import (
     cozo_query,
     partialclass,
@@ -38,6 +39,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("update_tool")
 @beartype
 def update_tool(
     *,

agents-api/agents_api/models/user/create_or_update_user.py

@@ -12,6 +12,7 @@
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import CreateOrUpdateUserRequest, User
+from ...metrics.counters import increase_counter
 from ..utils import (
     cozo_query,
     partialclass,
@@ -45,6 +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")
 @beartype
 def create_or_update_user(
     *,

agents-api/agents_api/models/user/create_user.py

@@ -12,6 +12,7 @@
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import CreateUserRequest, User
+from ...metrics.counters import increase_counter
 from ..utils import (
     cozo_query,
     partialclass,
@@ -56,6 +57,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("create_user")
 @beartype
 def create_user(
     *,

agents-api/agents_api/models/user/patch_user.py

@@ -11,6 +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 ..utils import (
     cozo_query,
     partialclass,
@@ -50,6 +51,7 @@
     _kind="inserted",
 )
 @cozo_query
+@increase_counter("patch_user")
 @beartype
 def patch_user(
     *,