remove token limit

planning_ai/common/utils.py

@@ -1,10 +1,6 @@
 from pathlib import Path
-from typing import List
 
 import polars as pl
-from langchain_core.documents import Document
-
-from planning_ai.llms.llm import LLM
 
 pl.Config(
     fmt_str_lengths=9,
@@ -20,12 +16,3 @@ class Paths:
     RAW = DATA / "raw"
     STAGING = DATA / "staging"
     OUT = DATA / "out"
-
-
-class Consts:
-    TOKEN_MAX = 100_000
-
-
-def length_function(documents: List[Document]) -> int:
-    """Get number of tokens for input contents."""
-    return sum(LLM.get_num_tokens(doc.page_content) for doc in documents)

planning_ai/graph.py

@@ -4,9 +4,8 @@ from planning_ai.nodes.map_node import (
     collect_summaries,
     generate_summary,
     map_summaries,
-    should_collapse,
 )
-from planning_ai.nodes.reduce_node import collapse_summaries, generate_final_summary
+from planning_ai.nodes.reduce_node import generate_final_summary
 from planning_ai.states import OverallState
 
 
@@ -14,13 +13,11 @@ def create_graph():
     graph = StateGraph(OverallState)
     graph.add_node("generate_summary", generate_summary)
     graph.add_node("collect_summaries", collect_summaries)
-    graph.add_node("collapse_summaries", collapse_summaries)
     graph.add_node("generate_final_summary", generate_final_summary)
 
     graph.add_conditional_edges(START, map_summaries, ["generate_summary"])
     graph.add_edge("generate_summary", "collect_summaries")
-    graph.add_conditional_edges("collect_summaries", should_collapse)
-    graph.add_conditional_edges("collapse_summaries", should_collapse)
+    graph.add_edge("collect_summaries", "generate_final_summary")
     graph.add_edge("generate_final_summary", END)
 
     return graph.compile()

planning_ai/main.py

@@ -7,8 +7,6 @@ from langchain_text_splitters import CharacterTextSplitter
 from planning_ai.common.utils import Paths
 from planning_ai.graph import create_graph
 
-doc_title = "Cambridge Response Summary"
-
 
 def build_quarto_doc(doc_title, out):
     final = out["generate_final_summary"]
@@ -120,7 +118,7 @@ def main():
         loader_cls=TextLoader,
         recursive=True,
     )
-    docs = [doc for doc in loader.load()[:200] if doc.page_content]
+    docs = [doc for doc in loader.load()[:20] if doc.page_content]
     text_splitter = CharacterTextSplitter.from_tiktoken_encoder(
         chunk_size=1000, chunk_overlap=0
     )
@@ -145,5 +143,6 @@ def main():
 
 
 if __name__ == "__main__":
+    doc_title = "Cambridge Response Summary"
     out = main()
     build_quarto_doc(doc_title, out)

planning_ai/nodes/map_node.py

@@ -1,10 +1,7 @@
-from typing import Literal
-
 from langchain_core.documents import Document
 from langgraph.constants import Send
 
 from planning_ai.chains.map_chain import map_chain
-from planning_ai.common.utils import Consts, length_function
 from planning_ai.states import OverallState, SummaryState
 
 
@@ -48,13 +45,3 @@ def collect_summaries(state: OverallState):
             for idx, summary in enumerate(state["summaries"], start=1)
         ]
     }
-
-
-def should_collapse(
-    state: OverallState,
-) -> Literal["collapse_summaries", "generate_final_summary"]:
-    num_tokens = length_function(state["collapsed_summaries"])
-    if num_tokens > Consts.TOKEN_MAX:
-        return "collapse_summaries"
-    else:
-        return "generate_final_summary"

planning_ai/nodes/reduce_node.py

@@ -1,21 +1,7 @@
-from langchain.chains.combine_documents.reduce import collapse_docs, split_list_of_docs
-
 from planning_ai.chains.reduce_chain import reduce_chain
-from planning_ai.common.utils import Consts, length_function
 from planning_ai.states import OverallState
 
 
-def collapse_summaries(state: OverallState):
-    doc_lists = split_list_of_docs(
-        state["collapsed_summaries"], length_function, Consts.TOKEN_MAX
-    )
-    results = []
-    for doc_list in doc_lists:
-        results.append(collapse_docs(doc_list, reduce_chain.invoke))
-
-    return {"collapsed_summaries": results}
-
-
 def generate_final_summary(state: OverallState):
     response = reduce_chain.invoke({"context": state["collapsed_summaries"]})
     return {