Spaces:
Runtime error
Runtime error
File size: 3,543 Bytes
ed4d993 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
from typing import Any, Dict, List, Optional
from langchain_community.graphs.graph_store import GraphStore
class TigerGraph(GraphStore):
"""TigerGraph wrapper for graph operations.
*Security note*: Make sure that the database connection uses credentials
that are narrowly-scoped to only include necessary permissions.
Failure to do so may result in data corruption or loss, since the calling
code may attempt commands that would result in deletion, mutation
of data if appropriately prompted or reading sensitive data if such
data is present in the database.
The best way to guard against such negative outcomes is to (as appropriate)
limit the permissions granted to the credentials used with this tool.
See https://python.langchain.com/docs/security for more information.
"""
def __init__(self, conn: Any) -> None:
"""Create a new TigerGraph graph wrapper instance."""
self.set_connection(conn)
self.set_schema()
@property
def conn(self) -> Any:
return self._conn
@property
def schema(self) -> Dict[str, Any]:
return self._schema
def get_schema(self) -> str: # type: ignore[override]
if self._schema:
return str(self._schema)
else:
self.set_schema()
return str(self._schema)
def set_connection(self, conn: Any) -> None:
try:
from pyTigerGraph import TigerGraphConnection
except ImportError:
raise ImportError(
"Could not import pyTigerGraph python package. "
"Please install it with `pip install pyTigerGraph`."
)
if not isinstance(conn, TigerGraphConnection):
msg = "**conn** parameter must inherit from TigerGraphConnection"
raise TypeError(msg)
if conn.ai.nlqs_host is None:
msg = """**conn** parameter does not have nlqs_host parameter defined.
Define hostname of NLQS service."""
raise ConnectionError(msg)
self._conn: TigerGraphConnection = conn
self.set_schema()
def set_schema(self, schema: Optional[Dict[str, Any]] = None) -> None:
"""
Set the schema of the TigerGraph Database.
Auto-generates Schema if **schema** is None.
"""
self._schema = self.generate_schema() if schema is None else schema
def generate_schema(
self,
) -> Dict[str, List[Dict[str, Any]]]:
"""
Generates the schema of the TigerGraph Database and returns it
User can specify a **sample_ratio** (0 to 1) to determine the
ratio of documents/edges used (in relation to the Collection size)
to render each Collection Schema.
"""
return self._conn.getSchema(force=True)
def refresh_schema(self): # type: ignore[no-untyped-def]
self.generate_schema()
def query(self, query: str) -> Dict[str, Any]: # type: ignore[override]
"""Query the TigerGraph database."""
answer = self._conn.ai.query(query)
return answer
def register_query(
self,
function_header: str,
description: str,
docstring: str,
param_types: dict = {},
) -> List[str]:
"""
Wrapper function to register a custom GSQL query to the TigerGraph NLQS.
"""
return self._conn.ai.registerCustomQuery(
function_header, description, docstring, param_types
)
|