Reduce size of context, hack for now

cognitive_architecture/database/graph_database/graph.py

@@ -104,6 +104,8 @@ class Node(BaseModel):
     category: str
     color: str ="blue"
     memory_type: str
+    created_at: Optional[float] = None
+    summarized: Optional[bool] = None
 
 
 
@@ -112,7 +114,8 @@ class Edge(BaseModel):
     target: int
     description: str
     color: str= "blue"
-
+    created_at: Optional[float] = None
+    summarized: Optional[bool] = None
 
 class KnowledgeGraph(BaseModel):
     nodes: List[Node] = Field(..., default_factory=list)
@@ -170,6 +173,29 @@ def generate_graph(input) -> KnowledgeGraph:
     )
     return out
 
+class MemorySummary(BaseModel):
+    nodes: List[Node] = Field(..., default_factory=list)
+    edges: List[Edge] = Field(..., default_factory=list)
+
+async def generate_summary(input) -> MemorySummary:
+    out =  aclient.chat.completions.create(
+        model="gpt-4-1106-preview",
+        messages=[
+            {
+                "role": "user",
+                "content": f"""Use the given format summarize and reduce the following input: {input}. """,
+
+            },
+            {   "role":"system", "content": """You are a top-tier algorithm
+                designed for summarizing existing knowledge graphs in structured formats based on a knowledge graph.
+                ## 1. Strict Compliance
+                Adhere to the rules strictly. Non-compliance will result in termination.
+                ## 2. Don't forget your main goal is to reduce the number of nodes in the knowledge graph while preserving the information contained in it."""}
+        ],
+        response_model=MemorySummary,
+    )
+    return out
+
 class AbstractGraphDB(ABC):
 
     @abstractmethod
@@ -226,6 +252,55 @@ class Neo4jGraphDB(AbstractGraphDB):
         """
         return user_memory_cypher
 
+    async def retrieve_memory(self, user_id: str, memory_type: str, timestamp: float = None, summarized: bool = None):
+        if memory_type == 'SemanticMemory':
+            relationship = 'SEMANTIC_MEMORY'
+            memory_rel = 'HAS_KNOWLEDGE'
+        elif memory_type == 'EpisodicMemory':
+            relationship = 'EPISODIC_MEMORY'
+            memory_rel = 'HAS_EVENT'
+        elif memory_type == 'Buffer':
+            relationship = 'BUFFER'
+            memory_rel = 'CURRENTLY_HOLDING'
+        if timestamp is not None and summarized is not None:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_{relationship}]->(memory:{memory_type})
+            MATCH (memory)-[:{memory_rel}]->(item)
+            WHERE item.created_at >= {timestamp} AND item.summarized = {str(summarized).lower()}
+            RETURN item
+            """
+        elif timestamp is not None:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_{relationship}]->(memory:{memory_type})
+            MATCH (memory)-[:{memory_rel}]->(item)
+            WHERE item.created_at >= {timestamp}
+            RETURN item
+            """
+        elif summarized is not None:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_{relationship}]->(memory:{memory_type})
+            MATCH (memory)-[:{memory_rel}]->(item)
+            WHERE item.summarized = {str(summarized).lower()}
+            RETURN item
+            """
+            print(query)
+        else:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_{relationship}]->(memory:{memory_type})
+            MATCH (memory)-[:{memory_rel}]->(item)
+            RETURN item
+            """
+        output = self.query(query, params={"user_id": user_id})
+        print("Here is the output", output)
+
+        reduced_graph = await generate_summary(input = output)
+        return reduced_graph
+
+
+    # def add_summarized_memory(self, user_id: str, memory_type: str, timestamp: float = None, summarized: bool = None):
+    #
+
+
     def user_query_to_edges_and_nodes(self, input: str) ->KnowledgeGraph:
         return aclient.chat.completions.create(
             model=config.model,
@@ -342,6 +417,14 @@ class Neo4jGraphDB(AbstractGraphDB):
     async def generate_cypher_query_for_user_prompt_decomposition(self, user_id:str, query:str):
 
         graph: KnowledgeGraph = generate_graph(query)
+        import time
+        for node in graph.nodes:
+            node.created_at = time.time()
+            node.summarized = False
+
+        for edge in graph.edges:
+            edge.created_at = time.time()
+            edge.summarized = False
         graph_dic = graph.dict()
 
         node_variable_mapping = create_node_variable_mapping(graph_dic['nodes'])
@@ -411,31 +494,100 @@ class Neo4jGraphDB(AbstractGraphDB):
             return f"{memory_type} deleted for user ID: {user_id}"
         except Exception as e:
             return f"An error occurred: {str(e)}"
-    def retrieve_semantic_memory(self, user_id: str):
+
-        query = f"""
+    def retrieve_semantic_memory(self, user_id: str, timestamp: float = None, summarized: bool = None):
-        MATCH (user:User {{userId: '{user_id}' }})-[:HAS_SEMANTIC_MEMORY]->(semantic:SemanticMemory)
+        if timestamp is not None and summarized is not None:
-        MATCH (semantic)-[:HAS_KNOWLEDGE]->(knowledge)
+            query = f"""
-        RETURN knowledge
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_SEMANTIC_MEMORY]->(semantic:SemanticMemory)
-        """
+            MATCH (semantic)-[:HAS_KNOWLEDGE]->(knowledge)
+            WHERE knowledge.created_at >= {timestamp} AND knowledge.summarized = {str(summarized).lower()}
+            RETURN knowledge
+            """
+        elif timestamp is not None:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_SEMANTIC_MEMORY]->(semantic:SemanticMemory)
+            MATCH (semantic)-[:HAS_KNOWLEDGE]->(knowledge)
+            WHERE knowledge.created_at >= {timestamp}
+            RETURN knowledge
+            """
+        elif summarized is not None:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_SEMANTIC_MEMORY]->(semantic:SemanticMemory)
+            MATCH (semantic)-[:HAS_KNOWLEDGE]->(knowledge)
+            WHERE knowledge.summarized = {str(summarized).lower()}
+            RETURN knowledge
+            """
+        else:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_SEMANTIC_MEMORY]->(semantic:SemanticMemory)
+            MATCH (semantic)-[:HAS_KNOWLEDGE]->(knowledge)
+            RETURN knowledge
+            """
         return self.query(query, params={"user_id": user_id})
 
-    def retrieve_episodic_memory(self, user_id: str):
+    def retrieve_episodic_memory(self, user_id: str, timestamp: float = None, summarized: bool = None):
-        query = """
+        if timestamp is not None and summarized is not None:
-        MATCH (user:User {userId: $user_id})-[:HAS_EPISODIC_MEMORY]->(episodic:EpisodicMemory)
+            query = f"""
-        MATCH (episodic)-[:HAS_EVENT]->(event)
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_EPISODIC_MEMORY]->(episodic:EpisodicMemory)
-        RETURN event
+            MATCH (episodic)-[:HAS_EVENT]->(event)
-        """
+            WHERE event.created_at >= {timestamp} AND event.summarized = {str(summarized).lower()}
+            RETURN event
+            """
+        elif timestamp is not None:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_EPISODIC_MEMORY]->(episodic:EpisodicMemory)
+            MATCH (episodic)-[:HAS_EVENT]->(event)
+            WHERE event.created_at >= {timestamp}
+            RETURN event
+            """
+        elif summarized is not None:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_EPISODIC_MEMORY]->(episodic:EpisodicMemory)
+            MATCH (episodic)-[:HAS_EVENT]->(event)
+            WHERE event.summarized = {str(summarized).lower()}
+            RETURN event
+            """
+        else:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_EPISODIC_MEMORY]->(episodic:EpisodicMemory)
+            MATCH (episodic)-[:HAS_EVENT]->(event)
+            RETURN event
+            """
         return self.query(query, params={"user_id": user_id})
 
-    def retrieve_buffer_memory(self, user_id: str):
+
-        query = """
+    def retrieve_buffer_memory(self, user_id: str, timestamp: float = None, summarized: bool = None):
-        MATCH (user:User {userId: $user_id})-[:HAS_BUFFER]->(buffer:Buffer)
+        if timestamp is not None and summarized is not None:
-        MATCH (buffer)-[:CURRENTLY_HOLDING]->(item)
+            query = f"""
-        RETURN item
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_BUFFER]->(buffer:Buffer)
-        """
+            MATCH (buffer)-[:CURRENTLY_HOLDING]->(item)
+            WHERE item.created_at >= {timestamp} AND item.summarized = {str(summarized).lower()}
+            RETURN item
+            """
+        elif timestamp is not None:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_BUFFER]->(buffer:Buffer)
+            MATCH (buffer)-[:CURRENTLY_HOLDING]->(item)
+            WHERE item.created_at >= {timestamp}
+            RETURN item
+            """
+        elif summarized is not None:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_BUFFER]->(buffer:Buffer)
+            MATCH (buffer)-[:CURRENTLY_HOLDING]->(item)
+            WHERE item.summarized = {str(summarized).lower()}
+            RETURN item
+            """
+        else:
+            query = f"""
+            MATCH (user:User {{userId: '{user_id}' }})-[:HAS_BUFFER]->(buffer:Buffer)
+            MATCH (buffer)-[:CURRENTLY_HOLDING]->(item)
+            RETURN item
+            """
         return self.query(query, params={"user_id": user_id})
 
 
+
     def retrieve_public_memory(self, user_id: str):
         query = """
         MATCH (user:User {userId: $user_id})-[:HAS_PUBLIC_MEMORY]->(public:PublicMemory)
@@ -661,7 +813,8 @@ class Neo4jGraphDB(AbstractGraphDB):
                title: '{title}',
                summary: '{summary}',
                documentCategory: '{document_category}',
-               d_id: '{d_id}'
+               d_id: '{d_id}',
+               created_at: timestamp(),
            }})
 
            // Link the Document node to the {memory_node_type} node