feat: use o3mini for final document

planning_ai/chains/hallucination_chain.py

@@ -2,7 +2,7 @@ from langchain_core.prompts import ChatPromptTemplate
 from pydantic import BaseModel, Field
 
 from planning_ai.common.utils import Paths
-from planning_ai.llms.llm import LLM
+from planning_ai.llms.llm import GPT4o
 
 with open(Paths.PROMPTS / "hallucination.txt", "r") as f:
     reduce_template = f.read()
@@ -15,7 +15,7 @@ class HallucinationChecker(BaseModel):
     explanation: str = Field(..., description="Explain your reasoning for the score")
 
 
-SLLM = LLM.with_structured_output(HallucinationChecker, strict=True)
+SLLM = GPT4o.with_structured_output(HallucinationChecker, strict=True)
 
 hallucination_prompt = ChatPromptTemplate([("system", reduce_template)])
 hallucination_chain = hallucination_prompt | SLLM

planning_ai/chains/map_chain.py

@@ -5,7 +5,7 @@ from langchain_core.prompts import ChatPromptTemplate
 from pydantic import BaseModel, Field, create_model
 
 from planning_ai.common.utils import Paths
-from planning_ai.llms.llm import LLM
+from planning_ai.llms.llm import GPT4o
 from planning_ai.themes import THEMES_AND_POLICIES
 
 with open(Paths.PROMPTS / "map.txt", "r") as f:
@@ -61,7 +61,7 @@ def create_dynamic_map_chain(themes, prompt: str):
     PolicyEnum = create_policy_enum(policy_groups)
     DynamicBriefSummary = create_brief_summary_model(PolicyEnum)
 
-    SLLM = LLM.with_structured_output(DynamicBriefSummary, strict=True)
+    SLLM = GPT4o.with_structured_output(DynamicBriefSummary, strict=True)
 
     prompt = (
         f"{prompt}\n\nAvailable Policies:\n\n- "

planning_ai/chains/policy_chain.py

@@ -2,7 +2,7 @@ from langchain_core.prompts import ChatPromptTemplate
 from pydantic import BaseModel
 
 from planning_ai.common.utils import Paths
-from planning_ai.llms.llm import LLM
+from planning_ai.llms.llm import GPT4o
 
 with open(Paths.PROMPTS / "policy.txt", "r") as f:
     policy_template = f.read()
@@ -19,7 +19,7 @@ class PolicyList(BaseModel):
     policies: list[Policy]
 
 
-SLLM = LLM.with_structured_output(PolicyList, strict=True)
+SLLM = GPT4o.with_structured_output(PolicyList, strict=True)
 
 
 policy_prompt = ChatPromptTemplate([("system", policy_template)])

planning_ai/chains/reduce_chain.py

@@ -2,20 +2,23 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
 
 from planning_ai.common.utils import Paths
-from planning_ai.llms.llm import LLM
+from planning_ai.llms.llm import O3Mini
 
 with open(Paths.PROMPTS / "reduce.txt", "r") as f:
     reduce_template = f.read()
 
+with open(Paths.PROMPTS / "reduce_final.txt", "r") as f:
+    reduce_template_final = f.read()
+
 
 reduce_prompt = ChatPromptTemplate([("system", reduce_template)])
-reduce_chain = reduce_prompt | LLM | StrOutputParser()
+reduce_prompt_final = ChatPromptTemplate([("system", reduce_template_final)])
+reduce_chain = reduce_prompt | O3Mini | StrOutputParser()
+reduce_chain_final = reduce_prompt_final | O3Mini | StrOutputParser()
 
 
 if __name__ == "__main__":
     test_summary = """
-        Summary:
-
         The author expresses concern over the proposed mass development north-west of Cambridge,
         highlighting the significant growth in the area over the past twenty years,
         particularly with the establishment of Cambourne and the expansion of Papworth Everard.

planning_ai/chains/themes_chain.py

@@ -5,7 +5,7 @@ from langchain_core.prompts import ChatPromptTemplate
 from pydantic import BaseModel, Field
 
 from planning_ai.common.utils import Paths
-from planning_ai.llms.llm import LLM
+from planning_ai.llms.llm import GPT4o
 
 
 class Theme(Enum):
@@ -32,7 +32,7 @@ with open(Paths.PROMPTS / "themes.txt", "r") as f:
 
 themes_prompt = ChatPromptTemplate.from_messages([("system", themes_template)])
 
-SLLM = LLM.with_structured_output(ThemeSelector, strict=True)
+SLLM = GPT4o.with_structured_output(ThemeSelector, strict=True)
 
 themes_chain = themes_prompt | SLLM
 

planning_ai/llms/llm.py

@@ -3,4 +3,5 @@ from langchain_openai import ChatOpenAI
 
 load_dotenv()
 
-LLM = ChatOpenAI(temperature=0, model="gpt-4o-mini")
+GPT4o = ChatOpenAI(temperature=0, model="gpt-4o-mini")
+O3Mini = ChatOpenAI(model="o3-mini")