v1 of transformations

open_notebook/graphs/multipattern.py

@@ -0,0 +1,66 @@
+import operator
+import os
+from typing import List, Literal, Sequence
+
+from langchain_core.runnables import (
+    RunnableConfig,
+)
+from langgraph.graph import END, START, StateGraph
+from loguru import logger
+from typing_extensions import Annotated, TypedDict
+
+from open_notebook.graphs.utils import run_pattern
+
+
+class PatternChainState(TypedDict):
+    content_stack: Annotated[Sequence[str], operator.add]
+    transformations: List[str]
+    output: str
+
+
+def call_model(state: dict, config: RunnableConfig) -> dict:
+    model_name = config.get("configurable", {}).get(
+        "model_name", os.environ.get("DEFAULT_MODEL")
+    )
+    transformations = state["transformations"]
+    current_transformation = transformations.pop(0)
+    if current_transformation.startswith("patterns/"):
+        input_args = {"input_text": state["content_stack"][-1]}
+    else:
+        input_args = {
+            "input_text": state["content_stack"][-1],
+            "command": current_transformation,
+        }
+        current_transformation = "patterns/custom"
+
+    logger.warning(f"Processing transformation: {current_transformation}")
+    logger.debug(f"Using input: {input_args}")
+    transformation_result = run_pattern(
+        pattern_name=current_transformation,
+        model_name=model_name,
+        state=input_args,
+    )
+    return {
+        "content_stack": [transformation_result.content],
+        "output": transformation_result.content,
+        "transformations": state["transformations"],
+    }
+
+
+def transform_condition(state: PatternChainState) -> Literal["agent", END]:  # type: ignore
+    """
+    Checks whether there are more chunks to process.
+    """
+    if len(state["transformations"]) > 0:
+        return "agent"
+    return END
+
+
+agent_state = StateGraph(PatternChainState)
+agent_state.add_node("agent", call_model)
+agent_state.add_edge(START, "agent")
+agent_state.add_conditional_edges(
+    "agent",
+    transform_condition,
+)
+graph = agent_state.compile()