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()

View file

@ -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",

View file

@ -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,
}
)
)