add ui improvements to embed and transformation dialogs

open_notebook/graphs/source.py

@@ -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("transform_content", END)
-workflow.add_edge("embed_content", END)
 
 # Compile the graph
 source_graph = workflow.compile()

pages/components/source_panel.py

@@ -44,19 +44,20 @@ def source_panel(source_id: str, modal=False):
 
         with c2:
             transformations = Transformation.get_all()
-            transformation = st.selectbox(
+            with st.container(border=True):
-                "Run a transformation",
+                transformation = st.selectbox(
-                transformations["source_insights"],
+                    "Run a transformation",
-                key=f"transformation_{source.id}",
+                    transformations["source_insights"],
-                format_func=lambda x: x["name"],
+                    key=f"transformation_{source.id}",
-            )
+                    format_func=lambda x: x["name"],
-            st.caption(transformation["description"])
-            if st.button("Run"):
-                result = run_patterns(source.full_text, transformation["patterns"])
-                source.add_insight(
-                    transformation["insight_type"], surreal_clean(result)
                 )
-                st.rerun(scope="fragment" if modal else "app")
+                st.caption(transformation["description"])
+                if st.button("Run"):
+                    result = run_patterns(source.full_text, transformation["patterns"])
+                    source.add_insight(
+                        transformation["insight_type"], surreal_clean(result)
+                    )
+                    st.rerun(scope="fragment" if modal else "app")
 
             if source.embedded_chunks == 0 and st.button(
                 "Embed vectors",

pages/stream_app/source.py

@@ -47,7 +47,10 @@ def add_source(notebook_id):
         options=available_transformations,
         default=default_transformations,
     )
-    embed = st.checkbox("Embed content for vector search", value=False)
+    run_embed = st.checkbox(
+        "Embed content for vector search",
+        help="Creates an embedded content for vector search. Costs a little money and takes a little bit more time. You can do this later if you prefer.",
+    )
     if st.button("Process", key="add_source"):
         logger.debug("Adding source")
         with st.status("Processing...", expanded=True):
@@ -78,7 +81,7 @@ def add_source(notebook_id):
                             "content_state": req,
                             "notebook_id": notebook_id,
                             "transformations": apply_transformations,
-                            "embed": embed,
+                            "embed": run_embed,
                         }
                     )
                 )