src/backend/tests/integration/utils.py

import dataclasses
import os
import uuid
from typing import Any

import requests
from astrapy.admin import parse_api_endpoint
from langflow.api.v1.schemas import InputValueRequest
from langflow.custom import Component
from langflow.custom.eval import eval_custom_component_code
from langflow.field_typing import Embeddings
from langflow.graph import Graph
from langflow.processing.process import run_graph_internal


def check_env_vars(*env_vars):
    """Check if all specified environment variables are set.

    Args:
    *env_vars (str): The environment variables to check.

    Returns:
    bool: True if all environment variables are set, False otherwise.
    """
    return all(os.getenv(var) for var in env_vars)


def valid_nvidia_vectorize_region(api_endpoint: str) -> bool:
    """Check if the specified region is valid.

    Args:
        api_endpoint: The API endpoint to check.

    Returns:
        True if the region contains hosted nvidia models, False otherwise.
    """
    parsed_endpoint = parse_api_endpoint(api_endpoint)
    if not parsed_endpoint:
        msg = "Invalid ASTRA_DB_API_ENDPOINT"
        raise ValueError(msg)
    return parsed_endpoint.region == "us-east-2"


class MockEmbeddings(Embeddings):
    def __init__(self):
        self.embedded_documents = None
        self.embedded_query = None

    @staticmethod
    def mock_embedding(text: str):
        return [len(text) / 2, len(text) / 5, len(text) / 10]

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        self.embedded_documents = texts
        return [self.mock_embedding(text) for text in texts]

    def embed_query(self, text: str) -> list[float]:
        self.embedded_query = text
        return self.mock_embedding(text)


@dataclasses.dataclass
class JSONFlow:
    json: dict

    def get_components_by_type(self, component_type):
        result = [node["id"] for node in self.json["data"]["nodes"] if node["data"]["type"] == component_type]
        if not result:
            msg = (
                f"Component of type {component_type} not found, "
                f"available types: {', '.join({node['data']['type'] for node in self.json['data']['nodes']})}"
            )
            raise ValueError(msg)
        return result

    def get_component_by_type(self, component_type):
        components = self.get_components_by_type(component_type)
        if len(components) > 1:
            msg = f"Multiple components of type {component_type} found"
            raise ValueError(msg)
        return components[0]

    def set_value(self, component_id, key, value):
        done = False
        for node in self.json["data"]["nodes"]:
            if node["id"] == component_id:
                if key not in node["data"]["node"]["template"]:
                    msg = f"Component {component_id} does not have input {key}"
                    raise ValueError(msg)
                node["data"]["node"]["template"][key]["value"] = value
                node["data"]["node"]["template"][key]["load_from_db"] = False
                done = True
                break
        if not done:
            msg = f"Component {component_id} not found"
            raise ValueError(msg)


def download_flow_from_github(name: str, version: str) -> JSONFlow:
    response = requests.get(
        f"https://raw.githubusercontent.com/langflow-ai/langflow/v{version}/src/backend/base/langflow/initial_setup/starter_projects/{name}.json",
        timeout=10,
    )
    response.raise_for_status()
    as_json = response.json()
    return JSONFlow(json=as_json)


def download_component_from_github(module: str, file_name: str, version: str) -> Component:
    version_string = f"v{version}" if version != "main" else version
    response = requests.get(
        f"https://raw.githubusercontent.com/langflow-ai/langflow/{version_string}/src/backend/base/langflow/components/{module}/{file_name}.py",
        timeout=10,
    )
    response.raise_for_status()
    return Component(_code=response.text)


async def run_json_flow(
    json_flow: JSONFlow, run_input: Any | None = None, session_id: str | None = None
) -> dict[str, Any]:
    graph = Graph.from_payload(json_flow.json)
    return await run_flow(graph, run_input, session_id)


async def run_flow(graph: Graph, run_input: Any | None = None, session_id: str | None = None) -> dict[str, Any]:
    graph.prepare()
    graph_run_inputs = [InputValueRequest(input_value=run_input, type="chat")] if run_input else []

    flow_id = str(uuid.uuid4())

    results, _ = await run_graph_internal(graph, flow_id, session_id=session_id, inputs=graph_run_inputs)
    outputs = {}
    for r in results:
        for out in r.outputs:
            outputs |= out.results
    return outputs


@dataclasses.dataclass
class ComponentInputHandle:
    clazz: type
    inputs: dict
    output_name: str


async def run_single_component(
    clazz: type,
    inputs: dict | None = None,
    run_input: Any | None = None,
    session_id: str | None = None,
    input_type: str | None = "chat",
) -> dict[str, Any]:
    user_id = str(uuid.uuid4())
    flow_id = str(uuid.uuid4())
    graph = Graph(user_id=user_id, flow_id=flow_id)

    def _add_component(clazz: type, inputs: dict | None = None) -> str:
        raw_inputs = {}
        if inputs:
            for key, value in inputs.items():
                if not isinstance(value, ComponentInputHandle):
                    raw_inputs[key] = value
                if isinstance(value, Component):
                    msg = "Component inputs must be wrapped in ComponentInputHandle"
                    raise TypeError(msg)
        component = clazz(**raw_inputs, _user_id=user_id)
        component_id = graph.add_component(component)
        if inputs:
            for input_name, handle in inputs.items():
                if isinstance(handle, ComponentInputHandle):
                    handle_component_id = _add_component(handle.clazz, handle.inputs)
                    graph.add_component_edge(handle_component_id, (handle.output_name, input_name), component_id)
        return component_id

    component_id = _add_component(clazz, inputs)
    graph.prepare()
    graph_run_inputs = [InputValueRequest(input_value=run_input, type=input_type)] if run_input else []

    _, _ = await run_graph_internal(
        graph, flow_id, session_id=session_id, inputs=graph_run_inputs, outputs=[component_id]
    )
    return graph.get_vertex(component_id).built_object


def build_component_instance_for_tests(version: str, module: str, file_name: str, **kwargs):
    component = download_component_from_github(module, file_name, version)
    cc_class = eval_custom_component_code(component._code)
    return cc_class(**kwargs), component._code