From b67c7d64427d1c02816b71782cab54a24e6878ab Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Mon, 27 Mar 2023 17:32:50 -0300 Subject: [PATCH] fix: updates to json and test added --- tests/data/complex_example.json | 964 ++++++++++++++++---------------- tests/test_graph.py | 59 +- 2 files changed, 507 insertions(+), 516 deletions(-) diff --git a/tests/data/complex_example.json b/tests/data/complex_example.json index 9d478c687..8d46a97d6 100644 --- a/tests/data/complex_example.json +++ b/tests/data/complex_example.json @@ -1,16 +1,440 @@ { - "name": "New Flow ", + "name": "New Flow", "id": "0", "data": { "nodes": [ { "width": 384, - "height": 477, - "id": "dndnode_7", + "height": 351, + "id": "dndnode_3", "type": "genericNode", "position": { - "x": -211.61829328351757, - "y": 132.6841414309356 + "x": 612.9299322834961, + "y": 194.75070242078417 + }, + "data": { + "type": "ZeroShotAgent", + "node": { + "template": { + "_type": "zero-shot-react-description", + "llm_chain": { + "type": "LLMChain", + "required": true, + "placeholder": "", + "list": false, + "show": true, + "password": false, + "multiline": false + }, + "allowed_tools": { + "type": "Tool", + "required": false, + "placeholder": "", + "list": true, + "show": true, + "password": false, + "multiline": false, + "value": null + }, + "return_values": { + "type": "str", + "required": false, + "placeholder": "", + "list": true, + "show": false, + "password": false, + "multiline": false, + "value": [ + "output" + ] + } + }, + "description": "Agent for the MRKL chain.", + "base_classes": [ + "Agent", + "function" + ] + }, + "id": "dndnode_3", + "value": null + }, + "selected": false, + "positionAbsolute": { + "x": 612.9299322834961, + "y": 194.75070242078417 + }, + "dragging": false + }, + { + "width": 384, + "height": 463, + "id": "dndnode_27", + "type": "genericNode", + "position": { + "x": 86.29922452047686, + "y": 39.132143332238115 + }, + "data": { + "type": "Tool", + "node": { + "template": { + "name": { + "type": "str", + "required": true, + "list": false, + "show": true, + "placeholder": "", + "value": "Uppercase", + "password": false, + "multiline": false + }, + "description": { + "type": "str", + "required": true, + "list": false, + "show": true, + "placeholder": "", + "value": "Returns an uppercase version of the text passed.", + "password": false, + "multiline": false + }, + "func": { + "type": "function", + "required": true, + "list": false, + "show": true, + "value": "", + "multiline": false, + "password": false + }, + "_type": "Tool" + }, + "name": "Tool", + "func": "", + "description": "", + "base_classes": [ + "Tool" + ] + }, + "id": "dndnode_27", + "value": null + }, + "selected": false, + "positionAbsolute": { + "x": 86.29922452047686, + "y": 39.132143332238115 + }, + "dragging": false + }, + { + "width": 384, + "height": 463, + "id": "dndnode_28", + "type": "genericNode", + "position": { + "x": 1134.4549802672202, + "y": 287.9885910233929 + }, + "data": { + "type": "Tool", + "node": { + "template": { + "name": { + "type": "str", + "required": true, + "list": false, + "show": true, + "placeholder": "", + "value": "", + "password": false, + "multiline": false + }, + "description": { + "type": "str", + "required": true, + "list": false, + "show": true, + "placeholder": "", + "value": "", + "password": false, + "multiline": false + }, + "func": { + "type": "function", + "required": true, + "list": false, + "show": true, + "value": "", + "multiline": false, + "password": false + }, + "_type": "Tool" + }, + "name": "Tool", + "func": "", + "description": "", + "base_classes": [ + "Tool" + ] + }, + "id": "dndnode_28", + "value": null + }, + "positionAbsolute": { + "x": 1134.4549802672202, + "y": 287.9885910233929 + }, + "selected": false, + "dragging": false + }, + { + "width": 384, + "height": 357, + "id": "dndnode_40", + "type": "genericNode", + "position": { + "x": -366.4341715850213, + "y": 136.29836646158452 + }, + "data": { + "type": "PythonFunction", + "node": { + "template": { + "code": { + "required": true, + "placeholder": "", + "show": true, + "multiline": true, + "value": "\ndef upper_case(text: str) -> str:\n return text.upper()\n", + "name": "code", + "type": "str", + "list": false + }, + "_type": "python_function" + }, + "description": "Python function to be executed.", + "base_classes": [ + "function" + ] + }, + "id": "dndnode_40", + "value": null + }, + "selected": false, + "positionAbsolute": { + "x": -366.4341715850213, + "y": 136.29836646158452 + }, + "dragging": false + }, + { + "width": 384, + "height": 351, + "id": "dndnode_41", + "type": "genericNode", + "position": { + "x": 1642.7653281427417, + "y": 69.01105573790835 + }, + "data": { + "type": "ZeroShotAgent", + "node": { + "template": { + "_type": "zero-shot-react-description", + "llm_chain": { + "type": "LLMChain", + "required": true, + "placeholder": "", + "list": false, + "show": true, + "password": false, + "multiline": false + }, + "allowed_tools": { + "type": "Tool", + "required": false, + "placeholder": "", + "list": true, + "show": true, + "password": false, + "multiline": false, + "value": null + }, + "return_values": { + "type": "str", + "required": false, + "placeholder": "", + "list": true, + "show": false, + "password": false, + "multiline": false, + "value": [ + "output" + ] + } + }, + "description": "Agent for the MRKL chain.", + "base_classes": [ + "Agent", + "function" + ] + }, + "id": "dndnode_41", + "value": null + }, + "selected": false, + "positionAbsolute": { + "x": 1642.7653281427417, + "y": 69.01105573790835 + }, + "dragging": false + }, + { + "width": 384, + "height": 529, + "id": "dndnode_42", + "type": "genericNode", + "position": { + "x": -379.23467185725826, + "y": -551.3889442620921 + }, + "data": { + "type": "ZeroShotPrompt", + "node": { + "template": { + "prefix": { + "required": false, + "placeholder": "", + "show": true, + "multiline": true, + "value": "Answer the following questions as best you can. You have access to the following tools:", + "name": "prefix", + "type": "str", + "list": false + }, + "suffix": { + "required": true, + "placeholder": "", + "show": true, + "multiline": true, + "value": "Begin!\n\nQuestion: {input}\nThought:{agent_scratchpad}", + "name": "suffix", + "type": "str", + "list": false + }, + "format_instructions": { + "required": false, + "placeholder": "", + "show": true, + "multiline": true, + "value": "Use the following format:\n\nQuestion: the input question you must answer\nThought: you should always think about what to do\nAction: the action to take, should be one of [{tool_names}]\nAction Input: the input to the action\nObservation: the result of the action\n... (this Thought/Action/Action Input/Observation can repeat N times)\nThought: I now know the final answer\nFinal Answer: the final answer to the original input question", + "name": "format_instructions", + "type": "str", + "list": false + }, + "_type": "zero_shot" + }, + "description": "Prompt template for Zero Shot Agent.", + "base_classes": [ + "BasePromptTemplate" + ] + }, + "id": "dndnode_42", + "value": null + }, + "selected": false, + "positionAbsolute": { + "x": -379.23467185725826, + "y": -551.3889442620921 + }, + "dragging": false + }, + { + "width": 384, + "height": 391, + "id": "dndnode_43", + "type": "genericNode", + "position": { + "x": 100.76532814274174, + "y": -437.78894426209195 + }, + "data": { + "type": "LLMChain", + "node": { + "template": { + "_type": "llm_chain", + "memory": { + "type": "BaseMemory", + "required": false, + "placeholder": "", + "list": false, + "show": true, + "password": false, + "multiline": false, + "value": null + }, + "verbose": { + "type": "bool", + "required": false, + "placeholder": "", + "list": false, + "show": false, + "password": false, + "multiline": false, + "value": false + }, + "prompt": { + "type": "BasePromptTemplate", + "required": true, + "placeholder": "", + "list": false, + "show": true, + "password": false, + "multiline": false + }, + "llm": { + "type": "BaseLanguageModel", + "required": true, + "placeholder": "", + "list": false, + "show": true, + "password": false, + "multiline": false + }, + "output_key": { + "type": "str", + "required": false, + "placeholder": "", + "list": false, + "show": false, + "password": true, + "multiline": false, + "value": "text" + } + }, + "description": "Chain to run queries against LLMs.", + "base_classes": [ + "Chain" + ] + }, + "id": "dndnode_43", + "value": null + }, + "selected": false, + "positionAbsolute": { + "x": 100.76532814274174, + "y": -437.78894426209195 + }, + "dragging": false + }, + { + "width": 384, + "height": 477, + "id": "dndnode_44", + "type": "genericNode", + "position": { + "x": -841.2346718572583, + "y": 368.6110557379079 }, "data": { "type": "OpenAI", @@ -58,7 +482,10 @@ "value": "text-davinci-003", "options": [ "text-davinci-003", - "text-davinci-002" + "text-davinci-002", + "text-curie-001", + "text-babbage-001", + "text-ada-001" ] }, "temperature": { @@ -209,531 +636,86 @@ "BaseLanguageModel" ] }, - "id": "dndnode_7", - "value": null - }, - "selected": false, - "positionAbsolute": { - "x": -211.61829328351757, - "y": 132.6841414309356 - }, - "dragging": false - }, - { - "width": 384, - "height": 351, - "id": "dndnode_8", - "type": "genericNode", - "position": { - "x": 429.5817067164825, - "y": 552.6841414309356 - }, - "data": { - "type": "PAL-MATH", - "node": { - "template": { - "llm": { - "type": "BaseLLM", - "required": true, - "list": false, - "show": true - }, - "_type": "pal-math" - }, - "name": "PAL-MATH", - "description": "A language model that is really good at solving complex word math problems. Input should be a fully worded hard word math problem.", - "base_classes": [ - "Tool" - ] - }, - "id": "dndnode_8", - "value": null - }, - "selected": false, - "positionAbsolute": { - "x": 429.5817067164825, - "y": 552.6841414309356 - }, - "dragging": false - }, - { - "width": 384, - "height": 529, - "id": "dndnode_9", - "type": "genericNode", - "position": { - "x": -238.41829328351753, - "y": 734.6841414309356 - }, - "data": { - "type": "ZeroShotPrompt", - "node": { - "template": { - "_type": "zero_shot", - "prefix": { - "type": "str", - "required": false, - "placeholder": "", - "list": false, - "show": true, - "multiline": true, - "value": "Answer the following questions as best you can. You have access to the following tools:" - }, - "suffix": { - "type": "str", - "required": true, - "placeholder": "", - "list": false, - "show": true, - "multiline": true, - "value": "Begin!\n\nQuestion: {input}\nThought:{agent_scratchpad}" - }, - "format_instructions": { - "type": "str", - "required": false, - "placeholder": "", - "list": false, - "show": true, - "multiline": true, - "value": "Use the following format:\n\nQuestion: the input question you must answer\nThought: you should always think about what to do\nAction: the action to take, should be one of [{tool_names}]\nAction Input: the input to the action\nObservation: the result of the action\n... (this Thought/Action/Action Input/Observation can repeat N times)\nThought: I now know the final answer\nFinal Answer: the final answer to the original input question" - } - }, - "description": "Prompt template for Zero Shot Agent.", - "base_classes": [ - "BasePromptTemplate" - ] - }, - "id": "dndnode_9", - "value": null - }, - "selected": false, - "positionAbsolute": { - "x": -238.41829328351753, - "y": 734.6841414309356 - }, - "dragging": false - }, - { - "width": 384, - "height": 351, - "id": "dndnode_34", - "type": "genericNode", - "position": { - "x": 931.431035174925, - "y": 285.11234969236074 - }, - "data": { - "type": "ZeroShotAgent", - "node": { - "template": { - "_type": "zero-shot-react-description", - "llm_chain": { - "type": "LLMChain", - "required": true, - "placeholder": "", - "list": false, - "show": true, - "password": false, - "multiline": false - }, - "allowed_tools": { - "type": "Tool", - "required": false, - "placeholder": "", - "list": true, - "show": true, - "password": false, - "multiline": false, - "value": null - }, - "return_values": { - "type": "str", - "required": false, - "placeholder": "", - "list": true, - "show": false, - "password": false, - "multiline": false, - "value": [ - "output" - ] - } - }, - "description": "Agent for the MRKL chain.", - "base_classes": [ - "Agent", - "function" - ] - }, - "id": "dndnode_34", - "value": null - }, - "selected": false, - "positionAbsolute": { - "x": 931.431035174925, - "y": 285.11234969236074 - }, - "dragging": false - }, - { - "width": 384, - "height": 523, - "id": "dndnode_41", - "type": "genericNode", - "position": { - "x": 1456.9923517899285, - "y": 180.9267225160853 - }, - "data": { - "type": "BaseTool", - "node": { - "template": { - "name": { - "type": "str", - "required": true, - "list": false, - "show": true, - "placeholder": "", - "value": "" - }, - "description": { - "type": "str", - "required": true, - "list": false, - "show": true, - "placeholder": "", - "value": "" - }, - "func": { - "type": "function", - "required": true, - "list": false, - "show": true, - "value": "", - "multiline": true - }, - "_type": "BaseTool" - }, - "name": "PAL-MATH", - "description": "A language model that is really good at solving complex word math problems. Input should be a fully worded hard word math problem.", - "base_classes": [ - "Tool" - ] - }, - "id": "dndnode_41", - "value": null - }, - "selected": false, - "positionAbsolute": { - "x": 1456.9923517899285, - "y": 180.9267225160853 - }, - "dragging": false - }, - { - "width": 384, - "height": 351, - "id": "dndnode_42", - "type": "genericNode", - "position": { - "x": 1994.12084226133, - "y": -48.749412190849014 - }, - "data": { - "type": "ZeroShotAgent", - "node": { - "template": { - "_type": "zero-shot-react-description", - "llm_chain": { - "type": "LLMChain", - "required": true, - "placeholder": "", - "list": false, - "show": true, - "password": false, - "multiline": false - }, - "allowed_tools": { - "type": "Tool", - "required": false, - "placeholder": "", - "list": true, - "show": true, - "password": false, - "multiline": false, - "value": null - }, - "return_values": { - "type": "str", - "required": false, - "placeholder": "", - "list": true, - "show": false, - "password": false, - "multiline": false, - "value": [ - "output" - ] - } - }, - "description": "Agent for the MRKL chain.", - "base_classes": [ - "Agent", - "function" - ] - }, - "id": "dndnode_42", - "value": null - }, - "selected": false, - "positionAbsolute": { - "x": 1994.12084226133, - "y": -48.749412190849014 - }, - "dragging": false - }, - { - "width": 384, - "height": 391, - "id": "dndnode_43", - "type": "genericNode", - "position": { - "x": 394.51027997668166, - "y": 45.75513998834094 - }, - "data": { - "type": "LLMChain", - "node": { - "template": { - "_type": "llm_chain", - "memory": { - "type": "BaseMemory", - "required": false, - "placeholder": "", - "list": false, - "show": true, - "password": false, - "multiline": false, - "value": null - }, - "verbose": { - "type": "bool", - "required": false, - "placeholder": "", - "list": false, - "show": false, - "password": false, - "multiline": false, - "value": false - }, - "prompt": { - "type": "BasePromptTemplate", - "required": true, - "placeholder": "", - "list": false, - "show": true, - "password": false, - "multiline": false - }, - "llm": { - "type": "BaseLanguageModel", - "required": true, - "placeholder": "", - "list": false, - "show": true, - "password": false, - "multiline": false - }, - "output_key": { - "type": "str", - "required": false, - "placeholder": "", - "list": false, - "show": false, - "password": true, - "multiline": false, - "value": "text" - } - }, - "description": "Chain to run queries against LLMs.", - "base_classes": [ - "Chain" - ] - }, - "id": "dndnode_43", - "value": null - }, - "selected": false, - "positionAbsolute": { - "x": 394.51027997668166, - "y": 45.75513998834094 - }, - "dragging": false - }, - { - "width": 384, - "height": 391, - "id": "dndnode_44", - "type": "genericNode", - "position": { - "x": 1404.5102799766814, - "y": -328.24486001165906 - }, - "data": { - "type": "LLMChain", - "node": { - "template": { - "_type": "llm_chain", - "memory": { - "type": "BaseMemory", - "required": false, - "placeholder": "", - "list": false, - "show": true, - "password": false, - "multiline": false, - "value": null - }, - "verbose": { - "type": "bool", - "required": false, - "placeholder": "", - "list": false, - "show": false, - "password": false, - "multiline": false, - "value": false - }, - "prompt": { - "type": "BasePromptTemplate", - "required": true, - "placeholder": "", - "list": false, - "show": true, - "password": false, - "multiline": false - }, - "llm": { - "type": "BaseLanguageModel", - "required": true, - "placeholder": "", - "list": false, - "show": true, - "password": false, - "multiline": false - }, - "output_key": { - "type": "str", - "required": false, - "placeholder": "", - "list": false, - "show": false, - "password": true, - "multiline": false, - "value": "text" - } - }, - "description": "Chain to run queries against LLMs.", - "base_classes": [ - "Chain" - ] - }, "id": "dndnode_44", "value": null }, "selected": false, "positionAbsolute": { - "x": 1404.5102799766814, - "y": -328.24486001165906 + "x": -841.2346718572583, + "y": 368.6110557379079 }, "dragging": false } ], "edges": [ { - "source": "dndnode_7", - "sourceHandle": "OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel", - "target": "dndnode_8", - "targetHandle": "BaseLLM|llm|dndnode_8", + "source": "dndnode_27", + "sourceHandle": "Tool|dndnode_27|Tool", + "target": "dndnode_3", + "targetHandle": "Tool|allowed_tools|dndnode_3", "className": "animate-pulse", - "id": "reactflow__edge-dndnode_7OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel-dndnode_8BaseLLM|llm|dndnode_8" + "id": "reactflow__edge-dndnode_27Tool|dndnode_27|Tool-dndnode_3Tool|allowed_tools|dndnode_3" }, { - "source": "dndnode_8", - "sourceHandle": "PAL-MATH|dndnode_8|Tool", - "target": "dndnode_34", - "targetHandle": "Tool|allowed_tools|dndnode_34", + "source": "dndnode_3", + "sourceHandle": "ZeroShotAgent|dndnode_3|Agent|function", + "target": "dndnode_28", + "targetHandle": "function|func|dndnode_28", "className": "animate-pulse", - "id": "reactflow__edge-dndnode_8PAL-MATH|dndnode_8|Tool-dndnode_34Tool|allowed_tools|dndnode_34" + "id": "reactflow__edge-dndnode_3ZeroShotAgent|dndnode_3|Agent|function-dndnode_28function|func|dndnode_28" }, { - "source": "dndnode_34", - "sourceHandle": "ZeroShotAgent|dndnode_34|Agent|function", + "source": "dndnode_40", + "sourceHandle": "PythonFunction|dndnode_40|function", + "target": "dndnode_27", + "targetHandle": "function|func|dndnode_27", + "className": "animate-pulse", + "id": "reactflow__edge-dndnode_40PythonFunction|dndnode_40|function-dndnode_27function|func|dndnode_27" + }, + { + "source": "dndnode_28", + "sourceHandle": "Tool|dndnode_28|Tool", "target": "dndnode_41", - "targetHandle": "function|func|dndnode_41", + "targetHandle": "Tool|allowed_tools|dndnode_41", "className": "animate-pulse", - "id": "reactflow__edge-dndnode_34ZeroShotAgent|dndnode_34|Agent|function-dndnode_41function|func|dndnode_41" + "id": "reactflow__edge-dndnode_28Tool|dndnode_28|Tool-dndnode_41Tool|allowed_tools|dndnode_41" }, { - "source": "dndnode_41", - "sourceHandle": "BaseTool|dndnode_41|Tool", - "target": "dndnode_42", - "targetHandle": "Tool|allowed_tools|dndnode_42", - "className": "animate-pulse", - "id": "reactflow__edge-dndnode_41BaseTool|dndnode_41|Tool-dndnode_42Tool|allowed_tools|dndnode_42" - }, - { - "source": "dndnode_9", - "sourceHandle": "ZeroShotPrompt|dndnode_9|BasePromptTemplate", + "source": "dndnode_42", + "sourceHandle": "ZeroShotPrompt|dndnode_42|BasePromptTemplate", "target": "dndnode_43", "targetHandle": "BasePromptTemplate|prompt|dndnode_43", "className": "animate-pulse", - "id": "reactflow__edge-dndnode_9ZeroShotPrompt|dndnode_9|BasePromptTemplate-dndnode_43BasePromptTemplate|prompt|dndnode_43" + "id": "reactflow__edge-dndnode_42ZeroShotPrompt|dndnode_42|BasePromptTemplate-dndnode_43BasePromptTemplate|prompt|dndnode_43" }, { - "source": "dndnode_7", - "sourceHandle": "OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel", + "source": "dndnode_44", + "sourceHandle": "OpenAI|dndnode_44|BaseOpenAI|BaseLLM|BaseLanguageModel", "target": "dndnode_43", "targetHandle": "BaseLanguageModel|llm|dndnode_43", "className": "animate-pulse", - "id": "reactflow__edge-dndnode_7OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel-dndnode_43BaseLanguageModel|llm|dndnode_43" + "id": "reactflow__edge-dndnode_44OpenAI|dndnode_44|BaseOpenAI|BaseLLM|BaseLanguageModel-dndnode_43BaseLanguageModel|llm|dndnode_43" }, { "source": "dndnode_43", "sourceHandle": "LLMChain|dndnode_43|Chain", - "target": "dndnode_34", - "targetHandle": "LLMChain|llm_chain|dndnode_34", + "target": "dndnode_3", + "targetHandle": "LLMChain|llm_chain|dndnode_3", "className": "animate-pulse", - "id": "reactflow__edge-dndnode_43LLMChain|dndnode_43|Chain-dndnode_34LLMChain|llm_chain|dndnode_34" + "id": "reactflow__edge-dndnode_43LLMChain|dndnode_43|Chain-dndnode_3LLMChain|llm_chain|dndnode_3" }, { - "source": "dndnode_44", - "sourceHandle": "LLMChain|dndnode_44|Chain", - "target": "dndnode_42", - "targetHandle": "LLMChain|llm_chain|dndnode_42", + "source": "dndnode_43", + "sourceHandle": "LLMChain|dndnode_43|Chain", + "target": "dndnode_41", + "targetHandle": "LLMChain|llm_chain|dndnode_41", "className": "animate-pulse", - "id": "reactflow__edge-dndnode_44LLMChain|dndnode_44|Chain-dndnode_42LLMChain|llm_chain|dndnode_42" - }, - { - "source": "dndnode_9", - "sourceHandle": "ZeroShotPrompt|dndnode_9|BasePromptTemplate", - "target": "dndnode_44", - "targetHandle": "BasePromptTemplate|prompt|dndnode_44", - "className": "animate-pulse", - "id": "reactflow__edge-dndnode_9ZeroShotPrompt|dndnode_9|BasePromptTemplate-dndnode_44BasePromptTemplate|prompt|dndnode_44" - }, - { - "source": "dndnode_7", - "sourceHandle": "OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel", - "target": "dndnode_44", - "targetHandle": "BaseLanguageModel|llm|dndnode_44", - "className": "animate-pulse", - "id": "reactflow__edge-dndnode_7OpenAI|dndnode_7|BaseOpenAI|BaseLLM|BaseLanguageModel-dndnode_44BaseLanguageModel|llm|dndnode_44" + "id": "reactflow__edge-dndnode_43LLMChain|dndnode_43|Chain-dndnode_41LLMChain|llm_chain|dndnode_41" } ], "viewport": { - "x": -12.25513998834083, - "y": 135.5224300058294, + "x": 250.11733592862913, + "y": 349.94447213104604, "zoom": 0.5 } }, diff --git a/tests/test_graph.py b/tests/test_graph.py index 3bdf72c45..e0888c0c4 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -81,36 +81,45 @@ def test_get_node_neighbors_complex(): # Get root node root = get_root_node(graph) assert root is not None - neighbors = graph.get_node_neighbors(root) + neighbors = graph.get_nodes_with_target(root) assert neighbors is not None - assert isinstance(neighbors, dict) + # Neighbors should be a list of nodes + assert isinstance(neighbors, list) # Root Node is an Agent, it requires an LLMChain and tools # We need to check if there is a Chain in the one of the neighbors' - # data attribute in the type key - assert any( - "Chain" in neighbor.data["type"] for neighbor, val in neighbors.items() if val + assert any("Chain" in neighbor.data["type"] for neighbor in neighbors) + # assert Tool is in the neighbors + assert any("Tool" in neighbor.data["type"] for neighbor in neighbors) + # Now on to the Chain's neighbors + chain = next(neighbor for neighbor in neighbors if "Chain" in neighbor.data["type"]) + chain_neighbors = graph.get_nodes_with_target(chain) + assert chain_neighbors is not None + # Check if there is a LLM in the chain's neighbors + assert any("OpenAI" in neighbor.data["type"] for neighbor in chain_neighbors) + # Chain should have a Prompt as a neighbor + assert any("Prompt" in neighbor.data["type"] for neighbor in chain_neighbors) + # Now on to the Tool's neighbors + tool = next(neighbor for neighbor in neighbors if "Tool" in neighbor.data["type"]) + tool_neighbors = graph.get_nodes_with_target(tool) + assert tool_neighbors is not None + # Check if there is an Agent in the tool's neighbors + assert any("Agent" in neighbor.data["type"] for neighbor in tool_neighbors) + # This Agent has a Tool that has a PythonFunction as func + agent = next( + neighbor for neighbor in tool_neighbors if "Agent" in neighbor.data["type"] ) - # assert BaseTool is in the neighbors - assert any( - "BaseTool" in neighbor.data["type"] - for neighbor, val in neighbors.items() - if val - ) - # Now on to the BaseTool's neighbors - base_tool = next( - neighbor - for neighbor, val in neighbors.items() - if "BaseTool" in neighbor.data["type"] and val - ) - base_tool_neighbors = graph.get_node_neighbors(base_tool) - assert base_tool_neighbors is not None - assert isinstance(base_tool_neighbors, dict) - # Check if there is an ZeroShotAgent in the base_tool's neighbors - assert any( - "ZeroShotAgent" in neighbor.data["type"] - for neighbor, val in base_tool_neighbors.items() - if val + agent_neighbors = graph.get_nodes_with_target(agent) + assert agent_neighbors is not None + # Check if there is a Tool in the agent's neighbors + assert any("Tool" in neighbor.data["type"] for neighbor in agent_neighbors) + # This Tool has a PythonFunction as func + tool = next( + neighbor for neighbor in agent_neighbors if "Tool" in neighbor.data["type"] ) + tool_neighbors = graph.get_nodes_with_target(tool) + assert tool_neighbors is not None + # Check if there is a PythonFunction in the tool's neighbors + assert any("PythonFunction" in neighbor.data["type"] for neighbor in tool_neighbors) def test_get_node():