Skip to content

Commit

Permalink
♻️ Move schemas to own package (#1041)
Browse files Browse the repository at this point in the history
  • Loading branch information
awtkns authored Jul 13, 2023
1 parent 3a214dc commit 709821b
Show file tree
Hide file tree
Showing 33 changed files with 57 additions and 51 deletions.
3 changes: 2 additions & 1 deletion platform/reworkd_platform/db/crud/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@

from reworkd_platform.db.crud.base import BaseCrud
from reworkd_platform.db.models.agent import AgentRun, AgentTask
from reworkd_platform.schemas import Loop_Step, UserBase
from reworkd_platform.schemas.agent import Loop_Step
from reworkd_platform.schemas.user import UserBase
from reworkd_platform.settings import settings
from reworkd_platform.web.api.errors import MaxLoopsError, MultipleSummaryError

Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/db/crud/edge.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from sqlalchemy.ext.asyncio import AsyncSession

from reworkd_platform.db.models.workflow import WorkflowEdgeModel
from reworkd_platform.web.api.workflow.schemas import EdgeUpsert
from reworkd_platform.schemas.workflow.base import EdgeUpsert


class EdgeCRUD:
Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/db/crud/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from sqlalchemy.ext.asyncio import AsyncSession

from reworkd_platform.db.models.workflow import NodeBlockModel, WorkflowNodeModel
from reworkd_platform.web.api.workflow.schemas import NodeUpsert
from reworkd_platform.schemas.workflow.base import NodeUpsert


class NodeCRUD:
Expand Down
6 changes: 3 additions & 3 deletions platform/reworkd_platform/db/crud/workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,13 @@
from reworkd_platform.db.crud.node import NodeCRUD
from reworkd_platform.db.dependencies import get_db_session
from reworkd_platform.db.models.workflow import WorkflowModel
from reworkd_platform.schemas import UserBase
from reworkd_platform.web.api.dependencies import get_current_user
from reworkd_platform.web.api.workflow.schemas import (
from reworkd_platform.schemas.user import UserBase
from reworkd_platform.schemas.workflow.base import (
Workflow,
WorkflowFull,
WorkflowUpdate,
)
from reworkd_platform.web.api.dependencies import get_current_user


class WorkflowCRUD(BaseCrud):
Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/db/models/workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from sqlalchemy.orm import mapped_column, relationship

from reworkd_platform.db.base import Base, TrackedModel, UserMixin
from reworkd_platform.web.api.workflow.schemas import Block, Edge, Node, Workflow
from reworkd_platform.schemas.workflow.base import Block, Edge, Node, Workflow


class WorkflowModel(TrackedModel, UserMixin):
Expand Down
Empty file.
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional
from typing import Literal, Dict, Optional, Any, List

from pydantic import BaseModel, Field, validator

Expand All @@ -10,7 +10,6 @@
"gpt-3.5-turbo-16k",
"gpt-4",
]

Loop_Step = Literal[
"start",
"analyze",
Expand All @@ -19,7 +18,6 @@
"summarize",
"chat",
]

LLM_MODEL_MAX_TOKENS: Dict[LLM_Model, int] = {
"gpt-3.5-turbo": 4000,
"gpt-3.5-turbo-16k": 16000,
Expand Down Expand Up @@ -87,9 +85,3 @@ class RunCount(BaseModel):
count: int
first_run: Optional[datetime]
last_run: Optional[datetime]


class UserBase(BaseModel):
id: str
name: Optional[str]
email: Optional[str]
9 changes: 9 additions & 0 deletions platform/reworkd_platform/schemas/user.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
from typing import Optional

from pydantic import BaseModel


class UserBase(BaseModel):
id: str
name: Optional[str]
email: Optional[str]
Empty file.
Empty file.
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from loguru import logger
from requests import RequestException

from reworkd_platform.web.api.workflow.schemas import Block, BlockIOBase
from reworkd_platform.schemas.workflow.base import Block, BlockIOBase


class UrlStatusCheckBlockInput(BlockIOBase):
Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/services/kafka/event_schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from pydantic import BaseModel

from reworkd_platform.web.api.workflow.schemas import Node
from reworkd_platform.schemas.workflow.base import Node


class WorkflowTaskEvent(BaseModel):
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from tiktoken import Encoding

from reworkd_platform.schemas import LLM_MODEL_MAX_TOKENS
from reworkd_platform.schemas.agent import LLM_MODEL_MAX_TOKENS
from reworkd_platform.web.api.agent.model_settings import WrappedChatOpenAI


Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/services/worker/exec.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
from loguru import logger
from networkx import topological_sort

from reworkd_platform.schemas.workflow.base import WorkflowFull
from reworkd_platform.services.kafka.event_schemas import WorkflowTaskEvent
from reworkd_platform.services.kafka.producers.task_producer import WorkflowTaskProducer
from reworkd_platform.services.sockets import websockets
from reworkd_platform.web.api.workflow.blocks.web import get_block_runner
from reworkd_platform.web.api.workflow.schemas import WorkflowFull


class ExecutionEngine:
Expand Down
3 changes: 2 additions & 1 deletion platform/reworkd_platform/tests/agent/create_model_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@

import pytest

from reworkd_platform.schemas import ModelSettings, UserBase
from reworkd_platform.schemas.agent import ModelSettings
from reworkd_platform.schemas.user import UserBase
from reworkd_platform.web.api.agent.model_settings import create_model


Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/tests/test_api_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import pytest

from reworkd_platform.schemas import LLM_Model
from reworkd_platform.schemas.agent import LLM_Model
from reworkd_platform.web.api.agent.api_utils import rotate_keys

ITERATIONS = 10000
Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/tests/test_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
InvalidRequestError,
)

from reworkd_platform.schemas import ModelSettings
from reworkd_platform.schemas.agent import ModelSettings
from reworkd_platform.web.api.agent.helpers import openai_error_handler
from reworkd_platform.web.api.errors import OpenAIError

Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/tests/test_schemas.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import pytest

from reworkd_platform.schemas import ModelSettings
from reworkd_platform.schemas.agent import ModelSettings


@pytest.mark.parametrize(
Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/tests/test_token_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import tiktoken

from reworkd_platform.schemas import LLM_MODEL_MAX_TOKENS
from reworkd_platform.schemas.agent import LLM_MODEL_MAX_TOKENS
from reworkd_platform.services.tokenizer.token_service import TokenService

encoding = tiktoken.get_encoding("cl100k_base")
Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/tests/workflow/test_edge_crud.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from pytest_mock import MockFixture

from reworkd_platform.db.crud.edge import EdgeCRUD
from reworkd_platform.web.api.workflow.schemas import EdgeUpsert
from reworkd_platform.schemas.workflow.base import EdgeUpsert


class TestEdgeCRUD:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

from reworkd_platform.db.crud.workflow import WorkflowCRUD
from reworkd_platform.db.models.workflow import WorkflowModel
from reworkd_platform.schemas import UserBase
from reworkd_platform.schemas.user import UserBase


@pytest.mark.asyncio
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@

from fastapi import Depends

from reworkd_platform.schemas import AgentRun, UserBase
from reworkd_platform.schemas.agent import AgentRun
from reworkd_platform.schemas.user import UserBase
from reworkd_platform.services.tokenizer.dependencies import get_token_service
from reworkd_platform.services.tokenizer.token_service import TokenService
from reworkd_platform.settings import settings
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from loguru import logger
from pydantic import ValidationError

from reworkd_platform.schemas import ModelSettings
from reworkd_platform.schemas.agent import ModelSettings
from reworkd_platform.services.tokenizer.token_service import TokenService
from reworkd_platform.web.api.agent.agent_service.agent_service import AgentService
from reworkd_platform.web.api.agent.analysis import Analysis, AnalysisArguments
Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/web/api/agent/api_utils.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import random
from typing import Optional

from reworkd_platform.schemas import LLM_Model
from reworkd_platform.schemas.agent import LLM_Model

PRIMARY_KEY_RATE = 0.7
SECONDARY_KEY_RATE = 1 - PRIMARY_KEY_RATE
Expand Down
14 changes: 7 additions & 7 deletions platform/reworkd_platform/web/api/agent/dependancies.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,17 +5,17 @@

from reworkd_platform.db.crud.agent import AgentCRUD
from reworkd_platform.db.dependencies import get_db_session
from reworkd_platform.schemas import (
AgentChat,
AgentRun,
from reworkd_platform.schemas.agent import (
Loop_Step,
AgentRunCreate,
AgentSummarize,
AgentRun,
AgentTaskAnalyze,
AgentTaskCreate,
AgentTaskExecute,
Loop_Step,
UserBase,
AgentTaskCreate,
AgentSummarize,
AgentChat,
)
from reworkd_platform.schemas.user import UserBase
from reworkd_platform.services.pinecone.pinecone import PineconeMemory
from reworkd_platform.services.vecs.dependencies import get_supabase_vecs
from reworkd_platform.services.vecs.vecs import VecsMemory
Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/web/api/agent/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
InvalidRequestError,
)

from reworkd_platform.schemas import ModelSettings
from reworkd_platform.schemas.agent import ModelSettings
from reworkd_platform.web.api.errors import OpenAIError

T = TypeVar("T")
Expand Down
3 changes: 2 additions & 1 deletion platform/reworkd_platform/web/api/agent/model_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@
from langchain.chat_models import ChatOpenAI
from pydantic import Field

from reworkd_platform.schemas import LLM_Model, ModelSettings, UserBase
from reworkd_platform.schemas.agent import LLM_Model, ModelSettings
from reworkd_platform.schemas.user import UserBase
from reworkd_platform.settings import settings
from reworkd_platform.web.api.agent.api_utils import rotate_keys

Expand Down
8 changes: 4 additions & 4 deletions platform/reworkd_platform/web/api/agent/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,13 @@
from fastapi.responses import StreamingResponse as FastAPIStreamingResponse
from pydantic import BaseModel

from reworkd_platform.schemas import (
AgentChat,
from reworkd_platform.schemas.agent import (
AgentRun,
AgentSummarize,
AgentTaskAnalyze,
AgentTaskCreate,
AgentTaskExecute,
AgentTaskCreate,
AgentSummarize,
AgentChat,
NewTasksResponse,
)
from reworkd_platform.web.api.agent.agent_service.agent_service import AgentService
Expand Down
2 changes: 1 addition & 1 deletion platform/reworkd_platform/web/api/dependencies.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from reworkd_platform.db.crud.user import UserCrud
from reworkd_platform.db.dependencies import get_db_session
from reworkd_platform.db.models.user import User
from reworkd_platform.schemas import UserBase
from reworkd_platform.schemas.user import UserBase


def user_crud(
Expand Down
3 changes: 2 additions & 1 deletion platform/reworkd_platform/web/api/models/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@
from fastapi import APIRouter, Depends
from pydantic import BaseModel, Field

from reworkd_platform.schemas import LLM_MODEL_MAX_TOKENS, UserBase
from reworkd_platform.schemas.agent import LLM_MODEL_MAX_TOKENS
from reworkd_platform.schemas.user import UserBase
from reworkd_platform.web.api.dependencies import get_current_user

router = APIRouter()
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from reworkd_platform.web.api.workflow.blocks.web.status_check import (
from reworkd_platform.schemas.workflow.base import Block
from reworkd_platform.schemas.workflow.blocks.status_check import (
UrlStatusCheckBlock,
)
from reworkd_platform.web.api.workflow.schemas import Block


def get_block_runner(block: Block) -> Block:
Expand Down
8 changes: 4 additions & 4 deletions platform/reworkd_platform/web/api/workflow/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,14 @@
from fastapi import APIRouter, Depends, HTTPException

from reworkd_platform.db.crud.workflow import WorkflowCRUD
from reworkd_platform.services.kafka.producers.task_producer import WorkflowTaskProducer
from reworkd_platform.services.networkx import validate_connected_and_acyclic
from reworkd_platform.services.worker.exec import ExecutionEngine
from reworkd_platform.web.api.workflow.schemas import (
from reworkd_platform.schemas.workflow.base import (
Workflow,
WorkflowUpdate,
WorkflowFull,
)
from reworkd_platform.services.kafka.producers.task_producer import WorkflowTaskProducer
from reworkd_platform.services.networkx import validate_connected_and_acyclic
from reworkd_platform.services.worker.exec import ExecutionEngine

router = APIRouter()

Expand Down

0 comments on commit 709821b

Please sign in to comment.