test iterate chain

planning_ai/chains/iterate_chain.py

@@ -0,0 +1,100 @@
+from typing import List, Literal, TypedDict
+
+from langchain_community.document_loaders import DirectoryLoader, TextLoader
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableConfig
+from langgraph.graph import END, START, StateGraph
+
+from planning_ai.common.utils import Paths
+from planning_ai.llms.llm import LLM
+
+# Initial summary
+summarize_prompt = ChatPromptTemplate(
+    [
+        ("human", "Write a concise summary of the following: {context}"),
+    ]
+)
+initial_summary_chain = summarize_prompt | LLM | StrOutputParser()
+
+refine_template = """
+Produce a final summary.
+
+Existing summary up to this point:
+{existing_answer}
+
+New context:
+------------
+{context}
+------------
+
+Given the new context, refine the original summary.
+"""
+refine_prompt = ChatPromptTemplate([("human", refine_template)])
+
+refine_summary_chain = refine_prompt | LLM | StrOutputParser()
+
+
+# We will define the state of the graph to hold the document
+# contents and summary. We also include an index to keep track
+# of our position in the sequence of documents.
+class State(TypedDict):
+    contents: List[str]
+    index: int
+    summary: str
+
+
+# We define functions for each node, including a node that generates
+# the initial summary:
+async def generate_initial_summary(state: State, config: RunnableConfig):
+    summary = await initial_summary_chain.ainvoke(
+        state["contents"][0],
+        config,
+    )
+    return {"summary": summary, "index": 1}
+
+
+# And a node that refines the summary based on the next document
+async def refine_summary(state: State, config: RunnableConfig):
+    content = state["contents"][state["index"]]
+    summary = await refine_summary_chain.ainvoke(
+        {"existing_answer": state["summary"], "context": content},
+        config,
+    )
+
+    return {"summary": summary, "index": state["index"] + 1}
+
+
+# Here we implement logic to either exit the application or refine
+# the summary.
+def should_refine(state: State) -> Literal["refine_summary", END]:
+    if state["index"] >= len(state["contents"]):
+        return END
+    else:
+        return "refine_summary"
+
+
+graph = StateGraph(State)
+graph.add_node("generate_initial_summary", generate_initial_summary)
+graph.add_node("refine_summary", refine_summary)
+
+graph.add_edge(START, "generate_initial_summary")
+graph.add_conditional_edges("generate_initial_summary", should_refine)
+graph.add_conditional_edges("refine_summary", should_refine)
+app = graph.compile()
+
+loader = DirectoryLoader(
+    path=str(Paths.STAGING),
+    show_progress=True,
+    use_multithreading=True,
+    loader_cls=TextLoader,
+    recursive=True,
+)
+docs = [doc for doc in loader.load()[:20] if doc.page_content]
+
+async for step in app.astream(
+    {"contents": [doc.page_content for doc in docs]},
+    stream_mode="values",
+):
+    if summary := step.get("summary"):
+        print(summary)