From 23f84c9ac9520507c7eed1994fbcc4b2b24efb44 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 15 Jan 2024 13:18:32 -0800
Subject: [PATCH] Fix tests to run in py 3.9

---
 Dockerfile                 |  4 ++--
 tests/test_pregel.py       |  8 ++++----
 tests/test_pregel_async.py | 17 +++++++++++++----
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 2095ad194..d29a3288c 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM python:3.9-slim
+FROM python:3.9
 
 # Set the working directory to /app
 WORKDIR /app
@@ -7,6 +7,6 @@ WORKDIR /app
 COPY . .
 
 # Install any needed packages specified in requirements.txt
-RUN pip install poetry && poetry config virtualenvs.create false && poetry install --with test,lint,typing
+RUN pip install poetry && poetry config virtualenvs.create false && poetry install --with test,lint,typing,dev
 
 RUN poetry run pytest
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 7553e4bc6..65494bd16 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -2,7 +2,7 @@ import operator
 import time
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from typing import Annotated, Generator, TypedDict
+from typing import Annotated, Generator, Optional, TypedDict, Union
 
 import pytest
 from langchain_core.runnables import RunnablePassthrough
@@ -584,7 +584,7 @@ def test_conditional_graph() -> None:
         ]
     )
 
-    def agent_parser(input: str) -> AgentFinish | AgentAction:
+    def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
         if input.startswith("finish"):
             _, answer = input.split(":")
             return AgentFinish(return_values={"answer": answer}, log=input)
@@ -786,7 +786,7 @@ def test_conditional_graph_state() -> None:
 
     class AgentState(TypedDict):
         input: str
-        agent_outcome: AgentAction | AgentFinish | None
+        agent_outcome: Optional[Union[AgentAction, AgentFinish]]
         intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
 
     # Assemble the tools
@@ -808,7 +808,7 @@ def test_conditional_graph_state() -> None:
         ]
     )
 
-    def agent_parser(input: str) -> AgentFinish | AgentAction:
+    def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
         if input.startswith("finish"):
             _, answer = input.split(":")
             return {
diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py
index cbd5406ad..e776c512b 100644
--- a/tests/test_pregel_async.py
+++ b/tests/test_pregel_async.py
@@ -1,7 +1,16 @@
 import asyncio
 import operator
 from contextlib import asynccontextmanager, contextmanager
-from typing import Annotated, Any, AsyncGenerator, AsyncIterator, Generator, TypedDict
+from typing import (
+    Annotated,
+    Any,
+    AsyncGenerator,
+    AsyncIterator,
+    Generator,
+    Optional,
+    TypedDict,
+    Union,
+)
 
 import pytest
 from langchain_core.runnables import RunnablePassthrough
@@ -621,7 +630,7 @@ async def test_conditional_graph() -> None:
         ]
     )
 
-    async def agent_parser(input: str) -> AgentFinish | AgentAction:
+    async def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
         if input.startswith("finish"):
             _, answer = input.split(":")
             return AgentFinish(return_values={"answer": answer}, log=input)
@@ -831,7 +840,7 @@ async def test_conditional_graph_state() -> None:
 
     class AgentState(TypedDict):
         input: str
-        agent_outcome: AgentAction | AgentFinish | None
+        agent_outcome: Optional[Union[AgentAction, AgentFinish]]
         intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
 
     # Assemble the tools
@@ -853,7 +862,7 @@ async def test_conditional_graph_state() -> None:
         ]
     )
 
-    def agent_parser(input: str) -> AgentFinish | AgentAction:
+    def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
         if input.startswith("finish"):
             _, answer = input.split(":")
             return {