"""Question answering over a graph."""from__future__importannotationsimportrefromtypingimportAny,Dict,List,Optionalfromlangchain.chains.baseimportChainfromlangchain.chains.llmimportLLMChainfromlangchain_core.callbacksimportCallbackManagerForChainRunfromlangchain_core.language_modelsimportBaseLanguageModelfromlangchain_core.promptsimportBasePromptTemplatefrompydanticimportFieldfromlangchain_community.chains.graph_qa.promptsimport(AQL_FIX_PROMPT,AQL_GENERATION_PROMPT,AQL_QA_PROMPT,)fromlangchain_community.graphs.arangodb_graphimportArangoGraph
[docs]classArangoGraphQAChain(Chain):"""Chain for question-answering against a graph by generating AQL statements. *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. """graph:ArangoGraph=Field(exclude=True)aql_generation_chain:LLMChainaql_fix_chain:LLMChainqa_chain:LLMChaininput_key:str="query"#: :meta private:output_key:str="result"#: :meta private:# Specifies the maximum number of AQL Query Results to returntop_k:int=10# Specifies the set of AQL Query Examples that promote few-shot-learningaql_examples:str=""# Specify whether to return the AQL Query in the output dictionaryreturn_aql_query:bool=False# Specify whether to return the AQL JSON Result in the output dictionaryreturn_aql_result:bool=False# Specify the maximum amount of AQL Generation attempts that should be mademax_aql_generation_attempts:int=3allow_dangerous_requests:bool=False"""Forced user opt-in to acknowledge that the chain can make dangerous requests. *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,**kwargs:Any)->None:"""Initialize the chain."""super().__init__(**kwargs)ifself.allow_dangerous_requestsisnotTrue:raiseValueError("In order to use this chain, you must acknowledge that it can make ""dangerous requests by setting `allow_dangerous_requests` to `True`.""You must narrowly scope the permissions of the database connection ""to only include necessary permissions. Failure to do so may result ""in data corruption or loss or reading sensitive data if such data is ""present in the database.""Only use this chain if you understand the risks and have taken the ""necessary precautions. ""See https://python.langchain.com/docs/security for more information.")@propertydefinput_keys(self)->List[str]:return[self.input_key]@propertydefoutput_keys(self)->List[str]:return[self.output_key]@propertydef_chain_type(self)->str:return"graph_aql_chain"
[docs]@classmethoddeffrom_llm(cls,llm:BaseLanguageModel,*,qa_prompt:BasePromptTemplate=AQL_QA_PROMPT,aql_generation_prompt:BasePromptTemplate=AQL_GENERATION_PROMPT,aql_fix_prompt:BasePromptTemplate=AQL_FIX_PROMPT,**kwargs:Any,)->ArangoGraphQAChain:"""Initialize from LLM."""qa_chain=LLMChain(llm=llm,prompt=qa_prompt)aql_generation_chain=LLMChain(llm=llm,prompt=aql_generation_prompt)aql_fix_chain=LLMChain(llm=llm,prompt=aql_fix_prompt)returncls(qa_chain=qa_chain,aql_generation_chain=aql_generation_chain,aql_fix_chain=aql_fix_chain,**kwargs,)
def_call(self,inputs:Dict[str,Any],run_manager:Optional[CallbackManagerForChainRun]=None,)->Dict[str,Any]:""" Generate an AQL statement from user input, use it retrieve a response from an ArangoDB Database instance, and respond to the user input in natural language. Users can modify the following ArangoGraphQAChain Class Variables: :var top_k: The maximum number of AQL Query Results to return :type top_k: int :var aql_examples: A set of AQL Query Examples that are passed to the AQL Generation Prompt Template to promote few-shot-learning. Defaults to an empty string. :type aql_examples: str :var return_aql_query: Whether to return the AQL Query in the output dictionary. Defaults to False. :type return_aql_query: bool :var return_aql_result: Whether to return the AQL Query in the output dictionary. Defaults to False :type return_aql_result: bool :var max_aql_generation_attempts: The maximum amount of AQL Generation attempts to be made prior to raising the last AQL Query Execution Error. Defaults to 3. :type max_aql_generation_attempts: int """_run_manager=run_managerorCallbackManagerForChainRun.get_noop_manager()callbacks=_run_manager.get_child()user_input=inputs[self.input_key]########################## Generate AQL Query #aql_generation_output=self.aql_generation_chain.run({"adb_schema":self.graph.schema,"aql_examples":self.aql_examples,"user_input":user_input,},callbacks=callbacks,)#########################aql_query=""aql_error=""aql_result=Noneaql_generation_attempt=1while(aql_resultisNoneandaql_generation_attempt<self.max_aql_generation_attempts+1):###################### Extract AQL Query #pattern=r"```(?i:aql)?(.*?)```"matches=re.findall(pattern,aql_generation_output,re.DOTALL)ifnotmatches:_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)raiseValueError(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)###################### Execute AQL Query #fromarangoimportAQLQueryExecuteErrortry:aql_result=self.graph.query(aql_query,self.top_k)exceptAQLQueryExecuteErrorase: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)######################### Retry AQL Generation #aql_generation_output=self.aql_fix_chain.run({"adb_schema":self.graph.schema,"aql_query":aql_query,"aql_error":aql_error,},callbacks=callbacks,)#############################################aql_generation_attempt+=1ifaql_resultisNone:m=f""" Maximum amount of AQL Query Generation attempts reached. Unable to execute the AQL Query due to the following error:{aql_error} """raiseValueError(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)######################### Interpret AQL Result #result=self.qa_chain({"adb_schema":self.graph.schema,"user_input":user_input,"aql_query":aql_query,"aql_result":aql_result,},callbacks=callbacks,)######################### Return results #result={self.output_key:result[self.qa_chain.output_key]}ifself.return_aql_query:result["aql_query"]=aql_queryifself.return_aql_result:result["aql_result"]=aql_resultreturnresult