Skip to content

Commit

Permalink
feat: Updated agents-api/agents_api/models/entry/t
Browse files Browse the repository at this point in the history
  • Loading branch information
sweep-ai[bot] committed Apr 13, 2024
1 parent 4199c15 commit 3cf6f77
Showing 1 changed file with 20 additions and 0 deletions.
20 changes: 20 additions & 0 deletions agents-api/agents_api/models/entry/test_entry_queries.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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.
Expand All @@ -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()

Expand All @@ -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()

Expand All @@ -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()
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 3cf6f77

Please sign in to comment.