Source code for langchain_community.utilities.graphql

import json
from typing import Any, Callable, Dict, Optional

from langchain_core.pydantic_v1 import BaseModel, Extra, root_validator


class GraphQLAPIWrapper(BaseModel):
    """封装了GraphQL API。

    要使用，您应该已安装``gql`` python包。
    此封装将使用GraphQL API执行查询。"""

    custom_headers: Optional[Dict[str, str]] = None
    fetch_schema_from_transport: Optional[bool] = None
    graphql_endpoint: str
    gql_client: Any  #: :meta private:
    gql_function: Callable[[str], Any]  #: :meta private:

    class Config:
        """此pydantic对象的配置。"""

        extra = Extra.forbid

    @root_validator(pre=True)
    def validate_environment(cls, values: Dict) -> Dict:
        """验证Python包是否存在于环境中。"""
        try:
            from gql import Client, gql
            from gql.transport.requests import RequestsHTTPTransport
        except ImportError as e:
            raise ImportError(
                "Could not import gql python package. "
                f"Try installing it with `pip install gql`. Received error: {e}"
            )
        headers = values.get("custom_headers")
        transport = RequestsHTTPTransport(
            url=values["graphql_endpoint"],
            headers=headers,
        )
        fetch_schema_from_transport = values.get("fetch_schema_from_transport", True)
        client = Client(
            transport=transport, fetch_schema_from_transport=fetch_schema_from_transport
        )
        values["gql_client"] = client
        values["gql_function"] = gql
        return values

    def run(self, query: str) -> str:
        """运行一个GraphQL查询并获取结果。"""
        result = self._execute_query(query)
        return json.dumps(result, indent=2)

    def _execute_query(self, query: str) -> Dict[str, Any]:
        """执行一个GraphQL查询并返回结果。"""
        document_node = self.gql_function(query)
        result = self.gql_client.execute(document_node)
        return result