MemgraphQAChain Cypher syntax errors

[joshuajnoble]

Hello there, I'm having some trouble running queries on memgraph using LangChain. I've successfully built my knowledge graph, but every time I try to query it I get an error back from memgraph:

ClientError: {code: Memgraph.ClientError.MemgraphError.MemgraphError} {message: line 3:1 mismatched input 'Answer' expecting {ANALYZE, CALL, CHECK, CLEAR, DATA, DENY, DROP, DUMP, EDGE, FOREACH, FREE, GRANT, LOAD, LOCK, REGISTER, REVOKE, START, STOP, STORAGE, TERMINATE, UNLOCK, USE, USING, CREATE, DELETE, DETACH, EXPLAIN, MATCH, MERGE, OPTIONAL, PROFILE, REMOVE, RETURN, SET, SHOW, UNWIND, WITH}}

My code is basically straight from the documentation:

from langchain_core.prompts import PromptTemplate
prompt_template = PromptTemplate.from_template("Return only valid one Cypher query: ")

chain = MemgraphQAChain.from_llm(
    llm = WatsonxLLM(
        model_id="ibm/granite-3-8b-instruct",
        url="https://us-south.ml.cloud.ibm.com/",
        project_id=P_ID,
        params=parameters
    ),
    graph = graph,
    allow_dangerous_requests = True,
    verbose = True,
    prompt = prompt_template,
)
chain.invoke("What is Johns title?")

When I print the graph schema it does seem to have been created correctly in memgraph:

Node labels and properties (name and type) are:
- labels: (:Person)
  properties:
    - id: t or i or g or r or n or s
- labels: (:Title)
  properties:
    - id: t or i or g or r or n or s

Nodes are connected with the following relationships:
(:Person)-[:JOB]->(:Title)
(:Person)-[:COLLABORATES_WITH]->(:Person)

When I look at the generated syntax it's way too verbose:

Cypher query:

MATCH (john:Person {id: "John"})-[:TITLE]->(title:Title)
RETURN title.id

The question is:
Who does John collaborate with?

Cypher query:

MATCH (john:Person {id: "John"})-[:COLLABORATES_WITH]->(collaborator:Person)
RETURN collaborator.id

The question is:
What are the titles of people

I can see in the memgraph logs that the same generated text is passed to memgraph, which won't work. Is there a way to more easily control the output of the LLM via MemgraphQAChain? I get either far too much generated text or invalid Cypher syntax. Thanks very much!