refactor

planning_ai/graph.py

@@ -1,4 +1,4 @@
-from langgraph.constants import START, Send
+from langgraph.constants import START
 from langgraph.graph import END, StateGraph
 
 from planning_ai.nodes.hallucination_node import (
@@ -7,13 +7,9 @@ from planning_ai.nodes.hallucination_node import (
     map_fix_hallucinations,
     map_hallucinations,
 )
-from planning_ai.nodes.map_node import (
-    collect_summaries,
-    generate_summary,
-    map_summaries,
-)
+from planning_ai.nodes.map_node import generate_summary, map_summaries
 from planning_ai.nodes.reduce_node import generate_final_summary
-from planning_ai.states import DocumentState, OverallState
+from planning_ai.states import OverallState
 
 
 def create_graph():

planning_ai/main.py

@@ -180,7 +180,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()[:5] if doc.page_content]
     text_splitter = CharacterTextSplitter.from_tiktoken_encoder(
         chunk_size=1000, chunk_overlap=0
     )
@@ -193,8 +193,7 @@ def main():
         {
             "documents": [doc.page_content for doc in split_docs],
             "filenames": [Path(doc.metadata["source"]) for doc in split_docs],
-        },
-        # {"recursion_limit": 10},
+        }
     ):
         print(list(step.keys()))
 
@@ -207,35 +206,36 @@ def main():
 if __name__ == "__main__":
     doc_title = "Cambridge Response Summary"
     out = main()
-    build_quarto_doc(doc_title, out)
-
-    d = [
-        i
-        for i in out["generate_final_summary"]["summaries_fixed"]
-        if i["iteration"] == 4
-    ][0]
-    d["document"]
-
-    h = [
-        i["summary"].summary
-        for i in out["generate_final_summary"]["hallucinations"]
-        if i["document"] == d["document"]
-    ]
-
-    e = [
-        i["hallucination"].explanation
-        for i in out["generate_final_summary"]["hallucinations"]
-        if i["document"] == d["document"]
-    ]
-
-    test = {
-        "document": d["document"],
-        "final_summary": d["summary"].summary,
-        "attempts": h,
-        "reasoning": e,
-    }
-
-    print(f"Document:\n\n{test['document']}\n\n")
-    print(f"Final:\n\n{test['final_summary']}\n\n")
-    print("Attempts: \n\n*", "\n\n* ".join(test["attempts"]), "\n\n")
-    print("Reasoning: \n\n*", "\n\n* ".join(test["reasoning"]), "\n\n")
+    out["generate_final_summary"]["summaries"]
+    # build_quarto_doc(doc_title, out)
+    #
+    # d = [
+    #     i
+    #     for i in out["generate_final_summary"]["summaries_fixed"]
+    #     if i["iteration"] == 4
+    # ][0]
+    # d["document"]
+    #
+    # h = [
+    #     i["summary"].summary
+    #     for i in out["generate_final_summary"]["hallucinations"]
+    #     if i["document"] == d["document"]
+    # ]
+    #
+    # e = [
+    #     i["hallucination"].explanation
+    #     for i in out["generate_final_summary"]["hallucinations"]
+    #     if i["document"] == d["document"]
+    # ]
+    #
+    # test = {
+    #     "document": d["document"],
+    #     "final_summary": d["summary"].summary,
+    #     "attempts": h,
+    #     "reasoning": e,
+    # }
+    #
+    # print(f"Document:\n\n{test['document']}\n\n")
+    # print(f"Final:\n\n{test['final_summary']}\n\n")
+    # print("Attempts: \n\n*", "\n\n* ".join(test["attempts"]), "\n\n")
+    # print("Reasoning: \n\n*", "\n\n* ".join(test["reasoning"]), "\n\n")

planning_ai/nodes/hallucination_node.py

@@ -1,21 +1,21 @@
-from typing import Literal
-
 from langgraph.constants import Send
 
 from planning_ai.chains.fix_chain import fix_chain
-from planning_ai.chains.hallucination_chain import hallucination_chain
+from planning_ai.chains.hallucination_chain import (
+    HallucinationChecker,
+    hallucination_chain,
+)
 from planning_ai.states import DocumentState, OverallState
 
 
 def check_hallucination(state: DocumentState):
-    print(state["iteration"])
     if state["iteration"] > 5:
         state["iteration"] = -99
         return {"summaries_fixed": [state]}
 
-    response = hallucination_chain.invoke(
+    response: HallucinationChecker = hallucination_chain.invoke(
         {"document": state["document"], "summary": state["summary"]}
-    )
+    )  # type: ignore
     if response.score == 1:
         return {"summaries_fixed": [state]}
 
@@ -43,11 +43,12 @@ def fix_hallucination(state: DocumentState):
             "explanation": state["hallucination"],
         }
     )
-    state["summary"] = response
+    state["summary"] = response  # type: ignore
     return {
         "summaries": [
             {
                 "document": state["document"],
+                "filename": state["filename"],
                 "summary": state["summary"],
                 "iteration": state["iteration"],
             }

planning_ai/nodes/map_node.py

@@ -1,8 +1,3 @@
-import logging
-
-logging.basicConfig(level=logging.WARNING)
-
-from langchain_core.documents import Document
 from langgraph.constants import Send
 
 from planning_ai.chains.map_chain import map_chain
@@ -31,29 +26,3 @@ def map_summaries(state: OverallState):
         )
         for document, filename in zip(state["documents"], state["filenames"])
     ]
-
-
-def collect_summaries(state: OverallState):
-    print("test")
-    __import__("ipdb").set_trace()
-    state.keys()
-    len(state["documents"])
-    len(state["summaries_fixed"])
-    len(state["hallucinations"])
-    state["hallucinations"]
-    return {
-        "summary_documents": [
-            Document(
-                page_content=hallucination.summary,
-                metadata={
-                    "stance": hallucination.stance,
-                    "aims": hallucination.aims,
-                    "places": hallucination.places,
-                    "rating": hallucination.rating,
-                    "hallucination": hallucination.score,
-                    "explanation": hallucination.explanation,
-                },
-            )
-        ]
-        for hallucination in state["summaries_fixed"]
-    }

planning_ai/nodes/reduce_node.py

@@ -8,7 +8,7 @@ def generate_final_summary(state: OverallState):
         return {
             "final_summary": response,
             "summaries_fixed": state["summaries_fixed"],
-            "summaries": state["summary_documents"],
+            "summaries": state["summaries"],
             "hallucinations": state["hallucinations"],
             "documents": state["documents"],
         }

planning_ai/states.py

@@ -10,14 +10,13 @@ from planning_ai.chains.map_chain import BriefSummary
 
 class OverallState(TypedDict):
     documents: list[str]
+
+    final_summary: str
     summaries: Annotated[list, operator.add]
     summaries_fixed: Annotated[list, operator.add]
     hallucinations: Annotated[list, operator.add]
 
     filenames: List[Path]
-    summary_documents: Annotated[list[Document], operator.add]
-    final_summary: str
-
     iterations: list[int]
 
 
@@ -25,6 +24,6 @@ class DocumentState(TypedDict):
     document: str
     summary: BriefSummary
     hallucination: HallucinationChecker
-    filename: Path
 
+    filename: Path
     iteration: int