-
Notifications
You must be signed in to change notification settings - Fork 870
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Reimplement tests for queries and operations (#443)
* Reimplement tests for queries and operations Reimplement tests for queries and operations in `agents-api/agents_api/models` subdirectories using the new code. * **agents-api/agents_api/models/agent/test_agent_queries.py** - Uncomment import statements and test functions. - Update test functions to use new query functions. - Ensure tests pass with new code. * **agents-api/agents_api/models/docs/test_docs_queries.py** - Uncomment import statements and test functions. - Update test functions to use new query functions. - Ensure tests pass with new code. * **agents-api/agents_api/models/entry/test_entry_queries.py** - Uncomment import statements and test functions. - Update test functions to use new query functions. - Ensure tests pass with new code. * **agents-api/agents_api/models/execution/test_execution_queries.py** - Uncomment import statements and test functions. - Update test functions to use new query functions. - Ensure tests pass with new code. * **agents-api/agents_api/models/session/test_session_queries.py** - Uncomment import statements and test functions. - Update test functions to use new query functions. - Ensure tests pass with new code. * **agents-api/agents_api/models/task/test_task_queries.py** - Uncomment import statements and test functions. - Update test functions to use new query functions. - Ensure tests pass with new code. * **agents-api/agents_api/models/tools/test_tool_queries.py** - Uncomment import statements and test functions. - Update test functions to use new query functions. - Ensure tests pass with new code. * **agents-api/agents_api/models/user/test_user_queries.py** - Uncomment import statements and test functions. - Update test functions to use new query functions. - Ensure tests pass with new code. --- For more details, open the [Copilot Workspace session](https://copilot-workspace.githubnext.com/julep-ai/julep?shareId=XXXX-XXXX-XXXX-XXXX). * wip: reimplement model tests Signed-off-by: Diwank Tomer <[email protected]> --------- Signed-off-by: Diwank Tomer <[email protected]> Co-authored-by: Diwank Tomer <[email protected]>
- Loading branch information
Showing
16 changed files
with
2,982 additions
and
3,022 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,22 @@ | ||
| from typing import Type | ||
|
|
||
| from beartype.vale import Is | ||
| from beartype.vale._core._valecore import BeartypeValidator | ||
| from pydantic import BaseModel | ||
|
|
||
|
|
||
| def dict_like(pydantic_model_class: Type[BaseModel]) -> BeartypeValidator: | ||
| required_fields_set: set[str] = set( | ||
| [ | ||
| field | ||
| for field, info in pydantic_model_class.model_fields.items() | ||
| if info.is_required() | ||
| ] | ||
| ) | ||
|
|
||
| validator = Is[ | ||
| lambda x: isinstance(x, pydantic_model_class) | ||
| or required_fields_set.issubset(set(x.keys())) | ||
| ] | ||
|
|
||
| return validator |
339 changes: 172 additions & 167 deletions
339
agents-api/agents_api/models/agent/test_agent_queries.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,167 +1,172 @@ | ||
| # # Tests for agent queries | ||
| # from uuid import uuid4 | ||
|
|
||
| # from cozo_migrate.api import init, apply | ||
| # from pycozo import Client | ||
| # from ward import test | ||
|
|
||
| # from .create_agent import create_agent_query | ||
| # from .delete_agent import delete_agent_query | ||
| # from .get_agent import get_agent_query | ||
| # from .list_agents import list_agents_query | ||
| # from .update_agent import update_agent_query | ||
|
|
||
| # MODEL = "julep-ai/samantha-1-turbo" | ||
|
|
||
|
|
||
| # def cozo_client(migrations_dir: str = "./migrations"): | ||
| # # Create a new client for each test | ||
| # # and initialize the schema. | ||
| # client = Client() | ||
|
|
||
| # init(client) | ||
| # apply(client, migrations_dir=migrations_dir, all_=True) | ||
|
|
||
| # return client | ||
|
|
||
|
|
||
| # @test("model: create agent") | ||
| # def _(): | ||
| # client = cozo_client() | ||
| # agent_id = uuid4() | ||
| # developer_id = uuid4() | ||
|
|
||
| # create_agent_query( | ||
| # agent_id=agent_id, | ||
| # model=MODEL, | ||
| # developer_id=developer_id, | ||
| # name="test agent", | ||
| # about="test agent about", | ||
| # client=client, | ||
| # ) | ||
|
|
||
|
|
||
| # @test("model: create agent with instructions") | ||
| # def _(): | ||
| # client = cozo_client() | ||
| # agent_id = uuid4() | ||
| # developer_id = uuid4() | ||
|
|
||
| # create_agent_query( | ||
| # agent_id=agent_id, | ||
| # model=MODEL, | ||
| # developer_id=developer_id, | ||
| # name="test agent", | ||
| # about="test agent about", | ||
| # instructions=[ | ||
| # "test instruction", | ||
| # ], | ||
| # client=client, | ||
| # ) | ||
|
|
||
|
|
||
| # @test("model: get agent not exists") | ||
| # def _(): | ||
| # client = cozo_client() | ||
| # agent_id = uuid4() | ||
| # developer_id = uuid4() | ||
|
|
||
| # result = get_agent_query( | ||
| # agent_id=agent_id, developer_id=developer_id, client=client | ||
| # ) | ||
|
|
||
| # assert len(result["id"]) == 0 | ||
|
|
||
|
|
||
| # @test("model: get agent exists") | ||
| # def _(): | ||
| # client = cozo_client() | ||
| # agent_id = uuid4() | ||
| # developer_id = uuid4() | ||
|
|
||
| # result = create_agent_query( | ||
| # agent_id=agent_id, | ||
| # model=MODEL, | ||
| # developer_id=developer_id, | ||
| # name="test agent", | ||
| # about="test agent about", | ||
| # default_settings={"temperature": 1.5}, | ||
| # client=client, | ||
| # ) | ||
|
|
||
| # result = get_agent_query( | ||
| # agent_id=agent_id, developer_id=developer_id, client=client | ||
| # ) | ||
|
|
||
| # assert len(result["id"]) == 1 | ||
| # assert "temperature" in result["default_settings"][0] | ||
| # assert result["default_settings"][0]["temperature"] == 1.5 | ||
|
|
||
|
|
||
| # @test("model: delete agent") | ||
| # def _(): | ||
| # client = cozo_client() | ||
| # agent_id = uuid4() | ||
| # developer_id = uuid4() | ||
|
|
||
| # # Create the agent | ||
| # result = create_agent_query( | ||
| # agent_id=agent_id, | ||
| # model=MODEL, | ||
| # developer_id=developer_id, | ||
| # name="test agent", | ||
| # about="test agent about", | ||
| # client=client, | ||
| # ) | ||
|
|
||
| # # Delete the agent | ||
| # result = delete_agent_query( | ||
| # agent_id=agent_id, developer_id=developer_id, client=client | ||
| # ) | ||
|
|
||
| # # Check that the agent is deleted | ||
| # result = get_agent_query( | ||
| # agent_id=agent_id, developer_id=developer_id, client=client | ||
| # ) | ||
|
|
||
| # assert len(result["id"]) == 0 | ||
|
|
||
|
|
||
| # @test("model: update agent") | ||
| # def _(): | ||
| # client = cozo_client() | ||
| # agent_id = uuid4() | ||
| # developer_id = uuid4() | ||
|
|
||
| # create_agent_query( | ||
| # agent_id=agent_id, | ||
| # model=MODEL, | ||
| # developer_id=developer_id, | ||
| # name="test agent", | ||
| # about="test agent about", | ||
| # client=client, | ||
| # ) | ||
|
|
||
| # result = update_agent_query( | ||
| # agent_id=agent_id, | ||
| # developer_id=developer_id, | ||
| # name="updated agent", | ||
| # about="updated agent about", | ||
| # default_settings={"temperature": 1.5}, | ||
| # client=client, | ||
| # ) | ||
|
|
||
| # data = result.iloc[0].to_dict() | ||
|
|
||
| # assert data["updated_at"] > data["created_at"] | ||
|
|
||
|
|
||
| # @test("model: list agents") | ||
| # def _(): | ||
| # """Tests listing all agents associated with a developer in the database. Verifies that the correct list of agents is retrieved.""" | ||
| # client = cozo_client() | ||
| # developer_id = uuid4() | ||
|
|
||
| # result = list_agents_query(developer_id=developer_id, client=client) | ||
|
|
||
| # assert len(result["id"]) == 0 | ||
| # Tests for agent queries | ||
| from uuid import uuid4 | ||
|
|
||
| from cozo_migrate.api import apply, init | ||
| from pycozo import Client | ||
| from ward import test | ||
|
|
||
| from agents_api.autogen.openapi_model import Agent | ||
|
|
||
| from .create_agent import create_agent | ||
| from .delete_agent import delete_agent | ||
| from .get_agent import get_agent | ||
| from .list_agents import list_agents | ||
| from .update_agent import update_agent | ||
|
|
||
| MODEL = "julep-ai/samantha-1-turbo" | ||
|
|
||
|
|
||
| def cozo_client(migrations_dir: str = "./migrations"): | ||
| # Create a new client for each test | ||
| # and initialize the schema. | ||
| client = Client() | ||
|
|
||
| init(client) | ||
| apply(client, migrations_dir=migrations_dir, all_=True) | ||
|
|
||
| return client | ||
|
|
||
|
|
||
| @test("model: create agent") | ||
| def _(): | ||
| client = cozo_client() | ||
| agent_id = uuid4() | ||
| developer_id = uuid4() | ||
|
|
||
| create_agent( | ||
| agent_id=agent_id, | ||
| developer_id=developer_id, | ||
| data={ | ||
| "model": MODEL, | ||
| "name": "test agent", | ||
| "about": "test agent about", | ||
| }, | ||
| client=client, | ||
| ) | ||
|
|
||
|
|
||
| @test("model: create agent with instructions") | ||
| def _(): | ||
| client = cozo_client() | ||
| agent_id = uuid4() | ||
| developer_id = uuid4() | ||
|
|
||
| create_agent( | ||
| agent_id=agent_id, | ||
| developer_id=developer_id, | ||
| data={ | ||
| "model": MODEL, | ||
| "name": "test agent", | ||
| "about": "test agent about", | ||
| "instructions": ["test instruction"], | ||
| }, | ||
| client=client, | ||
| ) | ||
|
|
||
|
|
||
| @test("model: get agent not exists") | ||
| def _(): | ||
| client = cozo_client() | ||
| agent_id = uuid4() | ||
| developer_id = uuid4() | ||
|
|
||
| result = get_agent(agent_id=agent_id, developer_id=developer_id, client=client) | ||
|
|
||
| assert result is None | ||
|
|
||
|
|
||
| @test("model: get agent exists") | ||
| def _(): | ||
| client = cozo_client() | ||
| agent_id = uuid4() | ||
| developer_id = uuid4() | ||
|
|
||
| create_agent( | ||
| agent_id=agent_id, | ||
| developer_id=developer_id, | ||
| data={ | ||
| "model": MODEL, | ||
| "name": "test agent", | ||
| "about": "test agent about", | ||
| "default_settings": {"temperature": 1.5}, | ||
| }, | ||
| client=client, | ||
| ) | ||
|
|
||
| result = get_agent(agent_id=agent_id, developer_id=developer_id, client=client) | ||
|
|
||
| assert result is not None | ||
| assert isinstance(result, Agent) | ||
| assert result.default_settings.temperature == 1.5 | ||
|
|
||
|
|
||
| @test("model: delete agent") | ||
| def _(): | ||
| client = cozo_client() | ||
| agent_id = uuid4() | ||
| developer_id = uuid4() | ||
|
|
||
| # Create the agent | ||
| create_agent( | ||
| agent_id=agent_id, | ||
| developer_id=developer_id, | ||
| data={ | ||
| "model": MODEL, | ||
| "name": "test agent", | ||
| "about": "test agent about", | ||
| }, | ||
| client=client, | ||
| ) | ||
|
|
||
| # Delete the agent | ||
| delete_agent(agent_id=agent_id, developer_id=developer_id, client=client) | ||
|
|
||
| # Check that the agent is deleted | ||
| result = get_agent(agent_id=agent_id, developer_id=developer_id, client=client) | ||
|
|
||
| assert result is None | ||
|
|
||
|
|
||
| @test("model: update agent") | ||
| def _(): | ||
| client = cozo_client() | ||
| agent_id = uuid4() | ||
| developer_id = uuid4() | ||
|
|
||
| create_agent( | ||
| agent_id=agent_id, | ||
| developer_id=developer_id, | ||
| data={ | ||
| "model": MODEL, | ||
| "name": "test agent", | ||
| "about": "test agent about", | ||
| }, | ||
| client=client, | ||
| ) | ||
|
|
||
| result = update_agent( | ||
| agent_id=agent_id, | ||
| developer_id=developer_id, | ||
| data={ | ||
| "name": "updated agent", | ||
| "about": "updated agent about", | ||
| "default_settings": {"temperature": 1.5}, | ||
| }, | ||
| client=client, | ||
| ) | ||
|
|
||
| assert result is not None | ||
| assert isinstance(result, Agent) | ||
| assert result.default_settings.temperature == 1.5 | ||
|
|
||
|
|
||
| @test("model: list agents") | ||
| def _(): | ||
| """Tests listing all agents associated with a developer in the database. Verifies that the correct list of agents is retrieved.""" | ||
| client = cozo_client() | ||
| developer_id = uuid4() | ||
|
|
||
| result = list_agents(developer_id=developer_id, client=client) | ||
|
|
||
| assert isinstance(result, list) | ||
| assert all(isinstance(agent, Agent) for agent in result) |
Oops, something went wrong.