Refactor imports and update function names
This commit is contained in:
Gabriel Luiz Freitas Almeida
parent
a41ada87ee
commit
abe4925cc3
8 changed files with 156 additions and 314 deletions
|
|
@ -1,18 +1,17 @@
|
|||
from collections import namedtuple
|
||||
import time
|
||||
import uuid
|
||||
from langflow.processing.process import Result
|
||||
from langflow.services.auth.utils import get_password_hash
|
||||
from langflow.services.database.models.api_key.api_key import ApiKey
|
||||
from langflow.services.getters import get_settings_service
|
||||
from langflow.services.database.utils import session_getter
|
||||
from langflow.services.getters import get_db_service
|
||||
from collections import namedtuple
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from langflow.interface.tools.constants import CUSTOM_TOOLS
|
||||
from langflow.processing.process import Result
|
||||
from langflow.services.auth.utils import get_password_hash
|
||||
from langflow.services.database.models.api_key.model import ApiKey
|
||||
from langflow.services.database.utils import session_getter
|
||||
from langflow.services.deps import get_db_service, get_settings_service
|
||||
from langflow.template.frontend_node.chains import TimeTravelGuideChainNode
|
||||
|
||||
import time
|
||||
|
||||
|
||||
def run_post(client, flow_id, headers, post_data):
|
||||
response = client.post(
|
||||
|
|
@ -31,10 +30,7 @@ def poll_task_status(client, headers, href, max_attempts=20, sleep_time=1):
|
|||
href,
|
||||
headers=headers,
|
||||
)
|
||||
if (
|
||||
task_status_response.status_code == 200
|
||||
and task_status_response.json()["status"] == "SUCCESS"
|
||||
):
|
||||
if task_status_response.status_code == 200 and task_status_response.json()["status"] == "SUCCESS":
|
||||
return task_status_response.json()
|
||||
time.sleep(sleep_time)
|
||||
return None # Return None if task did not complete in time
|
||||
|
|
@ -127,11 +123,7 @@ def created_api_key(active_user):
|
|||
)
|
||||
db_manager = get_db_service()
|
||||
with session_getter(db_manager) as session:
|
||||
if (
|
||||
existing_api_key := session.query(ApiKey)
|
||||
.filter(ApiKey.api_key == api_key.api_key)
|
||||
.first()
|
||||
):
|
||||
if existing_api_key := session.query(ApiKey).filter(ApiKey.api_key == api_key.api_key).first():
|
||||
return existing_api_key
|
||||
session.add(api_key)
|
||||
session.commit()
|
||||
|
|
@ -190,9 +182,7 @@ def test_process_flow_invalid_id(client, monkeypatch, created_api_key):
|
|||
}
|
||||
|
||||
invalid_id = uuid.uuid4()
|
||||
response = client.post(
|
||||
f"api/v1/process/{invalid_id}", headers=headers, json=post_data
|
||||
)
|
||||
response = client.post(f"api/v1/process/{invalid_id}", headers=headers, json=post_data)
|
||||
|
||||
assert response.status_code == 404
|
||||
assert f"Flow {invalid_id} not found" in response.json()["detail"]
|
||||
|
|
@ -233,9 +223,7 @@ def test_process_flow_without_autologin(client, flow, monkeypatch, created_api_k
|
|||
|
||||
monkeypatch.setattr(endpoints, "process_graph_cached", mock_process_graph_cached)
|
||||
monkeypatch.setattr(crud, "update_total_uses", mock_update_total_uses)
|
||||
monkeypatch.setattr(
|
||||
endpoints, "process_graph_cached_task", mock_process_graph_cached_task
|
||||
)
|
||||
monkeypatch.setattr(endpoints, "process_graph_cached_task", mock_process_graph_cached_task)
|
||||
|
||||
api_key = created_api_key.api_key
|
||||
headers = {"x-api-key": api_key}
|
||||
|
|
@ -420,110 +408,77 @@ def test_various_prompts(client, prompt, expected_input_variables):
|
|||
|
||||
|
||||
def test_get_vertices_flow_not_found(client, logged_in_headers):
|
||||
response = client.get(
|
||||
"/api/v1/build/nonexistent_id/vertices", headers=logged_in_headers
|
||||
)
|
||||
assert (
|
||||
response.status_code == 500
|
||||
) # Or whatever status code you've set for invalid ID
|
||||
response = client.get("/api/v1/build/nonexistent_id/vertices", headers=logged_in_headers)
|
||||
assert response.status_code == 500 # Or whatever status code you've set for invalid ID
|
||||
|
||||
|
||||
def test_get_vertices(client, added_flow_with_prompt_and_history, logged_in_headers):
|
||||
flow_id = added_flow_with_prompt_and_history["id"]
|
||||
response = client.get(
|
||||
f"/api/v1/build/{flow_id}/vertices", headers=logged_in_headers
|
||||
)
|
||||
response = client.get(f"/api/v1/build/{flow_id}/vertices", headers=logged_in_headers)
|
||||
assert response.status_code == 200
|
||||
assert "ids" in response.json()
|
||||
# The response should contain the list in this order
|
||||
# ['ConversationBufferMemory-Lu2Nb', 'PromptTemplate-5Q0W8', 'ChatOpenAI-vy7fV', 'LLMChain-UjBh1']
|
||||
# The important part is before the - (ConversationBufferMemory, PromptTemplate, ChatOpenAI, LLMChain)
|
||||
ids = [id.split("-")[0] for id in response.json()["ids"]]
|
||||
assert ids == [
|
||||
"ConversationBufferMemory",
|
||||
"PromptTemplate",
|
||||
"ChatOpenAI",
|
||||
"LLMChain",
|
||||
]
|
||||
ids = [inner_id.split("-")[0] for _id in response.json()["ids"] for inner_id in _id]
|
||||
assert ids == ["ChatOpenAI", "PromptTemplate", "ConversationBufferMemory", "LLMChain"]
|
||||
|
||||
|
||||
def test_build_vertex_invalid_flow_id(client, logged_in_headers):
|
||||
response = client.post(
|
||||
"/api/v1/build/nonexistent_id/vertices/vertex_id", headers=logged_in_headers
|
||||
)
|
||||
response = client.post("/api/v1/build/nonexistent_id/vertices/vertex_id", headers=logged_in_headers)
|
||||
assert response.status_code == 500
|
||||
|
||||
|
||||
def test_build_vertex_invalid_vertex_id(
|
||||
client, added_flow_with_prompt_and_history, logged_in_headers
|
||||
):
|
||||
def test_build_vertex_invalid_vertex_id(client, added_flow_with_prompt_and_history, logged_in_headers):
|
||||
flow_id = added_flow_with_prompt_and_history["id"]
|
||||
response = client.post(
|
||||
f"/api/v1/build/{flow_id}/vertices/invalid_vertex_id", headers=logged_in_headers
|
||||
)
|
||||
response = client.post(f"/api/v1/build/{flow_id}/vertices/invalid_vertex_id", headers=logged_in_headers)
|
||||
assert response.status_code == 500
|
||||
|
||||
|
||||
def test_build_all_vertices_in_sequence_with_chat_input(
|
||||
client, added_flow_chat_input, logged_in_headers
|
||||
):
|
||||
def test_build_all_vertices_in_sequence_with_chat_input(client, added_flow_chat_input, logged_in_headers):
|
||||
flow_id = added_flow_chat_input["id"]
|
||||
|
||||
# First, get all the vertices in the correct sequence
|
||||
response = client.get(
|
||||
f"/api/v1/build/{flow_id}/vertices", headers=logged_in_headers
|
||||
)
|
||||
response = client.get(f"/api/v1/build/{flow_id}/vertices", headers=logged_in_headers)
|
||||
assert response.status_code == 200
|
||||
assert "ids" in response.json()
|
||||
vertex_ids = response.json()["ids"]
|
||||
|
||||
# Now, iterate through each vertex and build it
|
||||
for vertex_id in vertex_ids:
|
||||
response = client.post(
|
||||
f"/api/v1/build/{flow_id}/vertices/{vertex_id}", headers=logged_in_headers
|
||||
)
|
||||
response = client.post(f"/api/v1/build/{flow_id}/vertices/{vertex_id}", headers=logged_in_headers)
|
||||
json_response = response.json()
|
||||
assert (
|
||||
response.status_code == 200
|
||||
), f"Failed at vertex {vertex_id}: {json_response}"
|
||||
assert response.status_code == 200, f"Failed at vertex {vertex_id}: {json_response}"
|
||||
assert "valid" in json_response
|
||||
assert json_response["valid"], json_response["params"]
|
||||
|
||||
|
||||
def test_build_all_vertices_in_sequence_with_two_outputs(
|
||||
client, added_flow_two_outputs, logged_in_headers
|
||||
):
|
||||
def test_build_all_vertices_in_sequence_with_two_outputs(client, added_flow_two_outputs, logged_in_headers):
|
||||
"""This tests the case where a node has two outputs, one of which is Text and the other (in this case) is
|
||||
a LLMChain. We need to make sure the correct output is passed in both cases."""
|
||||
flow_id = added_flow_two_outputs["id"]
|
||||
|
||||
# First, get all the vertices in the correct sequence
|
||||
response = client.get(
|
||||
f"/api/v1/build/{flow_id}/vertices", headers=logged_in_headers
|
||||
)
|
||||
response = client.get(f"/api/v1/build/{flow_id}/vertices", headers=logged_in_headers)
|
||||
assert response.status_code == 200
|
||||
assert "ids" in response.json()
|
||||
vertex_ids = response.json()["ids"]
|
||||
|
||||
# Now, iterate through each vertex and build it
|
||||
for vertex_id in vertex_ids:
|
||||
response = client.post(
|
||||
f"/api/v1/build/{flow_id}/vertices/{vertex_id}", headers=logged_in_headers
|
||||
)
|
||||
response = client.post(f"/api/v1/build/{flow_id}/vertices/{vertex_id}", headers=logged_in_headers)
|
||||
json_response = response.json()
|
||||
assert (
|
||||
response.status_code == 200
|
||||
), f"Failed at vertex {vertex_id}: {json_response}"
|
||||
assert response.status_code == 200, f"Failed at vertex {vertex_id}: {json_response}"
|
||||
assert "valid" in json_response
|
||||
assert json_response["valid"], json_response["params"]
|
||||
|
||||
|
||||
def test_basic_chat_in_process(client, added_flow, created_api_key):
|
||||
def test_basic_chat_in_process(client, flow, created_api_key):
|
||||
# Run the /api/v1/process/{flow_id} endpoint
|
||||
headers = {"x-api-key": created_api_key.api_key}
|
||||
post_data = {"inputs": {"text": "Hi, My name is Gabriel"}}
|
||||
response = client.post(
|
||||
f"api/v1/process/{added_flow.get('id')}",
|
||||
f"api/v1/process/{flow.get('id')}",
|
||||
headers=headers,
|
||||
json=post_data,
|
||||
)
|
||||
|
|
@ -540,7 +495,7 @@ def test_basic_chat_in_process(client, added_flow, created_api_key):
|
|||
"session_id": response.json()["session_id"],
|
||||
}
|
||||
response = client.post(
|
||||
f"api/v1/process/{added_flow.get('id')}",
|
||||
f"api/v1/process/{flow.get('id')}",
|
||||
headers=headers,
|
||||
json=post_data,
|
||||
)
|
||||
|
|
@ -548,12 +503,12 @@ def test_basic_chat_in_process(client, added_flow, created_api_key):
|
|||
assert "Gabriel" in response.json()["result"]["text"]
|
||||
|
||||
|
||||
def test_basic_chat_different_session_ids(client, added_flow, created_api_key):
|
||||
def test_basic_chat_different_session_ids(client, flow, created_api_key):
|
||||
# Run the /api/v1/process/{flow_id} endpoint
|
||||
headers = {"x-api-key": created_api_key.api_key}
|
||||
post_data = {"inputs": {"text": "Hi, My name is Gabriel"}}
|
||||
response = client.post(
|
||||
f"api/v1/process/{added_flow.get('id')}",
|
||||
f"api/v1/process/{flow.get('id')}",
|
||||
headers=headers,
|
||||
json=post_data,
|
||||
)
|
||||
|
|
@ -570,7 +525,7 @@ def test_basic_chat_different_session_ids(client, added_flow, created_api_key):
|
|||
"inputs": {"text": "What is my name?"},
|
||||
}
|
||||
response = client.post(
|
||||
f"api/v1/process/{added_flow.get('id')}",
|
||||
f"api/v1/process/{flow.get('id')}",
|
||||
headers=headers,
|
||||
json=post_data,
|
||||
)
|
||||
|
|
@ -579,9 +534,9 @@ def test_basic_chat_different_session_ids(client, added_flow, created_api_key):
|
|||
assert session_id1 != response.json()["session_id"]
|
||||
|
||||
|
||||
def test_basic_chat_with_two_session_ids_and_names(client, added_flow, created_api_key):
|
||||
def test_basic_chat_with_two_session_ids_and_names(client, flow, created_api_key):
|
||||
headers = {"x-api-key": created_api_key.api_key}
|
||||
flow_id = added_flow.get("id")
|
||||
flow_id = flow.get("id")
|
||||
names = ["Gabriel", "John"]
|
||||
session_ids = []
|
||||
|
||||
|
|
@ -606,9 +561,7 @@ def test_basic_chat_with_two_session_ids_and_names(client, added_flow, created_a
|
|||
|
||||
|
||||
@pytest.mark.async_test
|
||||
def test_vector_store_in_process(
|
||||
distributed_client, added_vector_store, created_api_key
|
||||
):
|
||||
def test_vector_store_in_process(distributed_client, added_vector_store, created_api_key):
|
||||
# Run the /api/v1/process/{flow_id} endpoint
|
||||
headers = {"x-api-key": created_api_key.api_key}
|
||||
post_data = {"inputs": {"input": "What is Langflow?"}}
|
||||
|
|
@ -627,13 +580,13 @@ def test_vector_store_in_process(
|
|||
|
||||
# Test function without loop
|
||||
@pytest.mark.async_test
|
||||
def test_async_task_processing(distributed_client, added_flow, created_api_key):
|
||||
def test_async_task_processing(distributed_client, flow, created_api_key):
|
||||
headers = {"x-api-key": created_api_key.api_key}
|
||||
post_data = {"inputs": {"text": "Hi, My name is Gabriel"}}
|
||||
|
||||
flow = flow.model_dump()
|
||||
# Run the /api/v1/process/{flow_id} endpoint with sync=False
|
||||
response = distributed_client.post(
|
||||
f"api/v1/process/{added_flow.get('id')}",
|
||||
f"api/v1/process/{flow.get('id')}",
|
||||
headers=headers,
|
||||
json={**post_data, "sync": False},
|
||||
)
|
||||
|
|
@ -659,9 +612,7 @@ def test_async_task_processing(distributed_client, added_flow, created_api_key):
|
|||
|
||||
# Test function without loop
|
||||
@pytest.mark.async_test
|
||||
def test_async_task_processing_vector_store(
|
||||
client, added_vector_store, created_api_key
|
||||
):
|
||||
def test_async_task_processing_vector_store(client, added_vector_store, created_api_key):
|
||||
headers = {"x-api-key": created_api_key.api_key}
|
||||
post_data = {"inputs": {"input": "How do I upload examples?"}}
|
||||
|
||||
|
|
@ -690,6 +641,4 @@ def test_async_task_processing_vector_store(
|
|||
# Validate that the task completed successfully and the result is as expected
|
||||
assert "result" in task_status_json, task_status_json
|
||||
assert "output" in task_status_json["result"], task_status_json["result"]
|
||||
assert "Langflow" in task_status_json["result"]["output"], task_status_json[
|
||||
"result"
|
||||
]
|
||||
assert "Langflow" in task_status_json["result"]["output"], task_status_json["result"]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue