Merge branch 'dev' into f/chat-mmr

julep-ai · Dec 12, 2024 · b14f82e · b14f82e
2 parents d43cc8b + 682da5a
commit b14f82e
Show file tree

Hide file tree

Showing 29 changed files with 165 additions and 77 deletions.
diff --git a/agents-api/agents_api/metrics/counters.py b/agents-api/agents_api/metrics/counters.py
@@ -1,31 +1,87 @@
 import inspect
+import time
 from functools import wraps
 from typing import Awaitable, Callable, ParamSpec, TypeVar
 
-from prometheus_client import Counter
+from prometheus_client import Counter, Histogram, Summary
+from prometheus_client.utils import INF
 
 P = ParamSpec("P")
 T = TypeVar("T")
 
 
-def increase_counter(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]]):
-        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)
-
-        @wraps(func)
-        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
-            metric.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
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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/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
@@ -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/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
@@ -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
@@ -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
@@ -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
@@ -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(
     *,
@@ -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/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
@@ -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(
     *,