Skip to content

neurotrace.core.graph_memory

neurotrace.core.graph_memory

GraphMemoryAdapter

Bases: BaseGraphMemoryAdapter

Source code in neurotrace/core/graph_memory.py 
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
class GraphMemoryAdapter(BaseGraphMemoryAdapter):
    def __init__(
        self,
        llm: Union[BaseLLM, BaseChatModel],
        graph_database: GraphStore,
        triplets_indexer: GraphTripletIndexerBase = None,
    ):
        self.llm = llm
        self.graph = graph_database

        self.triplets_indexer = triplets_indexer or GraphTripletIndexer(self.llm)
        self.qa_chain = GraphCypherQAChain.from_llm(
            llm=self.llm, graph=self.graph, verbose=True, allow_dangerous_requests=True  # Prints the generated Cypher
        )

    def insert_triplets(
        self, triples, sender: Literal["user", "agent"] = "user", timestamp: str = None, tags: List[str] = None
    ):
        if not triples:
            print("No triplets to insert")
            return

        if timestamp is None:
            timestamp = datetime.now().isoformat()

        if len(triples) > 3:
            print("Too many triplets", triples)
            return

        s, r, o = triples
        query = f"""
        MERGE (a:Entity {{name: $s}})
          ON CREATE SET a.created_at = $timestamp

        MERGE (b:Entity {{name: $o}})
          ON CREATE SET b.created_at = $timestamp

        MERGE (a)-[r:`{r.upper().replace(" ", "_")}`]->(b)
          ON CREATE SET r.created_at = $timestamp, r.source = $sender
          ON CREATE SET r.tags = $tags
        """
        self.graph.query(
            query, params={"s": s, "o": o, "r": r, "timestamp": timestamp, "sender": sender, "tags": tags or []}
        )

    def add_conversation(
        self, summarised_text: str, sender: Literal["user", "agent"] = "agent", tags: List[str] = None
    ):
        """
        Add a conversation to the graph memory.

        Args:
            sender (Literal["user", "agent"]): The sender of the message.
            :param summarised_text:
        """
        triplets = self.triplets_indexer.extract(summarised_text)

        for triplet in triplets:
            self.insert_triplets(triplet, sender=sender, tags=tags)

    def ask_graph(self, query: str) -> Dict[str, Any]:
        """
        Ask a question to the graph memory and get an answer.

        Args:
            query (str): The question to ask.

        Returns:
            str: The answer from the graph memory.
        """
        return self.qa_chain.invoke({"query": query})

    def get_all_relation_types(self) -> list[str]:
        result = self.graph.query("CALL db.relationshipTypes()")
        # Output is like: [{'relationshipType': 'WORKS_AT'}, ...]
        return [record["relationshipType"] for record in result]

add_conversation(summarised_text, sender='agent', tags=None)

Add a conversation to the graph memory.

Parameters:

Name Type Description Default
sender Literal['user', 'agent']

The sender of the message.

 'agent'

param summarised_text:

 required
Source code in neurotrace/core/graph_memory.py 
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def add_conversation(
    self, summarised_text: str, sender: Literal["user", "agent"] = "agent", tags: List[str] = None
):
    """
    Add a conversation to the graph memory.

    Args:
        sender (Literal["user", "agent"]): The sender of the message.
        :param summarised_text:
    """
    triplets = self.triplets_indexer.extract(summarised_text)

    for triplet in triplets:
        self.insert_triplets(triplet, sender=sender, tags=tags)

ask_graph(query)

Ask a question to the graph memory and get an answer.

Parameters:

Name Type Description Default
query str

The question to ask.

 required

Returns:

Name Type Description
str Dict[str, Any]

The answer from the graph memory.
Source code in neurotrace/core/graph_memory.py 
 98
 99
100
101
102
103
104
105
106
107
108
def ask_graph(self, query: str) -> Dict[str, Any]:
    """
    Ask a question to the graph memory and get an answer.

    Args:
        query (str): The question to ask.

    Returns:
        str: The answer from the graph memory.
    """
    return self.qa_chain.invoke({"query": query})