Add test for drawing lance example

libs/langgraph/poetry.lock

@@ -1760,13 +1760,13 @@ langchain-core = ">=0.2.2rc1,<0.3"
 
 [[package]]
 name = "langchain-core"
-version = "0.2.22"
+version = "0.2.25"
 description = "Building applications with LLMs through composability"
 optional = false
 python-versions = "<4.0,>=3.8.1"
 files = [
-    {file = "langchain_core-0.2.22-py3-none-any.whl", hash = "sha256:7731a86440c0958b3186c003fb9b26b2d5a682a6344bda7bfb9174e2898f8b43"},
-    {file = "langchain_core-0.2.22.tar.gz", hash = "sha256:582d6f929a43b830139444e4124123cd415331ad62f25757b1406252958cdcac"},
+    {file = "langchain_core-0.2.25-py3-none-any.whl", hash = "sha256:03d61b2a7f4b5f98df248c1b1f0ccd95c9d5ef2269e174133724365cd2a7ee1e"},
+    {file = "langchain_core-0.2.25.tar.gz", hash = "sha256:e64106a7d0e37e4d35b767f79e6c62b56e825f08f9e8cc4368bcea9955257a7e"},
 ]
 
 [package.dependencies]

libs/langgraph/tests/test_pregel.py

@@ -9097,3 +9097,130 @@ def test_remove_message_from_node():
     output = app.invoke([HumanMessage(content="Hi")])
     assert len(output) == 2
     assert output[-1].content == "How can I help you?"
+
+
+def test_xray_lance(snapshot: SnapshotAssertion):
+    from langchain_core.messages import AnyMessage, HumanMessage
+    from langchain_core.pydantic_v1 import BaseModel, Field
+
+    class Analyst(BaseModel):
+        affiliation: str = Field(
+            description="Primary affiliation of the investment analyst.",
+        )
+        name: str = Field(
+            description="Name of the investment analyst.",
+            pattern=r"^[a-zA-Z0-9_-]{1,64}$",
+        )
+        role: str = Field(
+            description="Role of the investment analyst in the context of the topic.",
+        )
+        description: str = Field(
+            description="Description of the investment analyst focus, concerns, and motives.",
+        )
+
+        @property
+        def persona(self) -> str:
+            return f"Name: {self.name}\nRole: {self.role}\nAffiliation: {self.affiliation}\nDescription: {self.description}\n"
+
+    class Perspectives(BaseModel):
+        analysts: List[Analyst] = Field(
+            description="Comprehensive list of investment analysts with their roles and affiliations.",
+        )
+
+    class Section(BaseModel):
+        section_title: str = Field(..., title="Title of the section")
+        context: str = Field(
+            ..., title="Provide a clear summary of the focus area that you researched."
+        )
+        findings: str = Field(
+            ...,
+            title="Give a clear and detailed overview of your findings based upon the expert interview.",
+        )
+        thesis: str = Field(
+            ...,
+            title="Give a clear and specific investment thesis based upon these findings.",
+        )
+
+    class InterviewState(TypedDict):
+        messages: Annotated[List[AnyMessage], add_messages]
+        analyst: Analyst
+        section: Section
+
+    class ResearchGraphState(TypedDict):
+        analysts: List[Analyst]
+        topic: str
+        max_analysts: int
+        sections: List[Section]
+        interviews: Annotated[list, operator.add]
+
+    # Conditional edge
+    def route_messages(state):
+        return "ask_question"
+
+    def generate_question(state):
+        return ...
+
+    def generate_answer(state):
+        return ...
+
+    # Add nodes and edges
+    interview_builder = StateGraph(InterviewState)
+    interview_builder.add_node("ask_question", generate_question)
+    interview_builder.add_node("answer_question", generate_answer)
+
+    # Flow
+    interview_builder.add_edge(START, "ask_question")
+    interview_builder.add_edge("ask_question", "answer_question")
+    interview_builder.add_conditional_edges("answer_question", route_messages)
+
+    # Set up memory
+    memory = MemorySaver()
+
+    # Interview
+    interview_graph = interview_builder.compile(checkpointer=memory).with_config(
+        run_name="Conduct Interviews"
+    )
+
+    # View
+    assert interview_graph.get_graph().to_json() == snapshot
+
+    def run_all_interviews(state: ResearchGraphState):
+        """Edge to run the interview sub-graph using Send"""
+        return [
+            Send(
+                "conduct_interview",
+                {
+                    "analyst": Analyst(),
+                    "messages": [
+                        HumanMessage(
+                            content="So you said you were writing an article on ...?"
+                        )
+                    ],
+                },
+            )
+            for s in state["analysts"]
+        ]
+
+    def generate_sections(state: ResearchGraphState):
+        return ...
+
+    def generate_analysts(state: ResearchGraphState):
+        return ...
+
+    builder = StateGraph(ResearchGraphState)
+    builder.add_node("generate_analysts", generate_analysts)
+    builder.add_node("conduct_interview", interview_builder.compile())
+    builder.add_node("generate_sections", generate_sections)
+
+    builder.add_edge(START, "generate_analysts")
+    builder.add_conditional_edges(
+        "generate_analysts", run_all_interviews, ["conduct_interview"]
+    )
+    builder.add_edge("conduct_interview", "generate_sections")
+    builder.add_edge("generate_sections", END)
+
+    graph = builder.compile()
+
+    # View
+    assert graph.get_graph().to_json() == snapshot
+    assert graph.get_graph(xray=1).to_json() == snapshot