lib: Add more tests

- add more unit tests (courtesy of claude)
- move tests with large assertions to separate file

libs/langgraph/tests/__snapshots__/test_pregel.ambr

@@ -1,230 +1,4 @@
 # serializer version: 1
-# name: test_branch_then.1
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([__start__]):::first
-  	prepare(prepare)
-  	tool_two_slow(tool_two_slow)
-  	tool_two_fast(tool_two_fast)
-  	finish(finish)
-  	__end__([__end__]):::last
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_branch_then[duckdb]
-  '''
-  graph TD;
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  
-  '''
-# ---
-# name: test_branch_then[duckdb].1
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	prepare(prepare)
-  	tool_two_slow(tool_two_slow)
-  	tool_two_fast(tool_two_fast)
-  	finish(finish)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_branch_then[memory]
-  '''
-  graph TD;
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  
-  '''
-# ---
-# name: test_branch_then[memory].1
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	prepare(prepare)
-  	tool_two_slow(tool_two_slow)
-  	tool_two_fast(tool_two_fast)
-  	finish(finish)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_branch_then[postgres]
-  '''
-  graph TD;
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  
-  '''
-# ---
-# name: test_branch_then[postgres].1
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	prepare(prepare)
-  	tool_two_slow(tool_two_slow)
-  	tool_two_fast(tool_two_fast)
-  	finish(finish)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_branch_then[postgres_pipe]
-  '''
-  graph TD;
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  
-  '''
-# ---
-# name: test_branch_then[postgres_pipe].1
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	prepare(prepare)
-  	tool_two_slow(tool_two_slow)
-  	tool_two_fast(tool_two_fast)
-  	finish(finish)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_branch_then[postgres_pool]
-  '''
-  graph TD;
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  
-  '''
-# ---
-# name: test_branch_then[postgres_pool].1
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	prepare(prepare)
-  	tool_two_slow(tool_two_slow)
-  	tool_two_fast(tool_two_fast)
-  	finish(finish)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
-# name: test_branch_then[sqlite]
-  '''
-  graph TD;
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  
-  '''
-# ---
-# name: test_branch_then[sqlite].1
-  '''
-  %%{init: {'flowchart': {'curve': 'linear'}}}%%
-  graph TD;
-  	__start__([<p>__start__</p>]):::first
-  	prepare(prepare)
-  	tool_two_slow(tool_two_slow)
-  	tool_two_fast(tool_two_fast)
-  	finish(finish)
-  	__end__([<p>__end__</p>]):::last
-  	__start__ --> prepare;
-  	finish --> __end__;
-  	prepare -.-> tool_two_slow;
-  	tool_two_slow --> finish;
-  	prepare -.-> tool_two_fast;
-  	tool_two_fast --> finish;
-  	classDef default fill:#f2f0ff,line-height:1.2
-  	classDef first fill-opacity:0
-  	classDef last fill:#bfb6fc
-  
-  '''
-# ---
 # name: test_conditional_entrypoint_graph
   '{"title": "LangGraphInput"}'
 # ---

libs/langgraph/tests/agents.py

@@ -0,0 +1,43 @@
+from typing import Literal, Union
+
+from pydantic import BaseModel
+
+
+# define these objects to avoid importing langchain_core.agents
+# and therefore avoid relying on core Pydantic version
+class AgentAction(BaseModel):
+    tool: str
+    tool_input: Union[str, dict]
+    log: str
+    type: Literal["AgentAction"] = "AgentAction"
+
+    model_config = {
+        "json_schema_extra": {
+            "description": (
+                """Represents a request to execute an action by an agent.
+
+The action consists of the name of the tool to execute and the input to pass
+to the tool. The log is used to pass along extra information about the action."""
+            )
+        }
+    }
+
+
+class AgentFinish(BaseModel):
+    """Final return value of an ActionAgent.
+
+    Agents return an AgentFinish when they have reached a stopping condition.
+    """
+
+    return_values: dict
+    log: str
+    type: Literal["AgentFinish"] = "AgentFinish"
+    model_config = {
+        "json_schema_extra": {
+            "description": (
+                """Final return value of an ActionAgent.
+
+Agents return an AgentFinish when they have reached a stopping condition."""
+            )
+        }
+    }