jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Fix import statements and formatting issues

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-25 13:03:42 -03:00
commit a6a69a0f3a
24 changed files with 924 additions and 527 deletions
Download patch file Download diff file Expand all files Collapse all files

158
tests/test_endpoints.py
View file

@ -1,17 +1,18 @@
import time
import uuid
from collections import namedtuple
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
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(
@ -24,13 +25,16 @@ def run_post(client, flow_id, headers, post_data):
# Helper function to poll task status
def poll_task_status(client, headers, href, max_attempts=20, sleep_time=2):
def poll_task_status(client, headers, href, max_attempts=20, sleep_time=1):
for _ in range(max_attempts):
task_status_response = client.get(
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
@ -72,14 +76,11 @@ PROMPT_REQUEST = {
"text-ada-001",
],
"ChatOpenAI": [
"gpt-3.5-turbo-0613",
"gpt-3.5-turbo",
"gpt-3.5-turbo-16k-0613",
"gpt-3.5-turbo-16k",
"gpt-4-0613",
"gpt-4-32k-0613",
"gpt-4-1106-preview",
"gpt-4",
"gpt-4-32k",
"gpt-3.5-turbo",
"gpt-3.5-turbo-16k",
],
"Anthropic": [
"claude-v1",
@ -126,7 +127,11 @@ 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()
@ -185,7 +190,9 @@ 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"]
@ -226,7 +233,9 @@ 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}
@ -410,6 +419,105 @@ def test_various_prompts(client, prompt, expected_input_variables):
assert response.json()["input_variables"] == 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
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
)
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",
]
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
)
assert response.status_code == 500
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
)
assert response.status_code == 500
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
)
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
)
json_response = response.json()
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
):
"""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
)
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
)
json_response = response.json()
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):
# Run the /api/v1/process/{flow_id} endpoint
headers = {"x-api-key": created_api_key.api_key}
@ -498,7 +606,9 @@ 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?"}}
@ -549,7 +659,9 @@ 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?"}}
@ -578,4 +690,6 @@ def test_async_task_processing_vector_store(client, added_vector_store, created_
# 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"
]