Kuzu Agent - Issue Capturing Steps

[harryaultdev]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from typing import Any, Dict
from uuid import UUID
from langchain_core.agents import AgentAction
from langchain_core.callbacks import BaseCallbackHandler
from langchain_openai import AzureChatOpenAI
from langchain.chains import KuzuQAChain
from langchain_community.graphs import KuzuGraph

llm = AzureChatOpenAI(azure_deployment="gpt-4o")



chain = KuzuQAChain.from_llm(
    llm, graph=graph, verbose=True, return_intermediate_steps=True
)

result = chain.invoke({"query": "Who are my Employees?"})

print(f"Intermediate steps: {result['intermediate_steps']}")
print(f"Final answer: {result['result']}")

Description

Expecting to see the intermediate steps but they're not being returned? Alternatively is there a callback handler I can use? Or am I just being an idiot...

I basically just want to capture the Cypher generated which I can see in the logs but I need to store it in a variable to return in an API response.

Verbose Logging below.

Entering new KuzuQAChain chain...
Generated Cypher:
MATCH (e:Employee) RETURN e
Full Context:

System Info

System Information

OS: Windows
OS Version: 10.0.22631
Python Version: 3.11.9 (tags/v3.11.9:de54cf5, Apr 2 2024, 10:12:12) [MSC v.1938 64 bit (AMD64)]

Package Information

langchain_core: 0.2.12
langchain: 0.2.7
langchain_community: 0.2.7
langsmith: 0.1.84
langchain_experimental: 0.0.62
langchain_openai: 0.1.14
langchain_text_splitters: 0.2.2

Packages not installed (Not Necessarily a Problem)

The following packages were not found:

langgraph
langserve

[dosubot[bot]]

Hey there, @harryaultdev! I'm Dosu, a bot here to help you solve bugs, answer questions, and become a contributor. I'm here to assist you while waiting for a human maintainer.

To capture the Cypher generated by KuzuQAChain and store it in a variable for an API response, you can use the CallbackManagerForChainRun. Here's how you can modify your code to achieve this:

1. Create a custom callback handler to capture the generated Cypher.
2. Modify the _call method in KuzuQAChain to store the generated Cypher in a variable.

Step 1: Create a Custom Callback Handler

from langchain_core.callbacks import BaseCallbackHandler

class CustomCallbackHandler(BaseCallbackHandler):
    def __init__(self):
        self.generated_cypher = None

    def on_text(self, text: str, **kwargs: Any) -> None:
        if "Generated Cypher:" in text:
            self.generated_cypher = text.split("Generated Cypher:")[1].strip()

Step 2: Modify the _call Method in KuzuQAChain

from langchain_core.callbacks.manager import CallbackManagerForChainRun

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
    )
    # Extract Cypher code if it is wrapped in triple backticks
    # with the language marker "cypher"
    generated_cypher = remove_prefix(extract_cypher(generated_cypher), "cypher")

    # Store the generated Cypher in a variable
    self.generated_cypher = generated_cypher

    _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]}

Updated Example Code

from typing import Any, Dict
from langchain_core.callbacks.manager import CallbackManagerForChainRun
from langchain_openai import AzureChatOpenAI
from langchain.chains import KuzuQAChain
from langchain_community.graphs import KuzuGraph

# Initialize the custom callback handler
callback_handler = CustomCallbackHandler()

# Initialize the callback manager with the custom handler
callback_manager = CallbackManagerForChainRun(handlers=[callback_handler])

llm = AzureChatOpenAI(azure_deployment="gpt-4o")
graph = KuzuGraph()  # Initialize your graph here

chain = KuzuQAChain.from_llm(
    llm, graph=graph, verbose=True, return_intermediate_steps=True
)

# Pass the callback manager to the chain
result = chain.invoke({"query": "Who are my Employees?"}, run_manager=callback_manager)

print(f"Intermediate steps: {result['intermediate_steps']}")
print(f"Final answer: {result['result']}")
print(f"Generated Cypher: {callback_handler.generated_cypher}")

This setup will capture the generated Cypher statement and store it in the generated_cypher attribute of the CustomCallbackHandler instance, which you can then use in your API response ^[1].

To continue talking to Dosu, mention @dosu.

[NatashaKSS]

This is probably a hacky workaround ... but if you inspect the _call method in langchain_kuzu/chains/graph_qa/kuzu.py, you can see:

...
"""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]

self.graph.refresh_schema()
generated_cypher = self.cypher_generation_chain.run(
    {"question": question, "schema": self.graph.get_schema}, callbacks=callbacks
)
# Extract Cypher code if it is wrapped in triple backticks
# with the language marker "cypher"
generated_cypher = remove_prefix(extract_cypher(generated_cypher), "cypher")

_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]}

You can kinda get the generated cypher by doing:

# note: the run method was deprecated in langchain 0.1.0, just illustrating with an example
generated_cypher = self.cypher_generation_chain.run(
     {"question": question, "schema": self.graph.get_schema}, callbacks=callbacks
)

So if you have something like:

# Create the KuzuQAChain with verbosity enabled to see the generated Cypher queries
chain = KuzuQAChain.from_llm(
    llm=ChatOpenAI(model="gpt-4o-mini", temperature=0.3, api_key=os.getenv("OPENAI_API_KEY")),
    graph=graph,
    verbose=True,
    allow_dangerous_requests=True,
)

# with the user's question and the KuzuGraph instance, then you can do
generated_cypher_result = chain.cypher_generation_chain.invoke(
     {"question": question, "schema": graph.get_schema}, callbacks=callbacks
)

# Clean up the cypher result by removing markdown formatting
# note this might not handle multiline queries well
cypher_query = generated_cypher_result['text'].strip("```").strip("cypher").strip()

print(cypher_query)

I would prefer if the API just provided an optional result with the Cypher query though