Persisting Responses

/**
 * Save a question and response to the database
 *
 * @param {string} sessionId
 * @param {string} source
 * @param {string} input
 * @param {string} rephrasedQuestion
 * @param {string} output
 * @param {string[]} ids
 * @param {string | null} cypher
 * @returns {string}  The ID of the Message node
 */
export async function saveHistory(
  sessionId: string,
  source: string,
  input: string,
  rephrasedQuestion: string,
  output: string,
  ids: string[],
  cypher: string | null = null
): Promise<string> {
  const graph = await initGraph();
  const res = await graph.query<{ id: string }>(
    `
    MERGE (session:Session { id: $sessionId }) // (1)

    // <2> Create new response
    CREATE (response:Response {
      id: randomUuid(),
      createdAt: datetime(),
      source: $source,
      input: $input,
      output: $output,
      rephrasedQuestion: $rephrasedQuestion,
      cypher: $cypher,
      ids: $ids
    })
    CREATE (session)-[:HAS_RESPONSE]->(response)

    WITH session, response

    CALL {
    WITH session, response

      // <3> Remove existing :LAST_RESPONSE relationship if it exists
      MATCH (session)-[lrel:LAST_RESPONSE]->(last)
      DELETE lrel

      // <4? Create :NEXT relationship
      CREATE (last)-[:NEXT]->(response)
    }


    // <5> Create new :LAST_RESPONSE relationship
    CREATE (session)-[:LAST_RESPONSE]->(response)

    // <6> Create relationship to context nodes
    WITH response

    CALL {
      WITH response
      UNWIND $ids AS id
      MATCH (context)
      WHERE elementId(context) = id
      CREATE (response)-[:CONTEXT]->(context)

      RETURN count(*) AS count
    }

    RETURN DISTINCT response.id AS id
  `,
    {
      sessionId,
      source,
      input,
      output,
      rephrasedQuestion,
      cypher: cypher,
      ids,
    },
    "WRITE"
  );

  return res && res.length ? res[0].id : "";
}

export async function getHistory(
  sessionId: string,
  limit: number = 5
): Promise<ChatbotResponse[]> {
  const graph = await initGraph();
  const res = await graph.query<ChatbotResponse>(
    `
      MATCH (:Session {id: $sessionId})-[:LAST_RESPONSE]->(last)
      MATCH path = (start)-[:NEXT*0..${limit}]->(last)
      WHERE length(path) = 5 OR NOT EXISTS { ()-[:NEXT]->(start) }
      UNWIND nodes(path) AS response
      RETURN response.id AS id,
        response.input AS input,
        response.rephrasedQuestion AS rephrasedQuestion,
        response.output AS output,
        response.cypher AS cypher,
        response.createdAt AS createdAt,
        [ (response)-[:CONTEXT]->(n) | elementId(n) ] AS context
    `,
    { sessionId },
    "READ"
  );

  return res as ChatbotResponse[];
}

import { close } from "../graph";
import { getHistory, saveHistory } from "./history";
import { Neo4jGraph } from "@langchain/community/graphs/neo4j_graph";

describe("Conversation History", () => {
  let graph: Neo4jGraph;
  let ids: string[];

  beforeAll(async () => {
    graph = await Neo4jGraph.initialize({
      url: process.env.NEO4J_URI as string,
      username: process.env.NEO4J_USERNAME as string,
      password: process.env.NEO4J_PASSWORD as string,
    });

    // Delete responses
    await graph.query(
      'MATCH (s:Session)-[:HAS_RESPONSE]->(r:Response) WHERE s.id IN ["test-1", "test-2", "test-3"] DETACH DELETE s, r'
    );

    // Create some test sources to link to
    const [first] = (await graph.query(`
            MERGE (t1:TestSource {id: 1})
            MERGE (t2:TestSource {id: 2})
            MERGE (t3:TestSource {id: 3})

            RETURN [ elementId(t1), elementId(t2), elementId(t3) ] AS ids
        `)) as Record<string, any>[];

    ids = first.ids;
  });

  afterAll(async () => {
    await graph.close();
    await close();
  });

  it("should save conversation history", async () => {
    const sessionId = "test-1";
    const source = "cypher";
    const input = "Who directed The Matrix?";
    const rephrasedQuestion = "Director of The Matrix";
    const output = "The Matrix was directed by The Wachowskis";
    const cypher =
      'MATCH (p:Person)-[:DIRECTED]->(m:Movie {title: "The Matrix"}) RETURN p.name AS name';

    // Save message
    const id = await saveHistory(
      sessionId,
      source,
      input,
      rephrasedQuestion,
      output,
      ids,
      cypher
    );

    expect(id).toBeDefined();

    // Get History
    const history = await getHistory(sessionId, 5);

    expect(history?.length).toBeGreaterThanOrEqual(1);

    const returnedIds = history.map((m) => m.id);

    // Was message returned in the history
    expect(returnedIds).toContain(id);

    // Check sources
    const res = await graph.query(
      `
        MATCH (s:Session {id: $sessionId})-[:LAST_RESPONSE]->(r)
        RETURN r { .* } AS properties,
        [ (r)-[:CONTEXT]->(c) | elementId(c) ] AS context
      `,
      { sessionId }
    );

    expect(res).toBeDefined();
    expect(res?.length).toBeGreaterThanOrEqual(1);

    // Has context been linked?
    const first = res![0];

    expect(first.properties.id).toEqual(id);
    expect(first.properties.source).toEqual(source);
    expect(first.properties.input).toEqual(input);
    expect(first.properties.rephrasedQuestion).toEqual(rephrasedQuestion);
    expect(first.properties.cypher).toEqual(cypher);
    expect(first.properties.output).toEqual(output);

    // Have all context nodes been linked to?
    for (const id of ids) {
      expect(first.context).toContain(id);
    }
  });

  it("should save a chain of responses", async () => {
    const sessionId = "test-3";
    const source = "retriever";
    const firstInput = "Who directed Toy Story?";
    const secondInput = "Who acted in it?";
    const thirdInput = "What else have the acted in together?";

    // Save message
    const messages = [
      await saveHistory(sessionId, source, firstInput, "", "", []),
      await saveHistory(sessionId, source, secondInput, "", "", []),
      await saveHistory(sessionId, source, thirdInput, "", "", []),
    ];

    for (const message of messages) {
      expect(message).toBeDefined();
    }

    // Get History
    const history = await getHistory(sessionId, 5);

    const returnedIds = history.map((m) => m.id).join(",");

    // Responses should be returned in order
    expect(messages.join(",")).toBe(returnedIds);
  });
});