diff --git a/open_notebook/graphs/source.py b/open_notebook/graphs/source.py index dbce6da..9456bd7 100644 --- a/open_notebook/graphs/source.py +++ b/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("embed_content", END) +workflow.add_edge("transform_content", END) # Compile the graph source_graph = workflow.compile() diff --git a/pages/components/source_panel.py b/pages/components/source_panel.py index 83f2304..55f09df 100644 --- a/pages/components/source_panel.py +++ b/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( - "Run a transformation", - transformations["source_insights"], - 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) + with st.container(border=True): + transformation = st.selectbox( + "Run a transformation", + transformations["source_insights"], + key=f"transformation_{source.id}", + format_func=lambda x: x["name"], ) - 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", diff --git a/pages/stream_app/source.py b/pages/stream_app/source.py index d37efcc..d8b04f5 100644 --- a/pages/stream_app/source.py +++ b/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, } ) )