add ui improvements to embed and transformation dialogs

This commit is contained in:
LUIS NOVO 2024-11-11 18:17:08 -03:00
parent 532e606a49
commit 8cb6d835fe
3 changed files with 25 additions and 26 deletions

View file

@ -23,7 +23,7 @@ class SourceState(TypedDict):
notebook_id: str
source: Source
transformations: Annotated[list, operator.add]
embed: bool = False
embed: bool
class TransformationState(TypedDict):
@ -61,6 +61,11 @@ def save_source(state: SourceState) -> dict:
if state["notebook_id"]:
logger.debug(f"Adding source to notebook {state['notebook_id']}")
source.add_to_notebook(state["notebook_id"])
if state["embed"]:
logger.debug("Embedding content for vector search")
source.vectorize()
return {"source": source}
@ -103,14 +108,6 @@ async def transform_content(state: TransformationState) -> dict:
return {"transformations": [{"name": transformation["name"], "content": result}]}
async def embed_content(state: SourceState) -> dict:
source: Source = state["source"]
if state["embed"]:
logger.debug("Embedding content for vector search")
source.vectorize()
return {"source": source}
# Create and compile the workflow
workflow = StateGraph(SourceState)
@ -118,15 +115,13 @@ workflow = StateGraph(SourceState)
workflow.add_node("content_process", content_process)
workflow.add_node("save_source", save_source)
workflow.add_node("transform_content", transform_content)
workflow.add_node("embed_content", embed_content)
# Define the graph edges
workflow.add_edge(START, "content_process")
workflow.add_edge("content_process", "save_source")
workflow.add_conditional_edges(
"save_source", trigger_transformations, ["transform_content"]
)
workflow.add_edge("transform_content", "embed_content")
workflow.add_edge("embed_content", END)
workflow.add_edge("transform_content", END)
# Compile the graph
source_graph = workflow.compile()