269 lines
7 KiB
Python
269 lines
7 KiB
Python
from fastapi.testclient import TestClient
|
|
|
|
|
|
def test_zero_shot_agent(client: TestClient):
|
|
response = client.get("api/v1/all")
|
|
assert response.status_code == 200
|
|
json_response = response.json()
|
|
agents = json_response["agents"]
|
|
|
|
zero_shot_agent = agents["ZeroShotAgent"]
|
|
assert set(zero_shot_agent["base_classes"]) == {
|
|
"ZeroShotAgent",
|
|
"BaseSingleActionAgent",
|
|
"Agent",
|
|
"function",
|
|
}
|
|
template = zero_shot_agent["template"]
|
|
|
|
assert template["tools"] == {
|
|
"required": True,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "tools",
|
|
"type": "BaseTool",
|
|
"list": True,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
|
|
# Additional assertions for other template variables
|
|
assert template["callback_manager"] == {
|
|
"required": False,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": False,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "callback_manager",
|
|
"type": "BaseCallbackManager",
|
|
"list": False,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["llm"] == {
|
|
"required": True,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "llm",
|
|
"type": "BaseLanguageModel",
|
|
"list": False,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["output_parser"] == {
|
|
"required": False,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": False,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "output_parser",
|
|
"type": "AgentOutputParser",
|
|
"list": False,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["input_variables"] == {
|
|
"required": False,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": False,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "input_variables",
|
|
"type": "str",
|
|
"list": True,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["prefix"] == {
|
|
"required": False,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": True,
|
|
"value": "Answer the following questions as best you can. You have access to the following tools:",
|
|
"password": False,
|
|
"name": "prefix",
|
|
"type": "str",
|
|
"list": False,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["suffix"] == {
|
|
"required": False,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": True,
|
|
"value": "Begin!\n\nQuestion: {input}\nThought:{agent_scratchpad}",
|
|
"password": False,
|
|
"name": "suffix",
|
|
"type": "str",
|
|
"list": False,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
|
|
|
|
def test_json_agent(client: TestClient):
|
|
response = client.get("api/v1/all")
|
|
assert response.status_code == 200
|
|
json_response = response.json()
|
|
agents = json_response["agents"]
|
|
|
|
json_agent = agents["JsonAgent"]
|
|
assert json_agent["base_classes"] == ["AgentExecutor"]
|
|
template = json_agent["template"]
|
|
|
|
assert template["toolkit"] == {
|
|
"required": True,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "toolkit",
|
|
"type": "BaseToolkit",
|
|
"list": False,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["llm"] == {
|
|
"required": True,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "llm",
|
|
"type": "BaseLanguageModel",
|
|
"list": False,
|
|
"advanced": False,
|
|
"display_name": "LLM",
|
|
"info": "",
|
|
}
|
|
|
|
|
|
def test_csv_agent(client: TestClient):
|
|
response = client.get("api/v1/all")
|
|
assert response.status_code == 200
|
|
json_response = response.json()
|
|
agents = json_response["agents"]
|
|
|
|
csv_agent = agents["CSVAgent"]
|
|
assert csv_agent["base_classes"] == ["AgentExecutor"]
|
|
template = csv_agent["template"]
|
|
|
|
assert template["path"] == {
|
|
"required": True,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"value": "",
|
|
"suffixes": [".csv"],
|
|
"fileTypes": ["csv"],
|
|
"password": False,
|
|
"name": "path",
|
|
"type": "file",
|
|
"list": False,
|
|
"file_path": None,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["llm"] == {
|
|
"required": True,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "llm",
|
|
"type": "BaseLanguageModel",
|
|
"list": False,
|
|
"advanced": False,
|
|
"display_name": "LLM",
|
|
"info": "",
|
|
}
|
|
|
|
|
|
def test_initialize_agent(client: TestClient):
|
|
response = client.get("api/v1/all")
|
|
assert response.status_code == 200
|
|
json_response = response.json()
|
|
agents = json_response["agents"]
|
|
|
|
initialize_agent = agents["AgentInitializer"]
|
|
assert initialize_agent["base_classes"] == ["AgentExecutor", "function"]
|
|
template = initialize_agent["template"]
|
|
|
|
assert template["agent"] == {
|
|
"required": True,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"value": "zero-shot-react-description",
|
|
"password": False,
|
|
"options": [
|
|
"zero-shot-react-description",
|
|
"react-docstore",
|
|
"self-ask-with-search",
|
|
"conversational-react-description",
|
|
"openai-functions",
|
|
"openai-multi-functions",
|
|
],
|
|
"name": "agent",
|
|
"type": "str",
|
|
"list": True,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["memory"] == {
|
|
"required": False,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "memory",
|
|
"type": "BaseChatMemory",
|
|
"list": False,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["tools"] == {
|
|
"required": False,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "tools",
|
|
"type": "Tool",
|
|
"list": True,
|
|
"advanced": False,
|
|
"info": "",
|
|
}
|
|
assert template["llm"] == {
|
|
"required": True,
|
|
"dynamic": False,
|
|
"placeholder": "",
|
|
"show": True,
|
|
"multiline": False,
|
|
"password": False,
|
|
"name": "llm",
|
|
"type": "BaseLanguageModel",
|
|
"list": False,
|
|
"advanced": False,
|
|
"display_name": "LLM",
|
|
"info": "",
|
|
}
|