add hallucination checking cycle

planning_ai/chains/fix_chain.py

@@ -0,0 +1,26 @@
+from langchain_core.prompts import ChatPromptTemplate
+
+from planning_ai.chains.map_chain import SLLM
+
+with open("./planning_ai/chains/prompts/fix_hallucination.txt", "r") as f:
+    map_template = f.read()
+
+map_prompt = ChatPromptTemplate.from_messages([("system", map_template)])
+fix_chain = map_prompt | SLLM
+
+if __name__ == "__main__":
+    test_document = """
+    The Local Plan proposes a mass development north-west of Cambridge despite marked growth
+    in the last twenty years or so following the previous New Settlement Study. In this period,
+    the major settlement of Cambourne has been created - now over the projected 3,000 homes and
+    Papworth Everard has grown beyond recognition. This in itself is a matter of concern.
+    """
+
+    result = fix_chain.invoke(
+        {
+            "summary": "This plan is great because they are building a nuclear power plant.",
+            "explanation": "The original response does not mention a nuclear power plant, and appears to view the plan negatively",
+            "context": test_document,
+        }
+    )
+    print(result)

planning_ai/chains/hallucination_chain.py

@@ -0,0 +1,35 @@
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.pydantic_v1 import BaseModel, Field
+
+from planning_ai.llms.llm import LLM
+
+with open("./planning_ai/chains/prompts/hallucination.txt", "r") as f:
+    reduce_template = f.read()
+
+
+class HallucinationChecker(BaseModel):
+    """Grade the summary based upon the above criteria."""
+
+    score: int = Field(..., description="Score for the summary")
+    explanation: str = Field(..., description="Explain your reasoning for the score")
+
+
+SLLM = LLM.with_structured_output(HallucinationChecker)
+
+hallucination_prompt = ChatPromptTemplate([("human", reduce_template)])
+hallucination_chain = hallucination_prompt | SLLM
+
+if __name__ == "__main__":
+    test_document = """
+    The Local Plan proposes a mass development north-west of Cambridge despite marked growth
+    in the last twenty years or so following the previous New Settlement Study. In this period,
+    the major settlement of Cambourne has been created - now over the projected 3,000 homes and
+    Papworth Everard has grown beyond recognition. This in itself is a matter of concern.
+    """
+
+    result = hallucination_chain.invoke(
+        {
+            "summary": "The author fully supports the plan due to the nuclear power plant.",
+            "document": test_document,
+        }
+    )

planning_ai/chains/prompts/fix_hallucination.txt

@@ -0,0 +1,12 @@
+You are tasked with summarising a response to a planning application proposed by South Cambridgeshire Council. Below, we have provided an **incorrect summary** of the response, along with an **explanation** detailing why the summary is incorrect. Your job is to generate a **correct** summary based solely on the original response provided, avoiding the errors highlighted in the explanation.
+
+- **Incorrect Summary**:
+  {summary}
+
+- **Explanation of Errors**:
+  {explanation}
+
+- **Original Response**:
+  {context}
+
+**Your task**: Write a concise and accurate summary of the original response, taking into account the errors highlighted in the explanation.

planning_ai/chains/prompts/hallucination.txt

@@ -0,0 +1,15 @@
+You are grading text summaries of source documents focused on faithfulness and detection of any hallucinations.
+
+Ensure that the Assistant's Summary meets the following criteria: 
+(1) it does not contain information outside the score of the source document provided
+(2) the summary should be fully grounded in and based upon the source documents
+
+Score:
+A score of 1 means that the Assistant Summary meets the criteria. This is the highest (best) score. 
+A score of 0 means that the Assistant Summary does not the criteria. This is the lowest possible score you can give.
+
+Explain your reasoning step-by-step to ensure your reasoning and conclusion are correct. 
+
+Assistant's Summary: {summary}
+
+Source document: {document}

planning_ai/graph.py

@@ -19,10 +19,21 @@ from planning_ai.states import DocumentState, OverallState
 def create_graph():
     graph = StateGraph(OverallState)
     graph.add_node("generate_summary", generate_summary)
+    graph.add_node("check_hallucination", check_hallucination)
+    graph.add_node("fix_hallucination", fix_hallucination)
     graph.add_node("generate_final_summary", generate_final_summary)
 
     graph.add_conditional_edges(START, map_summaries, ["generate_summary"])
-    graph.add_edge("generate_summary", "generate_final_summary")
+    graph.add_conditional_edges(
+        "generate_summary", map_hallucinations, ["check_hallucination"]
+    )
+    graph.add_conditional_edges(
+        "check_hallucination", map_fix_hallucinations, ["fix_hallucination"]
+    )
+    graph.add_conditional_edges(
+        "fix_hallucination", map_hallucinations, ["check_hallucination"]
+    )
+    graph.add_edge("check_hallucination", "generate_final_summary")
     graph.add_edge("generate_final_summary", END)
 
     return graph.compile()

planning_ai/nodes/hallucination_node.py

@@ -0,0 +1,68 @@
+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.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(
+        {"document": state["document"], "summary": state["summary"]}
+    )
+    if response.score == 1:
+        return {"summaries_fixed": [state]}
+
+    return {
+        "hallucinations": [
+            {
+                "hallucination": response,
+                "document": state["document"],
+                "summary": state["summary"],
+                "iteration": state["iteration"] + 1,
+            }
+        ]
+    }
+
+
+def map_hallucinations(state: OverallState):
+    return [Send("check_hallucination", summary) for summary in state["summaries"]]
+
+
+def fix_hallucination(state: DocumentState):
+    response = fix_chain.invoke(
+        {
+            "context": state["document"],
+            "summary": state["summary"],
+            "explanation": state["hallucination"],
+        }
+    )
+    state["summary"] = response
+    return {
+        "summaries": [
+            {
+                "document": state["document"],
+                "summary": state["summary"],
+                "iteration": state["iteration"],
+            }
+        ]
+    }
+
+
+def map_fix_hallucinations(state: OverallState):
+    hallucinations = []
+    if "hallucinations" in state:
+        hallucinations = [
+            hallucination
+            for hallucination in state["hallucinations"]
+            if hallucination["hallucination"].score != 1
+        ]
+    return [
+        Send("fix_hallucination", hallucination) for hallucination in hallucinations
+    ]

planning_ai/nodes/reduce_node.py

@@ -3,11 +3,12 @@ from planning_ai.states import OverallState
 
 
 def generate_final_summary(state: OverallState):
-    __import__("ipdb").set_trace()
-    # response = reduce_chain.invoke({"context": state["summary_documents"]})
-    return {
-        # "final_summary": response,
-        "summaries": state["summary_documents"],
-        "hallucinations": state["hallucinations"],
-        "summary": state["summary_documents"],
-    }
+    if len(state["documents"]) == len(state["summaries_fixed"]):
+        response = reduce_chain.invoke({"context": state["summaries_fixed"]})
+        return {
+            "final_summary": response,
+            "summaries_fixed": state["summaries_fixed"],
+            "summaries": state["summary_documents"],
+            "hallucinations": state["hallucinations"],
+            "documents": state["documents"],
+        }

planning_ai/states.py

@@ -1,18 +1,30 @@
 import operator
 from pathlib import Path
-from typing import Annotated, List, TypedDict
+from typing import Annotated, List, Optional, TypedDict
 
 from langchain_core.documents import Document
 
+from planning_ai.chains.hallucination_chain import HallucinationChecker
+from planning_ai.chains.map_chain import BriefSummary
+
 
 class OverallState(TypedDict):
-    contents: List[str]
-    filenames: List[str]
+    documents: list[str]
     summaries: Annotated[list, operator.add]
-    collapsed_summaries: List[Document]
+    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]
 
-class SummaryState(TypedDict):
-    content: str
+
+class DocumentState(TypedDict):
+    document: str
+    summary: BriefSummary
+    hallucination: HallucinationChecker
     filename: Path
+
+    iteration: int