From 3cf6f77c24f7ec83921d19b46b4eb093aca1ca7b Mon Sep 17 00:00:00 2001
From: "sweep-ai[bot]" <128439645+sweep-ai[bot]@users.noreply.github.com>
Date: Sat, 13 Apr 2024 13:19:57 +0000
Subject: [PATCH] feat: Updated agents-api/agents_api/models/entry/t

---
 .../models/entry/test_entry_queries.py        | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/agents-api/agents_api/models/entry/test_entry_queries.py b/agents-api/agents_api/models/entry/test_entry_queries.py
index 7fb5375c..8c206139 100644
--- a/agents-api/agents_api/models/entry/test_entry_queries.py
+++ b/agents-api/agents_api/models/entry/test_entry_queries.py
@@ -1,3 +1,7 @@
+"""
+This module contains tests for entry queries against the CozoDB database.
+It verifies the functionality of adding, retrieving, and processing entries as defined in the schema.
+"""
 # Tests for entry queries
 from uuid import uuid4
 
@@ -19,6 +23,7 @@
 from .proc_mem_context import proc_mem_context_query
 
 
+# Initializes a new CozoDB client for testing, applying all migrations.
 def cozo_client(migrations_dir: str = "./migrations"):
     # Create a new client for each test
     # and initialize the schema.
@@ -32,6 +37,10 @@ def cozo_client(migrations_dir: str = "./migrations"):
 
 @test("model: create entry")
 def _():
+    """
+    Tests the addition of a new entry to the database.
+    Verifies that the entry can be successfully added using the add_entries_query function.
+    """
     client = cozo_client()
     session_id = uuid4()
 
@@ -46,6 +55,10 @@ def _():
 
 @test("model: get entries")
 def _():
+    """
+    Tests the retrieval of entries from the database.
+    Verifies that entries matching specific criteria can be successfully retrieved.
+    """
     client = cozo_client()
     session_id = uuid4()
 
@@ -66,11 +79,16 @@ def _():
 
     result = get_entries_query(session_id=session_id, client=client)
 
+    # Asserts that only one entry is retrieved, matching the session_id.
     assert len(result["entry_id"]) == 1
 
 
 @test("model: procedural memory context")
 def _():
+    """
+    Tests the procedural memory context in the database.
+    Verifies the functionality of retrieving relevant memory context based on embeddings.
+    """
     client = cozo_client()
     developer_id = uuid4()
     user_id = uuid4()
@@ -80,6 +98,7 @@ def _():
     user_doc_id = uuid4()
     agent_doc_id = uuid4()
 
+    # Setup: Creates a user, agent, session, function, and documents, then embeds tools and document snippets.
     # Create stuff
     test_entry = Entry(
         session_id=session_id,
@@ -157,6 +176,7 @@ def _():
         ),
     ]
 
+    # Executes the procedural memory context query to retrieve relevant memory context based on embeddings.
     # Run the query
     result = proc_mem_context_query(
         session_id=session_id,