don't write tag classes, create new ones on read (#905)

Author: dinmukhamedm

app-server/src/api/v1/tag.rs

@@ -1,12 +1,11 @@
 use std::sync::Arc;
 
 use crate::{
-    ch::spans::append_tags_to_span,
-    db::{DB, project_api_keys::ProjectApiKey, tags::TagSource},
+    ch::{spans::append_tags_to_span, tags::insert_tag},
+    db::{project_api_keys::ProjectApiKey, tags::TagSource},
     query_engine::QueryEngine,
     routes::types::ResponseResult,
     sql::{self, ClickhouseReadonlyClient},
-    tags::create_tag,
 };
 use actix_web::{
     HttpResponse, post,
@@ -41,7 +40,6 @@ pub enum TagRequest {
 #[post("tag")]
 pub async fn tag_trace(
     req: Json<TagRequest>,
-    db: web::Data<DB>,
     clickhouse: web::Data<clickhouse::Client>,
     clickhouse_ro: web::Data<Option<Arc<ClickhouseReadonlyClient>>>,
     query_engine: web::Data<Arc<QueryEngine>>,
@@ -90,13 +88,12 @@ pub async fn tag_trace(
     let futures = names
         .iter()
         .map(|name| {
-            create_tag(
-                &db.pool,
+            insert_tag(
                 clickhouse.clone(),
                 project_api_key.project_id,
-                span_id,
                 name.clone(),
                 TagSource::CODE,
+                span_id,
             )
         })
         .collect::<Vec<_>>();

app-server/src/ch/tags.rs

@@ -53,19 +53,19 @@ impl CHTag {
 pub async fn insert_tag(
     client: clickhouse::Client,
     project_id: Uuid,
-    id: Uuid,
     name: String,
     source: TagSource,
     span_id: Uuid,
-) -> Result<()> {
+) -> Result<Uuid> {
+    let id = Uuid::new_v4();
     let tag = CHTag::new(project_id, id, name, source, span_id);
     let ch_insert = client.insert("tags");
     match ch_insert {
         Ok(mut ch_insert) => {
             ch_insert.write(&tag).await?;
             let ch_insert_end_res = ch_insert.end().await;
             match ch_insert_end_res {
-                Ok(_) => Ok(()),
+                Ok(_) => Ok(id),
                 Err(e) => {
                     return Err(anyhow::anyhow!("Clickhouse tag insertion failed: {:?}", e));
                 }

app-server/src/db/tags.rs

@@ -1,7 +1,6 @@
-use anyhow::Result;
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
-use sqlx::{FromRow, PgPool};
+use sqlx::FromRow;
 use uuid::Uuid;
 
 #[derive(sqlx::Type, Serialize, Deserialize, Clone, PartialEq)]
@@ -50,18 +49,3 @@ pub struct SpanTag {
     pub user_id: Option<Uuid>,
     pub user_email: Option<String>,
 }
-
-pub async fn insert_tag_class(pool: &PgPool, project_id: Uuid, tag_name: &String) -> Result<()> {
-    sqlx::query(
-        "
-    INSERT INTO tag_classes (project_id, name)
-    VALUES ($1, $2)
-    ON CONFLICT (project_id, name) DO NOTHING
-    ",
-    )
-    .bind(project_id)
-    .bind(tag_name)
-    .execute(pool)
-    .await?;
-    Ok(())
-}

app-server/src/main.rs

@@ -60,7 +60,6 @@ mod db;
 mod evaluations;
 mod evaluators;
 mod features;
-mod tags;
 mod language_model;
 mod mq;
 mod names;

app-server/src/tags/mod.rs

@@ -28,7 +28,7 @@ use crate::{
         events::record_events,
         limits::update_workspace_limit_exceeded_by_project_id,
         provider::convert_span_to_provider_format,
-        utils::{get_llm_usage_for_span, prepare_span_for_recording, record_tags_to_db_and_ch},
+        utils::{get_llm_usage_for_span, prepare_span_for_recording, record_tags},
     },
 };
 
@@ -346,8 +346,7 @@ async fn process_batch(
     }
 
     for span in &stripped_spans {
-        if let Err(e) = record_tags_to_db_and_ch(
-            db.clone(),
+        if let Err(e) = record_tags(
             clickhouse.clone(),
             &span.tags,
             &span.span_id,

app-server/src/traces/consumer.rs

@@ -95,8 +95,7 @@ pub async fn get_llm_usage_for_span(
     }
 }
 
-pub async fn record_tags_to_db_and_ch(
-    db: Arc<DB>,
+pub async fn record_tags(
     clickhouse: clickhouse::Client,
     tags: &[String],
     span_id: &Uuid,
@@ -107,13 +106,12 @@ pub async fn record_tags_to_db_and_ch(
     }
 
     for tag_name in tags {
-        crate::tags::create_tag(
-            &db.pool,
+        crate::ch::tags::insert_tag(
             clickhouse.clone(),
             *project_id,
-            *span_id,
             tag_name.clone(),
             TagSource::CODE,
+            *span_id,
         )
         .await?;
     }

app-server/src/traces/utils.rs

@@ -24,7 +24,7 @@ export async function POST(
   const projectId = params.projectId;
   const spanId = params.spanId;
 
-  const body = (await req.json()) as { classId: string; name: string };
+  const body = (await req.json()) as { name: string };
 
   const res = await addSpanTag({
     spanId,

frontend/app/api/projects/[projectId]/spans/[spanId]/tags/route.ts

@@ -1,5 +1,7 @@
 import { and, eq } from 'drizzle-orm';
+import { prettifyError, ZodError } from 'zod';
 
+import { createOrUpdateTagClass } from '@/lib/actions/tags';
 import { db } from '@/lib/db/drizzle';
 import { tagClasses } from '@/lib/db/migrations/schema';
 
@@ -11,21 +13,22 @@ export async function POST(
   const params = await props.params;
   const projectId = params.projectId;
   const tagName = params.tagName;
-
   const body = await req.json();
 
-  const result = await db.update(tagClasses).set({
-    name: body.name,
-    color: body.color,
-  }).where(
-    and(eq(tagClasses.name, tagName), eq(tagClasses.projectId, projectId))
-  ).returning();
-
-  if (result.length === 0) {
-    return new Response('Tag class not found', { status: 404 });
+  try {
+    const result = await createOrUpdateTagClass({
+      projectId,
+      name: tagName,
+      color: body.color,
+    });
+    return Response.json(result, { status: 200 });
+  } catch (error) {
+    if (error instanceof ZodError) {
+      return Response.json({ error: prettifyError(error) }, { status: 400 });
+    }
+
+    return Response.json({ error: "Internal server error" }, { status: 500 });
   }
-
-  return Response.json(result[0], { status: 200 });
 }
 
 export async function DELETE(

frontend/app/api/projects/[projectId]/tag-classes/[tagName]/route.ts

@@ -15,25 +15,3 @@ export async function GET(req: Request, props: { params: Promise<{ projectId: st
 
   return new Response(JSON.stringify(res), { status: 200 });
 }
-
-export async function POST(req: Request, props: { params: Promise<{ projectId: string }> }): Promise<Response> {
-  const params = await props.params;
-  const projectId = params.projectId;
-
-  const body = await req.json();
-
-  const res = await db
-    .insert(tagClasses)
-    .values({
-      projectId,
-      name: body.name,
-      color: body.color,
-    })
-    .returning();
-
-  if (res.length === 0) {
-    return new Response(JSON.stringify({ error: "Failed to create tag class" }), { status: 500 });
-  }
-
-  return new Response(JSON.stringify(res[0]), { status: 200 });
-}