From aefe3abf9cca4bbe5dff8486424b3333f7cfd039 Mon Sep 17 00:00:00 2001 From: Diwank Tomer Date: Sun, 11 Aug 2024 18:51:50 -0400 Subject: [PATCH] feat(agents-api): Add settings to developers relation Signed-off-by: Diwank Tomer --- .../agents_api/common/protocol/developers.py | 22 ++++++ .../agents_api/dependencies/developer_id.py | 47 +++++++------ agents-api/agents_api/models/__init__.py | 1 + .../agents_api/models/developer/__init__.py | 19 ++++++ .../models/developer/get_developer.py | 68 +++++++++++++++++++ ...e_1723400730_add_settings_to_developers.py | 68 +++++++++++++++++++ agents-api/tests/fixtures.py | 4 +- agents-api/tests/test_developer_queries.py | 36 ++++++++++ 8 files changed, 242 insertions(+), 23 deletions(-) create mode 100644 agents-api/agents_api/common/protocol/developers.py create mode 100644 agents-api/agents_api/models/developer/__init__.py create mode 100644 agents-api/agents_api/models/developer/get_developer.py create mode 100644 agents-api/migrations/migrate_1723400730_add_settings_to_developers.py create mode 100644 agents-api/tests/test_developer_queries.py diff --git a/agents-api/agents_api/common/protocol/developers.py b/agents-api/agents_api/common/protocol/developers.py new file mode 100644 index 000000000..4c1cfdac4 --- /dev/null +++ b/agents-api/agents_api/common/protocol/developers.py @@ -0,0 +1,22 @@ +""" +This module defines session-related data structures and settings used across the agents API. +It includes definitions for session settings and session data models. +""" + +from uuid import UUID + +from pydantic import AwareDatetime, BaseModel, EmailStr + + +class Developer(BaseModel): + """ + Represents the data associated with a developer + """ + + id: UUID + email: EmailStr + active: bool + tags: list[str] + settings: dict + created_at: AwareDatetime + created_at: AwareDatetime diff --git a/agents-api/agents_api/dependencies/developer_id.py b/agents-api/agents_api/dependencies/developer_id.py index a088389b5..e9b51e850 100644 --- a/agents-api/agents_api/dependencies/developer_id.py +++ b/agents-api/agents_api/dependencies/developer_id.py @@ -1,44 +1,49 @@ -import uuid from typing import Annotated +from uuid import UUID from fastapi import Header -from pydantic import validate_email -from pydantic_core import PydanticCustomError +from ..common.protocol.developers import Developer from ..env import skip_check_developer_headers +from ..models.developer.get_developer import get_developer, verify_developer from .exceptions import InvalidHeaderFormat async def get_developer_id( - x_developer_id: Annotated[uuid.UUID | None, Header()] = None, -): + x_developer_id: Annotated[UUID | None, Header()] = None, +) -> UUID: if skip_check_developer_headers: - return x_developer_id or uuid.UUID("00000000-0000-0000-0000-000000000000") + return UUID("00000000-0000-0000-0000-000000000000") if not x_developer_id: - raise InvalidHeaderFormat("X-Developer-Id header invalid") + raise InvalidHeaderFormat("X-Developer-Id header required") if isinstance(x_developer_id, str): try: - x_developer_id = uuid.UUID(x_developer_id, version=4) - except ValueError: - raise InvalidHeaderFormat("X-Developer-Id must be a valid UUID") + x_developer_id = UUID(x_developer_id, version=4) + except ValueError as e: + raise InvalidHeaderFormat("X-Developer-Id must be a valid UUID") from e + + verify_developer(developer_id=x_developer_id) return x_developer_id -async def get_developer_email( - x_developer_email: Annotated[str | None, Header()] = None, -): +async def get_developer_data( + x_developer_id: Annotated[UUID | None, Header()] = None, +) -> Developer: if skip_check_developer_headers: - return x_developer_email or "unknown_user@mail.com" + x_developer_id = UUID("00000000-0000-0000-0000-000000000000") + + if not x_developer_id: + raise InvalidHeaderFormat("X-Developer-Id header required") - if not x_developer_email: - raise InvalidHeaderFormat("X-Developer-Email header invalid") + if isinstance(x_developer_id, str): + try: + x_developer_id = UUID(x_developer_id, version=4) + except ValueError as e: + raise InvalidHeaderFormat("X-Developer-Id must be a valid UUID") from e - try: - validate_email(x_developer_email) - except PydanticCustomError: - raise InvalidHeaderFormat("X-Developer-Email header invalid") + developer = get_developer(developer_id=x_developer_id) - return x_developer_email + return developer diff --git a/agents-api/agents_api/models/__init__.py b/agents-api/agents_api/models/__init__.py index e84bb5ecc..b1f918ee7 100644 --- a/agents-api/agents_api/models/__init__.py +++ b/agents-api/agents_api/models/__init__.py @@ -9,6 +9,7 @@ # ruff: noqa: F401, F403, F405 import agents_api.models.agent as agent +import agents_api.models.developer as developer import agents_api.models.docs as docs import agents_api.models.entry as entry import agents_api.models.execution as execution diff --git a/agents-api/agents_api/models/developer/__init__.py b/agents-api/agents_api/models/developer/__init__.py new file mode 100644 index 000000000..a7117c06b --- /dev/null +++ b/agents-api/agents_api/models/developer/__init__.py @@ -0,0 +1,19 @@ +""" +Module: agents_api/models/docs + +This module is responsible for managing document-related operations within the application, particularly for agents and possibly other entities. It serves as a core component of the document management system, enabling features such as document creation, listing, deletion, and embedding of snippets for enhanced search and retrieval capabilities. + +Main functionalities include: +- Creating new documents and associating them with agents or users. +- Listing documents based on various criteria, including ownership and metadata filters. +- Deleting documents by their unique identifiers. +- Embedding document snippets for retrieval purposes. + +The module interacts with other parts of the application, such as the agents and users modules, to provide a comprehensive document management system. Its role is crucial in enabling document search, retrieval, and management features within the context of agents and users. + +This documentation aims to provide clear, concise, and sufficient context for new developers or contributors to understand the module's role without needing to dive deep into the code immediately. +""" + +# ruff: noqa: F401, F403, F405 + +from .get_developer import get_developer, verify_developer diff --git a/agents-api/agents_api/models/developer/get_developer.py b/agents-api/agents_api/models/developer/get_developer.py new file mode 100644 index 000000000..c78517613 --- /dev/null +++ b/agents-api/agents_api/models/developer/get_developer.py @@ -0,0 +1,68 @@ +"""Module for retrieving document snippets from the CozoDB based on document IDs.""" + +from uuid import UUID + +from beartype import beartype +from fastapi import HTTPException +from pycozo.client import QueryException +from pydantic import ValidationError + +from ...common.protocol.developers import Developer +from ..utils import ( + cozo_query, + partialclass, + rewrap_exceptions, + verify_developer_id_query, + wrap_in_class, +) + + +@rewrap_exceptions({QueryException: partialclass(HTTPException, status_code=401)}) +@cozo_query +@beartype +def verify_developer( + *, + developer_id: UUID, +) -> tuple[str, dict]: + return (verify_developer_id_query(developer_id), {}) + + +@rewrap_exceptions( + { + QueryException: partialclass(HTTPException, status_code=403), + ValidationError: partialclass(HTTPException, status_code=500), + } +) +@wrap_in_class(Developer, one=True, transform=lambda d: {**d, "id": d["developer_id"]}) +@cozo_query +@beartype +def get_developer( + *, + developer_id: UUID, +) -> tuple[str, dict]: + developer_id = str(developer_id) + + query = """ + input[developer_id] <- [[to_uuid($developer_id)]] + ?[ + developer_id, + email, + active, + tags, + settings, + created_at, + updated_at, + ] := + input[developer_id], + *developers { + developer_id, + email, + active, + tags, + settings, + created_at, + updated_at, + } + """ + + return (query, {"developer_id": developer_id}) diff --git a/agents-api/migrations/migrate_1723400730_add_settings_to_developers.py b/agents-api/migrations/migrate_1723400730_add_settings_to_developers.py new file mode 100644 index 000000000..e10e71510 --- /dev/null +++ b/agents-api/migrations/migrate_1723400730_add_settings_to_developers.py @@ -0,0 +1,68 @@ +# /usr/bin/env python3 + +MIGRATION_ID = "add_settings_to_developers" +CREATED_AT = 1723400730.539554 + + +def up(client): + client.run( + """ + ?[ + developer_id, + email, + active, + tags, + settings, + created_at, + updated_at, + ] := *developers { + developer_id, + email, + active, + created_at, + updated_at, + }, + tags = [], + settings = {} + + :replace developers { + developer_id: Uuid, + => + email: String, + active: Bool default true, + tags: [String] default [], + settings: Json, + created_at: Float default now(), + updated_at: Float default now(), + } + """ + ) + + +def down(client): + client.run( + """ + ?[ + developer_id, + email, + active, + created_at, + updated_at, + ] := *developers { + developer_id, + email, + active, + created_at, + updated_at, + } + + :replace developers { + developer_id: Uuid, + => + email: String, + active: Bool default true, + created_at: Float default now(), + updated_at: Float default now(), + } + """ + ) diff --git a/agents-api/tests/fixtures.py b/agents-api/tests/fixtures.py index e04cb316d..f2b76e3e3 100644 --- a/agents-api/tests/fixtures.py +++ b/agents-api/tests/fixtures.py @@ -50,8 +50,8 @@ def test_developer_id(cozo_client=cozo_client): cozo_client.run( f""" - ?[developer_id, email] <- [["{str(developer_id)}", "developers@julep.ai"]] - :insert developers {{ developer_id, email }} + ?[developer_id, email, settings] <- [["{str(developer_id)}", "developers@julep.ai", {{}}]] + :insert developers {{ developer_id, email, settings }} """ ) diff --git a/agents-api/tests/test_developer_queries.py b/agents-api/tests/test_developer_queries.py new file mode 100644 index 000000000..569733fa5 --- /dev/null +++ b/agents-api/tests/test_developer_queries.py @@ -0,0 +1,36 @@ +# Tests for agent queries +from uuid import uuid4 + +from ward import raises, test + +from agents_api.common.protocol.developers import Developer +from agents_api.models.developer.get_developer import get_developer, verify_developer +from tests.fixtures import cozo_client, test_developer_id + + +@test("model: get developer") +def _(client=cozo_client, developer_id=test_developer_id): + developer = get_developer( + developer_id=developer_id, + client=client, + ) + + assert isinstance(developer, Developer) + assert developer.id + + +@test("model: verify developer exists") +def _(client=cozo_client, developer_id=test_developer_id): + verify_developer( + developer_id=developer_id, + client=client, + ) + + +@test("model: verify developer not exists") +def _(client=cozo_client): + with raises(Exception): + verify_developer( + developer_id=uuid4(), + client=client, + )