Merge pull request #1484 from Agenta-AI/v1/observability-sdk

Observability

.gitignore

@@ -59,4 +59,10 @@ agenta-web/cypress/videos/
 rabbitmq_data
 
 # docker compose override
-docker-compose.*override.yaml
+docker-compose.*override.yaml
+agenta-backend/agenta_backend/migrations/*.sh
+agenta-backend/agenta_backend/migrations/migrations/v0_11_0_to_v0_12_0/20240131194122_create_default_workspace_for_org.py
+agenta-backend/agenta_backend/migrations/migrations/v0_11_0_to_v0_12_0/20240203035221_move_org_existing_members_to_workspace.py
+agenta-backend/agenta_backend/migrations/migrations/v0_11_0_to_v0_12_0/20240201102156_update_org_invitation.py
+agenta-backend/agenta_backend/migrations/migrations/v0_11_0_to_v0_12_0/20240203035520_add_workspace_field_to_db_models.py
+docker-compose.oss.prod.yml

agenta-backend/agenta_backend/__init__.py

@@ -1,6 +1,6 @@
 import os
 
-if os.environ["FEATURE_FLAG"] in ["cloud"]:
+if os.environ["FEATURE_FLAG"] in ["cloud", "cloud-dev"]:
     import agenta_backend.cloud.__init__
 if os.environ["FEATURE_FLAG"] in ["ee"]:
     import agenta_backend.ee.__init__

agenta-backend/agenta_backend/main.py

@@ -10,21 +10,20 @@
     evaluation_router,
     human_evaluation_router,
     evaluators_router,
-    observability_router,
     testset_router,
     user_profile,
     variants_router,
     bases_router,
     configs_router,
     health_router,
 )
-from agenta_backend.utils.common import isCloudEE
+from agenta_backend.utils.common import isEE, isCloudProd, isCloudDev, isOss, isCloudEE
 from agenta_backend.models.db_engine import DBEngine
 from agenta_backend.open_api import open_api_tags_metadata
 
-if isCloudEE():
+if isEE() or isCloudProd():
     from agenta_backend.commons.services import templates_manager
-else:
+elif isCloudDev() or isOss():
     from agenta_backend.services import templates_manager
 
 from fastapi import FastAPI
@@ -100,9 +99,6 @@ async def lifespan(application: FastAPI, cache=True):
 app.include_router(
     environment_router.router, prefix="/environments", tags=["Environments"]
 )
-app.include_router(
-    observability_router.router, prefix="/observability", tags=["Observability"]
-)
 app.include_router(bases_router.router, prefix="/bases", tags=["Bases"])
 app.include_router(configs_router.router, prefix="/configs", tags=["Configs"])
 

agenta-backend/agenta_backend/models/api/api_models.py

@@ -1,9 +1,26 @@
+from enum import Enum
 from datetime import datetime
-from agenta_backend.models.db_models import ConfigDB
 from typing import Any, Dict, List, Optional
-from enum import Enum
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+
+from agenta_backend.models.db_models import ConfigDB
+
+
+class PaginationParam(BaseModel):
+    page: int = Field(default=1, ge=1)
+    pageSize: int = Field(default=10, ge=1)
+
+
+class SorterParams(BaseModel):
+    created_at: str = Field("desc")
+
+
+class WithPagination(BaseModel):
+    data: List[Any]
+    total: int
+    page: int
+    pageSize: int
 
 
 class Error(BaseModel):
@@ -18,7 +35,6 @@ class Result(BaseModel):
 
 
 class GetConfigResponse(BaseModel):
-    config_id: Optional[str]
     config_name: str
     current_version: int
     parameters: Dict[str, Any]

agenta-backend/agenta_backend/models/converters.py

@@ -3,16 +3,11 @@
 
 import json
 import logging
-from typing import List
+from typing import List, Tuple, Any
 
 from agenta_backend.services import db_manager
 from agenta_backend.utils.common import isCloudEE
 from agenta_backend.models.api.user_models import User
-from agenta_backend.models.api.observability_models import (
-    Span,
-    Trace,
-    Feedback as FeedbackOutput,
-)
 from agenta_backend.models.api.evaluation_model import (
     Evaluation,
     HumanEvaluation,
@@ -74,12 +69,9 @@
     )
 
 from agenta_backend.models.db_models import (
-    SpanDB,
-    TraceDB,
     TemplateDB,
     AggregatedResult,
     AppVariantRevisionsDB,
-    Feedback as FeedbackDB,
     EvaluationScenarioResult,
 )
 from agenta_backend.models.api.api_models import (
@@ -89,9 +81,12 @@
     TestSetOutput,
     TemplateImageInfo,
     AppVariantRevision,
+    PaginationParam,
+    WithPagination,
 )
 
-from beanie import Link
+from fastapi import Depends
+from beanie import Link, PydanticObjectId as ObjectId
 
 
 logger = logging.getLogger(__name__)
@@ -491,68 +486,6 @@ def testset_db_to_pydantic(test_set_db: TestSetDB) -> TestSetOutput:
     )
 
 
-def spans_db_to_pydantic(spans_db: List[SpanDB]) -> List[Span]:
-    return [
-        Span(
-            span_id=str(span_db.id),
-            parent_span_id=str(span_db.parent_span_id),
-            meta=span_db.meta,
-            event_name=span_db.event_name,
-            event_type=span_db.event_type,
-            start_time=span_db.start_time,
-            duration=span_db.duration,
-            status=span_db.status,
-            end_time=span_db.end_time,
-            inputs=span_db.inputs,
-            outputs=span_db.outputs,
-            prompt_template=span_db.prompt_template,
-            tokens_input=span_db.tokens_input,
-            tokens_output=span_db.tokens_output,
-            token_total=span_db.token_total,
-            cost=span_db.cost,
-            tags=span_db.tags,
-        ).dict(exclude_unset=True)
-        for span_db in spans_db
-    ]
-
-
-def feedback_db_to_pydantic(feedback_db: FeedbackDB) -> FeedbackOutput:
-    return FeedbackOutput(
-        feedback_id=str(feedback_db.uid),
-        feedback=feedback_db.feedback,
-        score=feedback_db.score,
-        meta=feedback_db.meta,
-        created_at=feedback_db.created_at,
-    ).dict(exclude_unset=True)
-
-
-def trace_db_to_pydantic(trace_db: TraceDB) -> Trace:
-    feedbacks = trace_db.feedbacks
-    if feedbacks is None:
-        result = []
-    else:
-        result = [
-            feedback_db_to_pydantic(feedback)
-            for feedback in feedbacks
-            if feedback is not None
-        ]
-
-    return Trace(
-        trace_id=str(trace_db.id),
-        app_id=trace_db.app_id,
-        variant_id=trace_db.variant_id,
-        cost=trace_db.cost,
-        latency=trace_db.latency,
-        status=trace_db.status,
-        token_consumption=trace_db.token_consumption,
-        tags=trace_db.tags,
-        start_time=trace_db.start_time,
-        end_time=trace_db.end_time,
-        feedbacks=result,
-        spans=[str(span) for span in trace_db.spans],
-    ).dict(exclude_unset=True)
-
-
 def user_db_to_pydantic(user_db: UserDB) -> User:
     return User(
         id=str(user_db.id),
@@ -571,3 +504,17 @@ def evaluator_config_db_to_pydantic(evaluator_config: EvaluatorConfigDB):
         created_at=evaluator_config.created_at,
         updated_at=evaluator_config.updated_at,
     )
+
+
+def get_paginated_data(
+    data: List[Any], data_count: int, query: PaginationParam = Depends()
+):
+    return WithPagination(
+        data=data, total=data_count, page=query.page, pageSize=query.pageSize
+    )
+
+
+def get_pagination_skip_limit(pagination: PaginationParam) -> Tuple[int, int]:
+    skip = (pagination.page - 1) * pagination.pageSize
+    limit = pagination.pageSize
+    return skip, limit

agenta-backend/agenta_backend/models/db_engine.py

@@ -9,6 +9,7 @@
 from agenta_backend.utils.common import isCloudEE
 
 if isCloudEE():
+    from agenta_backend.commons.observability.models.db import SpanDB
     from agenta_backend.commons.models.db_models import (
         APIKeyDB,
         WorkspaceDB,
@@ -47,8 +48,6 @@
     )
 
 from agenta_backend.models.db_models import (
-    SpanDB,
-    TraceDB,
     TemplateDB,
     AppVariantRevisionsDB,
 )
@@ -57,8 +56,6 @@
 document_models: List[Document] = [
     AppDB,
     UserDB,
-    SpanDB,
-    TraceDB,
     ImageDB,
     TestSetDB,
     TemplateDB,
@@ -76,7 +73,7 @@
 ]
 
 if isCloudEE():
-    document_models = document_models + [OrganizationDB, WorkspaceDB, APIKeyDB]
+    document_models = document_models + [SpanDB, OrganizationDB, WorkspaceDB, APIKeyDB]
 
 
 # Configure and set logging level

agenta-backend/agenta_backend/models/db_models.py

@@ -1,5 +1,3 @@
-from enum import Enum
-from uuid import uuid4
 from datetime import datetime
 from typing import Any, Dict, List, Optional
 
@@ -297,53 +295,3 @@ class EvaluationScenarioDB(Document):
 
     class Settings:
         name = "new_evaluation_scenarios"
-
-
-class SpanDB(Document):
-    parent_span_id: Optional[str]
-    meta: Optional[Dict[str, Any]]
-    event_name: str  # Function or execution name
-    event_type: Optional[str]
-    start_time: datetime
-    duration: Optional[int]
-    status: str  # initiated, completed, stopped, cancelled
-    end_time: datetime = Field(default=datetime.now())
-    inputs: Optional[List[str]]
-    outputs: Optional[List[str]]
-    prompt_template: Optional[str]
-    tokens_input: Optional[int]
-    tokens_output: Optional[int]
-    token_total: Optional[int]
-    cost: Optional[float]
-    tags: Optional[List[str]]
-
-    class Settings:
-        name = "spans"
-
-
-class Feedback(BaseModel):
-    uid: str = Field(default=str(uuid4()))
-    user_id: str
-    feedback: Optional[str]
-    score: Optional[float]
-    meta: Optional[Dict[str, Any]]
-    created_at: datetime
-    updated_at: datetime = Field(default=datetime.now())
-
-
-class TraceDB(Document):
-    app_id: Optional[str]
-    variant_id: str
-    spans: List[PydanticObjectId]
-    start_time: datetime
-    end_time: datetime = Field(default=datetime.now())
-    cost: Optional[float]
-    latency: float
-    status: str  # initiated, completed, stopped, cancelled, failed
-    token_consumption: Optional[int]
-    user: Link[UserDB]
-    tags: Optional[List[str]]
-    feedbacks: Optional[List[Feedback]]
-
-    class Settings:
-        name = "traces"