Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

wip: ArangoDB #1

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 30 additions & 31 deletions libs/community/langchain_community/chains/graph_qa/arangodb.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,10 @@ class ArangoGraphQAChain(Chain):
# Specify the maximum amount of AQL Generation attempts that should be made
max_aql_generation_attempts: int = 3

# Specify whether to execute the generated AQL Query
# If False, the AQL Query is only explained & returned, not executed
execute_aql_query: bool = True

allow_dangerous_requests: bool = False
"""Forced user opt-in to acknowledge that the chain can make dangerous requests.

Expand Down Expand Up @@ -155,6 +159,11 @@ def _call(
AQL Query Execution Error. Defaults to 3.
:type max_aql_generation_attempts: int
"""
try:
from arango import AQLQueryExecuteError, AQLQueryExplainError
except ImportError:
raise ImportError("ArangoDB not installed, please install with `pip install python-arango`.")

_run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
callbacks = _run_manager.get_child()
user_input = inputs[self.input_key]
Expand All @@ -176,48 +185,34 @@ def _call(
aql_result = None
aql_generation_attempt = 1

while (
aql_result is None
and aql_generation_attempt < self.max_aql_generation_attempts + 1
):
aql_execution_func = self.graph.query if self.execute_aql_query else self.graph.explain

while aql_result is None and aql_generation_attempt < self.max_aql_generation_attempts + 1:
#####################
# Extract AQL Query #
pattern = r"```(?i:aql)?(.*?)```"
matches = re.findall(pattern, aql_generation_output, re.DOTALL)
if not matches:
_run_manager.on_text(
"Invalid Response: ", end="\n", verbose=self.verbose
)
_run_manager.on_text(
aql_generation_output, color="red", end="\n", verbose=self.verbose
)
_run_manager.on_text("Invalid Response: ", end="\n", verbose=self.verbose)
_run_manager.on_text(aql_generation_output, color="red", end="\n", verbose=self.verbose)
raise ValueError(f"Response is Invalid: {aql_generation_output}")

aql_query = matches[0]
#####################

_run_manager.on_text(
f"AQL Query ({aql_generation_attempt}):", verbose=self.verbose
)
_run_manager.on_text(
aql_query, color="green", end="\n", verbose=self.verbose
)
_run_manager.on_text(f"AQL Query ({aql_generation_attempt}):", verbose=self.verbose)
_run_manager.on_text(aql_query, color="green", end="\n", verbose=self.verbose)

#####################
# Execute AQL Query #
from arango import AQLQueryExecuteError
#############################
# Execute/Explain AQL Query #

try:
aql_result = self.graph.query(aql_query, self.top_k)
except AQLQueryExecuteError as e:
aql_result = aql_execution_func(aql_query, self.top_k)
except (AQLQueryExecuteError, AQLQueryExplainError) as e:
aql_error = e.error_message

_run_manager.on_text(
"AQL Query Execution Error: ", end="\n", verbose=self.verbose
)
_run_manager.on_text(
aql_error, color="yellow", end="\n\n", verbose=self.verbose
)
_run_manager.on_text("AQL Query Execution Error: ", end="\n", verbose=self.verbose)
_run_manager.on_text(aql_error, color="yellow", end="\n\n", verbose=self.verbose)

########################
# Retry AQL Generation #
Expand All @@ -243,10 +238,14 @@ def _call(
"""
raise ValueError(m)

_run_manager.on_text("AQL Result:", end="\n", verbose=self.verbose)
_run_manager.on_text(
str(aql_result), color="green", end="\n", verbose=self.verbose
)
text = "AQL Result:" if self.execute_aql_query else "AQL Explain:"
_run_manager.on_text(text, end="\n", verbose=self.verbose)
_run_manager.on_text(str(aql_result), color="green", end="\n", verbose=self.verbose)

if not self.execute_aql_query:
result = {self.output_key: aql_query}

return result

########################
# Interpret AQL Result #
Expand Down
Loading