forked from Archives/langchain
Add KuzuQAChain (#6454)
This PR adds `KuzuGraph` and `KuzuQAChain` for interacting with [Kùzu database](https://github.com/kuzudb/kuzu). Kùzu is an in-process property graph database management system (GDBMS) built for query speed and scalability. The `KuzuGraph` and `KuzuQAChain` provide the same functionality as the existing integration with NebulaGraph and Neo4j and enables query generation and question answering over Kùzu database. A notebook example and a simple test case have also been added. --------- Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>master
parent
6e07283dd5
commit
0fce8ef178
@ -0,0 +1,363 @@
|
|||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"attachments": {},
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"# KuzuQAChain\n",
|
||||||
|
"\n",
|
||||||
|
"This notebook shows how to use LLMs to provide a natural language interface to [Kùzu](https://kuzudb.com) database."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"attachments": {},
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"[Kùzu](https://kuzudb.com) is an in-process property graph database management system. You can simply install it with `pip`:\n",
|
||||||
|
"\n",
|
||||||
|
"```bash\n",
|
||||||
|
"pip install kuzu\n",
|
||||||
|
"```\n",
|
||||||
|
"\n",
|
||||||
|
"Once installed, you can simply import it and start creating a database on the local machine and connect to it:\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 1,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"import kuzu\n",
|
||||||
|
"db = kuzu.Database(\"test_db\")\n",
|
||||||
|
"conn = kuzu.Connection(db)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"attachments": {},
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"First, we create the schema for a simple movie database:"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 2,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"<kuzu.query_result.QueryResult at 0x1066ff410>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 2,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"conn.execute(\"CREATE NODE TABLE Movie (name STRING, PRIMARY KEY(name))\")\n",
|
||||||
|
"conn.execute(\"CREATE NODE TABLE Person (name STRING, birthDate STRING, PRIMARY KEY(name))\")\n",
|
||||||
|
"conn.execute(\"CREATE REL TABLE ActedIn (FROM Person TO Movie)\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"attachments": {},
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"Then we can insert some data."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 3,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"<kuzu.query_result.QueryResult at 0x107016210>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 3,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"conn.execute(\"CREATE (:Person {name: 'Al Pacino', birthDate: '1940-04-25'})\")\n",
|
||||||
|
"conn.execute(\"CREATE (:Person {name: 'Robert De Niro', birthDate: '1943-08-17'})\")\n",
|
||||||
|
"conn.execute(\"CREATE (:Movie {name: 'The Godfather'})\")\n",
|
||||||
|
"conn.execute(\"CREATE (:Movie {name: 'The Godfather: Part II'})\")\n",
|
||||||
|
"conn.execute(\"CREATE (:Movie {name: 'The Godfather Coda: The Death of Michael Corleone'})\")\n",
|
||||||
|
"conn.execute(\"MATCH (p:Person), (m:Movie) WHERE p.name = 'Al Pacino' AND m.name = 'The Godfather' CREATE (p)-[:ActedIn]->(m)\")\n",
|
||||||
|
"conn.execute(\"MATCH (p:Person), (m:Movie) WHERE p.name = 'Al Pacino' AND m.name = 'The Godfather: Part II' CREATE (p)-[:ActedIn]->(m)\")\n",
|
||||||
|
"conn.execute(\"MATCH (p:Person), (m:Movie) WHERE p.name = 'Al Pacino' AND m.name = 'The Godfather Coda: The Death of Michael Corleone' CREATE (p)-[:ActedIn]->(m)\")\n",
|
||||||
|
"conn.execute(\"MATCH (p:Person), (m:Movie) WHERE p.name = 'Robert De Niro' AND m.name = 'The Godfather: Part II' CREATE (p)-[:ActedIn]->(m)\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"attachments": {},
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Creating `KuzuQAChain`\n",
|
||||||
|
"\n",
|
||||||
|
"We can now create the `KuzuGraph` and `KuzuQAChain`. To create the `KuzuGraph` we simply need to pass the database object to the `KuzuGraph` constructor."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 4,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"from langchain.chat_models import ChatOpenAI\n",
|
||||||
|
"from langchain.graphs import KuzuGraph\n",
|
||||||
|
"from langchain.chains import KuzuQAChain"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 5,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"graph = KuzuGraph(db)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 6,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"chain = KuzuQAChain.from_llm(\n",
|
||||||
|
" ChatOpenAI(temperature=0), graph=graph, verbose=True\n",
|
||||||
|
")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"attachments": {},
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Refresh graph schema information\n",
|
||||||
|
"\n",
|
||||||
|
"If the schema of database changes, you can refresh the schema information needed to generate Cypher statements."
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 7,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# graph.refresh_schema()"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 8,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"Node properties: [{'properties': [('name', 'STRING')], 'label': 'Movie'}, {'properties': [('name', 'STRING'), ('birthDate', 'STRING')], 'label': 'Person'}]\n",
|
||||||
|
"Relationships properties: [{'properties': [], 'label': 'ActedIn'}]\n",
|
||||||
|
"Relationships: ['(:Person)-[:ActedIn]->(:Movie)']\n",
|
||||||
|
"\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"print(graph.get_schema)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"attachments": {},
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"## Querying the graph\n",
|
||||||
|
"\n",
|
||||||
|
"We can now use the `KuzuQAChain` to ask question of the graph"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 9,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"\u001b[1m> Entering new chain...\u001b[0m\n",
|
||||||
|
"Generated Cypher:\n",
|
||||||
|
"\u001b[32;1m\u001b[1;3mMATCH (p:Person)-[:ActedIn]->(m:Movie {name: 'The Godfather: Part II'}) RETURN p.name\u001b[0m\n",
|
||||||
|
"Full Context:\n",
|
||||||
|
"\u001b[32;1m\u001b[1;3m[{'p.name': 'Al Pacino'}, {'p.name': 'Robert De Niro'}]\u001b[0m\n",
|
||||||
|
"\n",
|
||||||
|
"\u001b[1m> Finished chain.\u001b[0m\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"'Al Pacino and Robert De Niro both played in The Godfather: Part II.'"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 9,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"chain.run(\"Who played in The Godfather: Part II?\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 10,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"\u001b[1m> Entering new chain...\u001b[0m\n",
|
||||||
|
"Generated Cypher:\n",
|
||||||
|
"\u001b[32;1m\u001b[1;3mMATCH (p:Person {name: 'Robert De Niro'})-[:ActedIn]->(m:Movie)\n",
|
||||||
|
"RETURN m.name\u001b[0m\n",
|
||||||
|
"Full Context:\n",
|
||||||
|
"\u001b[32;1m\u001b[1;3m[{'m.name': 'The Godfather: Part II'}]\u001b[0m\n",
|
||||||
|
"\n",
|
||||||
|
"\u001b[1m> Finished chain.\u001b[0m\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"'Robert De Niro played in The Godfather: Part II.'"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 10,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"chain.run(\"Robert De Niro played in which movies?\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 11,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"\u001b[1m> Entering new chain...\u001b[0m\n",
|
||||||
|
"Generated Cypher:\n",
|
||||||
|
"\u001b[32;1m\u001b[1;3mMATCH (p:Person {name: 'Robert De Niro'})-[:ActedIn]->(m:Movie)\n",
|
||||||
|
"RETURN p.birthDate\u001b[0m\n",
|
||||||
|
"Full Context:\n",
|
||||||
|
"\u001b[32;1m\u001b[1;3m[{'p.birthDate': '1943-08-17'}]\u001b[0m\n",
|
||||||
|
"\n",
|
||||||
|
"\u001b[1m> Finished chain.\u001b[0m\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"'Robert De Niro was born on August 17, 1943.'"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 11,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"chain.run(\"Robert De Niro is born in which year?\")"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 12,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"\u001b[1m> Entering new chain...\u001b[0m\n",
|
||||||
|
"Generated Cypher:\n",
|
||||||
|
"\u001b[32;1m\u001b[1;3mMATCH (p:Person)-[:ActedIn]->(m:Movie{name:'The Godfather: Part II'})\n",
|
||||||
|
"WITH p, m, p.birthDate AS birthDate\n",
|
||||||
|
"ORDER BY birthDate ASC\n",
|
||||||
|
"LIMIT 1\n",
|
||||||
|
"RETURN p.name\u001b[0m\n",
|
||||||
|
"Full Context:\n",
|
||||||
|
"\u001b[32;1m\u001b[1;3m[{'p.name': 'Al Pacino'}]\u001b[0m\n",
|
||||||
|
"\n",
|
||||||
|
"\u001b[1m> Finished chain.\u001b[0m\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"'The oldest actor who played in The Godfather: Part II is Al Pacino.'"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 12,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"chain.run(\"Who is the oldest actor who played in The Godfather: Part II?\")"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": "Python 3",
|
||||||
|
"language": "python",
|
||||||
|
"name": "python3"
|
||||||
|
},
|
||||||
|
"language_info": {
|
||||||
|
"codemirror_mode": {
|
||||||
|
"name": "ipython",
|
||||||
|
"version": 3
|
||||||
|
},
|
||||||
|
"file_extension": ".py",
|
||||||
|
"mimetype": "text/x-python",
|
||||||
|
"name": "python",
|
||||||
|
"nbconvert_exporter": "python",
|
||||||
|
"pygments_lexer": "ipython3",
|
||||||
|
"version": "3.11.4"
|
||||||
|
},
|
||||||
|
"orig_nbformat": 4
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 2
|
||||||
|
}
|
@ -0,0 +1,93 @@
|
|||||||
|
"""Question answering over a graph."""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
|
from pydantic import Field
|
||||||
|
|
||||||
|
from langchain.base_language import BaseLanguageModel
|
||||||
|
from langchain.callbacks.manager import CallbackManagerForChainRun
|
||||||
|
from langchain.chains.base import Chain
|
||||||
|
from langchain.chains.graph_qa.prompts import CYPHER_QA_PROMPT, KUZU_GENERATION_PROMPT
|
||||||
|
from langchain.chains.llm import LLMChain
|
||||||
|
from langchain.graphs.kuzu_graph import KuzuGraph
|
||||||
|
from langchain.prompts.base import BasePromptTemplate
|
||||||
|
|
||||||
|
|
||||||
|
class KuzuQAChain(Chain):
|
||||||
|
"""Chain for question-answering against a graph by generating Cypher statements for
|
||||||
|
Kùzu.
|
||||||
|
"""
|
||||||
|
|
||||||
|
graph: KuzuGraph = Field(exclude=True)
|
||||||
|
cypher_generation_chain: LLMChain
|
||||||
|
qa_chain: LLMChain
|
||||||
|
input_key: str = "query" #: :meta private:
|
||||||
|
output_key: str = "result" #: :meta private:
|
||||||
|
|
||||||
|
@property
|
||||||
|
def input_keys(self) -> List[str]:
|
||||||
|
"""Return the input keys.
|
||||||
|
|
||||||
|
:meta private:
|
||||||
|
"""
|
||||||
|
return [self.input_key]
|
||||||
|
|
||||||
|
@property
|
||||||
|
def output_keys(self) -> List[str]:
|
||||||
|
"""Return the output keys.
|
||||||
|
|
||||||
|
:meta private:
|
||||||
|
"""
|
||||||
|
_output_keys = [self.output_key]
|
||||||
|
return _output_keys
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_llm(
|
||||||
|
cls,
|
||||||
|
llm: BaseLanguageModel,
|
||||||
|
*,
|
||||||
|
qa_prompt: BasePromptTemplate = CYPHER_QA_PROMPT,
|
||||||
|
cypher_prompt: BasePromptTemplate = KUZU_GENERATION_PROMPT,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> KuzuQAChain:
|
||||||
|
"""Initialize from LLM."""
|
||||||
|
qa_chain = LLMChain(llm=llm, prompt=qa_prompt)
|
||||||
|
cypher_generation_chain = LLMChain(llm=llm, prompt=cypher_prompt)
|
||||||
|
|
||||||
|
return cls(
|
||||||
|
qa_chain=qa_chain,
|
||||||
|
cypher_generation_chain=cypher_generation_chain,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _call(
|
||||||
|
self,
|
||||||
|
inputs: Dict[str, Any],
|
||||||
|
run_manager: Optional[CallbackManagerForChainRun] = None,
|
||||||
|
) -> Dict[str, str]:
|
||||||
|
"""Generate Cypher statement, use it to look up in db and answer question."""
|
||||||
|
_run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
|
||||||
|
callbacks = _run_manager.get_child()
|
||||||
|
question = inputs[self.input_key]
|
||||||
|
|
||||||
|
generated_cypher = self.cypher_generation_chain.run(
|
||||||
|
{"question": question, "schema": self.graph.get_schema}, callbacks=callbacks
|
||||||
|
)
|
||||||
|
|
||||||
|
_run_manager.on_text("Generated Cypher:", end="\n", verbose=self.verbose)
|
||||||
|
_run_manager.on_text(
|
||||||
|
generated_cypher, color="green", end="\n", verbose=self.verbose
|
||||||
|
)
|
||||||
|
context = self.graph.query(generated_cypher)
|
||||||
|
|
||||||
|
_run_manager.on_text("Full Context:", end="\n", verbose=self.verbose)
|
||||||
|
_run_manager.on_text(
|
||||||
|
str(context), color="green", end="\n", verbose=self.verbose
|
||||||
|
)
|
||||||
|
|
||||||
|
result = self.qa_chain(
|
||||||
|
{"question": question, "context": context},
|
||||||
|
callbacks=callbacks,
|
||||||
|
)
|
||||||
|
return {self.output_key: result[self.qa_chain.output_key]}
|
@ -1,6 +1,7 @@
|
|||||||
"""Graph implementations."""
|
"""Graph implementations."""
|
||||||
|
from langchain.graphs.kuzu_graph import KuzuGraph
|
||||||
from langchain.graphs.nebula_graph import NebulaGraph
|
from langchain.graphs.nebula_graph import NebulaGraph
|
||||||
from langchain.graphs.neo4j_graph import Neo4jGraph
|
from langchain.graphs.neo4j_graph import Neo4jGraph
|
||||||
from langchain.graphs.networkx_graph import NetworkxEntityGraph
|
from langchain.graphs.networkx_graph import NetworkxEntityGraph
|
||||||
|
|
||||||
__all__ = ["NetworkxEntityGraph", "Neo4jGraph", "NebulaGraph"]
|
__all__ = ["NetworkxEntityGraph", "Neo4jGraph", "NebulaGraph", "KuzuGraph"]
|
||||||
|
@ -0,0 +1,90 @@
|
|||||||
|
from typing import Any, Dict, List
|
||||||
|
|
||||||
|
|
||||||
|
class KuzuGraph:
|
||||||
|
"""Kùzu wrapper for graph operations."""
|
||||||
|
|
||||||
|
def __init__(self, db: Any, database: str = "kuzu") -> None:
|
||||||
|
try:
|
||||||
|
import kuzu
|
||||||
|
except ImportError:
|
||||||
|
raise ImportError(
|
||||||
|
"Could not import Kùzu python package."
|
||||||
|
"Please install Kùzu with `pip install kuzu`."
|
||||||
|
)
|
||||||
|
self.db = db
|
||||||
|
self.conn = kuzu.Connection(self.db)
|
||||||
|
self.database = database
|
||||||
|
self.refresh_schema()
|
||||||
|
|
||||||
|
@property
|
||||||
|
def get_schema(self) -> str:
|
||||||
|
"""Returns the schema of the Kùzu database"""
|
||||||
|
return self.schema
|
||||||
|
|
||||||
|
def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:
|
||||||
|
"""Query Kùzu database"""
|
||||||
|
params_list = []
|
||||||
|
for param_name in params:
|
||||||
|
params_list.append([param_name, params[param_name]])
|
||||||
|
result = self.conn.execute(query, params_list)
|
||||||
|
column_names = result.get_column_names()
|
||||||
|
return_list = []
|
||||||
|
while result.has_next():
|
||||||
|
row = result.get_next()
|
||||||
|
return_list.append(dict(zip(column_names, row)))
|
||||||
|
return return_list
|
||||||
|
|
||||||
|
def refresh_schema(self) -> None:
|
||||||
|
"""Refreshes the Kùzu graph schema information"""
|
||||||
|
node_properties = []
|
||||||
|
node_table_names = self.conn._get_node_table_names()
|
||||||
|
for table_name in node_table_names:
|
||||||
|
current_table_schema = {"properties": [], "label": table_name}
|
||||||
|
properties = self.conn._get_node_property_names(table_name)
|
||||||
|
for property_name in properties:
|
||||||
|
property_type = properties[property_name]["type"]
|
||||||
|
list_type_flag = ""
|
||||||
|
if properties[property_name]["dimension"] > 0:
|
||||||
|
if "shape" in properties[property_name]:
|
||||||
|
for s in properties[property_name]["shape"]:
|
||||||
|
list_type_flag += "[%s]" % s
|
||||||
|
else:
|
||||||
|
for i in range(properties[property_name]["dimension"]):
|
||||||
|
list_type_flag += "[]"
|
||||||
|
property_type += list_type_flag
|
||||||
|
current_table_schema["properties"].append(
|
||||||
|
(property_name, property_type)
|
||||||
|
)
|
||||||
|
node_properties.append(current_table_schema)
|
||||||
|
|
||||||
|
relationships = []
|
||||||
|
rel_tables = self.conn._get_rel_table_names()
|
||||||
|
for table in rel_tables:
|
||||||
|
relationships.append(
|
||||||
|
"(:%s)-[:%s]->(:%s)" % (table["src"], table["name"], table["dst"])
|
||||||
|
)
|
||||||
|
|
||||||
|
rel_properties = []
|
||||||
|
for table in rel_tables:
|
||||||
|
current_table_schema = {"properties": [], "label": table["name"]}
|
||||||
|
properties_text = self.conn._connection.get_rel_property_names(
|
||||||
|
table["name"]
|
||||||
|
).split("\n")
|
||||||
|
for i, line in enumerate(properties_text):
|
||||||
|
# The first 3 lines defines src, dst and name, so we skip them
|
||||||
|
if i < 3:
|
||||||
|
continue
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
property_name, property_type = line.strip().split(" ")
|
||||||
|
current_table_schema["properties"].append(
|
||||||
|
(property_name, property_type)
|
||||||
|
)
|
||||||
|
rel_properties.append(current_table_schema)
|
||||||
|
|
||||||
|
self.schema = (
|
||||||
|
f"Node properties: {node_properties}\n"
|
||||||
|
f"Relationships properties: {rel_properties}\n"
|
||||||
|
f"Relationships: {relationships}\n"
|
||||||
|
)
|
@ -0,0 +1,56 @@
|
|||||||
|
import shutil
|
||||||
|
import tempfile
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
from langchain.graphs import KuzuGraph
|
||||||
|
|
||||||
|
EXPECTED_SCHEMA = """
|
||||||
|
Node properties: [{'properties': [('name', 'STRING')], 'label': 'Movie'}, {'properties': [('name', 'STRING'), ('birthDate', 'STRING')], 'label': 'Person'}]
|
||||||
|
Relationships properties: [{'properties': [], 'label': 'ActedIn'}]
|
||||||
|
Relationships: ['(:Person)-[:ActedIn]->(:Movie)']
|
||||||
|
""" # noqa: E501
|
||||||
|
|
||||||
|
|
||||||
|
class TestKuzu(unittest.TestCase):
|
||||||
|
def setUp(self) -> None:
|
||||||
|
try:
|
||||||
|
import kuzu
|
||||||
|
except ImportError as e:
|
||||||
|
raise ImportError(
|
||||||
|
"Cannot import Python package kuzu. Please install it by running "
|
||||||
|
"`pip install kuzu`."
|
||||||
|
) from e
|
||||||
|
|
||||||
|
self.tmpdir = tempfile.mkdtemp()
|
||||||
|
self.kuzu_database = kuzu.Database(self.tmpdir)
|
||||||
|
self.conn = kuzu.Connection(self.kuzu_database)
|
||||||
|
self.conn.execute("CREATE NODE TABLE Movie (name STRING, PRIMARY KEY(name))")
|
||||||
|
self.conn.execute("CREATE (:Movie {name: 'The Godfather'})")
|
||||||
|
self.conn.execute("CREATE (:Movie {name: 'The Godfather: Part II'})")
|
||||||
|
self.conn.execute(
|
||||||
|
"CREATE (:Movie {name: 'The Godfather Coda: The Death of Michael "
|
||||||
|
"Corleone'})"
|
||||||
|
)
|
||||||
|
self.kuzu_graph = KuzuGraph(self.kuzu_database)
|
||||||
|
|
||||||
|
def tearDown(self) -> None:
|
||||||
|
shutil.rmtree(self.tmpdir, ignore_errors=True)
|
||||||
|
|
||||||
|
def test_query(self) -> None:
|
||||||
|
result = self.kuzu_graph.query("MATCH (n:Movie) RETURN n.name ORDER BY n.name")
|
||||||
|
excepted_result = [
|
||||||
|
{"n.name": "The Godfather"},
|
||||||
|
{"n.name": "The Godfather Coda: The Death of Michael Corleone"},
|
||||||
|
{"n.name": "The Godfather: Part II"},
|
||||||
|
]
|
||||||
|
self.assertEqual(result, excepted_result)
|
||||||
|
|
||||||
|
def test_refresh_schema(self) -> None:
|
||||||
|
self.conn.execute(
|
||||||
|
"CREATE NODE TABLE Person (name STRING, birthDate STRING, PRIMARY "
|
||||||
|
"KEY(name))"
|
||||||
|
)
|
||||||
|
self.conn.execute("CREATE REL TABLE ActedIn (FROM Person TO Movie)")
|
||||||
|
self.kuzu_graph.refresh_schema()
|
||||||
|
schema = self.kuzu_graph.get_schema
|
||||||
|
self.assertEqual(schema, EXPECTED_SCHEMA)
|
Loading…
Reference in New Issue