From 4b1441514a4a03a387a3be41d2ffb630780a05eb Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Fri, 22 Mar 2024 15:25:21 -0400
Subject: [PATCH] update run names (#102)

Update run names for suggest
---
 backend/server/api/suggest.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/backend/server/api/suggest.py b/backend/server/api/suggest.py
index c7551e2..93b7705 100644
--- a/backend/server/api/suggest.py
+++ b/backend/server/api/suggest.py
@@ -72,7 +72,7 @@ class ExtractorDefinition(BaseModel):
 
 suggestion_chain = SUGGEST_PROMPT | model.with_structured_output(
     schema=ExtractorDefinition
-)
+).with_config({"run_name": "suggest"})
 
 UPDATE_PROMPT = ChatPromptTemplate.from_messages(
     [
@@ -94,7 +94,9 @@ class ExtractorDefinition(BaseModel):
     ]
 )
 
-UPDATE_CHAIN = UPDATE_PROMPT | model.with_structured_output(schema=ExtractorDefinition)
+UPDATE_CHAIN = (
+    UPDATE_PROMPT | model.with_structured_output(schema=ExtractorDefinition)
+).with_config({"run_name": "suggest_update"})
 
 
 # PUBLIC API
@@ -104,7 +106,6 @@ class ExtractorDefinition(BaseModel):
 async def suggest(request: SuggestExtractor) -> ExtractorDefinition:
     """Endpoint to create an extractor."""
     if len(request.jsonSchema) > 10:
-        print(f"Using update chain with {request.jsonSchema}")
         return await UPDATE_CHAIN.ainvoke(
             {"input": request.description, "json_schema": request.jsonSchema}
         )