diff --git a/platform/reworkd_platform/db/crud/agent.py b/platform/reworkd_platform/db/crud/agent.py index 3a66fc8d6e..4762bfebe2 100644 --- a/platform/reworkd_platform/db/crud/agent.py +++ b/platform/reworkd_platform/db/crud/agent.py @@ -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 diff --git a/platform/reworkd_platform/db/crud/edge.py b/platform/reworkd_platform/db/crud/edge.py index 829094e54e..17c4725e88 100644 --- a/platform/reworkd_platform/db/crud/edge.py +++ b/platform/reworkd_platform/db/crud/edge.py @@ -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: diff --git a/platform/reworkd_platform/db/crud/node.py b/platform/reworkd_platform/db/crud/node.py index 8f26873cc6..ddf9e4ea44 100644 --- a/platform/reworkd_platform/db/crud/node.py +++ b/platform/reworkd_platform/db/crud/node.py @@ -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: diff --git a/platform/reworkd_platform/db/crud/workflow.py b/platform/reworkd_platform/db/crud/workflow.py index f44dbf7a38..9e4b3578f1 100644 --- a/platform/reworkd_platform/db/crud/workflow.py +++ b/platform/reworkd_platform/db/crud/workflow.py @@ -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): diff --git a/platform/reworkd_platform/db/models/workflow.py b/platform/reworkd_platform/db/models/workflow.py index 14c17d28b4..35196271d5 100644 --- a/platform/reworkd_platform/db/models/workflow.py +++ b/platform/reworkd_platform/db/models/workflow.py @@ -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): diff --git a/platform/reworkd_platform/schemas/__init__.py b/platform/reworkd_platform/schemas/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/platform/reworkd_platform/schemas.py b/platform/reworkd_platform/schemas/agent.py similarity index 93% rename from platform/reworkd_platform/schemas.py rename to platform/reworkd_platform/schemas/agent.py index eef46458be..9edc98b63c 100644 --- a/platform/reworkd_platform/schemas.py +++ b/platform/reworkd_platform/schemas/agent.py @@ -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 @@ -10,7 +10,6 @@ "gpt-3.5-turbo-16k", "gpt-4", ] - Loop_Step = Literal[ "start", "analyze", @@ -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, @@ -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] diff --git a/platform/reworkd_platform/schemas/user.py b/platform/reworkd_platform/schemas/user.py new file mode 100644 index 0000000000..c78bdb9ade --- /dev/null +++ b/platform/reworkd_platform/schemas/user.py @@ -0,0 +1,9 @@ +from typing import Optional + +from pydantic import BaseModel + + +class UserBase(BaseModel): + id: str + name: Optional[str] + email: Optional[str] diff --git a/platform/reworkd_platform/schemas/workflow/__init__.py b/platform/reworkd_platform/schemas/workflow/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/platform/reworkd_platform/web/api/workflow/schemas.py b/platform/reworkd_platform/schemas/workflow/base.py similarity index 100% rename from platform/reworkd_platform/web/api/workflow/schemas.py rename to platform/reworkd_platform/schemas/workflow/base.py diff --git a/platform/reworkd_platform/schemas/workflow/blocks/__init__.py b/platform/reworkd_platform/schemas/workflow/blocks/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/platform/reworkd_platform/web/api/workflow/blocks/web/status_check.py b/platform/reworkd_platform/schemas/workflow/blocks/status_check.py similarity index 91% rename from platform/reworkd_platform/web/api/workflow/blocks/web/status_check.py rename to platform/reworkd_platform/schemas/workflow/blocks/status_check.py index 5846ca5ba7..09e3d4e339 100644 --- a/platform/reworkd_platform/web/api/workflow/blocks/web/status_check.py +++ b/platform/reworkd_platform/schemas/workflow/blocks/status_check.py @@ -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): diff --git a/platform/reworkd_platform/services/kafka/event_schemas.py b/platform/reworkd_platform/services/kafka/event_schemas.py index 30787ab401..1bd65acc6f 100644 --- a/platform/reworkd_platform/services/kafka/event_schemas.py +++ b/platform/reworkd_platform/services/kafka/event_schemas.py @@ -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): diff --git a/platform/reworkd_platform/services/tokenizer/token_service.py b/platform/reworkd_platform/services/tokenizer/token_service.py index d1eef326a8..e4aad0fbcd 100644 --- a/platform/reworkd_platform/services/tokenizer/token_service.py +++ b/platform/reworkd_platform/services/tokenizer/token_service.py @@ -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 diff --git a/platform/reworkd_platform/services/worker/exec.py b/platform/reworkd_platform/services/worker/exec.py index a5ea08946a..cfd97d318e 100644 --- a/platform/reworkd_platform/services/worker/exec.py +++ b/platform/reworkd_platform/services/worker/exec.py @@ -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: diff --git a/platform/reworkd_platform/tests/agent/create_model_test.py b/platform/reworkd_platform/tests/agent/create_model_test.py index 94639b551b..dd8dc0078c 100644 --- a/platform/reworkd_platform/tests/agent/create_model_test.py +++ b/platform/reworkd_platform/tests/agent/create_model_test.py @@ -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 diff --git a/platform/reworkd_platform/tests/test_api_utils.py b/platform/reworkd_platform/tests/test_api_utils.py index 27631c836e..c7a627bdaa 100644 --- a/platform/reworkd_platform/tests/test_api_utils.py +++ b/platform/reworkd_platform/tests/test_api_utils.py @@ -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 diff --git a/platform/reworkd_platform/tests/test_helpers.py b/platform/reworkd_platform/tests/test_helpers.py index 4b23975adc..18d0719565 100644 --- a/platform/reworkd_platform/tests/test_helpers.py +++ b/platform/reworkd_platform/tests/test_helpers.py @@ -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 diff --git a/platform/reworkd_platform/tests/test_schemas.py b/platform/reworkd_platform/tests/test_schemas.py index 1be72d5b79..095c30dd63 100644 --- a/platform/reworkd_platform/tests/test_schemas.py +++ b/platform/reworkd_platform/tests/test_schemas.py @@ -1,6 +1,6 @@ import pytest -from reworkd_platform.schemas import ModelSettings +from reworkd_platform.schemas.agent import ModelSettings @pytest.mark.parametrize( diff --git a/platform/reworkd_platform/tests/test_token_service.py b/platform/reworkd_platform/tests/test_token_service.py index 15f3a4cc1c..8609c0730e 100644 --- a/platform/reworkd_platform/tests/test_token_service.py +++ b/platform/reworkd_platform/tests/test_token_service.py @@ -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") diff --git a/platform/reworkd_platform/tests/workflow/test_edge_crud.py b/platform/reworkd_platform/tests/workflow/test_edge_crud.py index 48b5e8cb47..8a7f243a33 100644 --- a/platform/reworkd_platform/tests/workflow/test_edge_crud.py +++ b/platform/reworkd_platform/tests/workflow/test_edge_crud.py @@ -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: diff --git a/platform/reworkd_platform/tests/workflow/test_workflow_crud.py b/platform/reworkd_platform/tests/workflow/test_workflow_crud.py index 613709d958..5d73e52372 100644 --- a/platform/reworkd_platform/tests/workflow/test_workflow_crud.py +++ b/platform/reworkd_platform/tests/workflow/test_workflow_crud.py @@ -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 diff --git a/platform/reworkd_platform/web/api/agent/agent_service/agent_service_provider.py b/platform/reworkd_platform/web/api/agent/agent_service/agent_service_provider.py index 8096fbbe51..4241c44448 100644 --- a/platform/reworkd_platform/web/api/agent/agent_service/agent_service_provider.py +++ b/platform/reworkd_platform/web/api/agent/agent_service/agent_service_provider.py @@ -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 diff --git a/platform/reworkd_platform/web/api/agent/agent_service/open_ai_agent_service.py b/platform/reworkd_platform/web/api/agent/agent_service/open_ai_agent_service.py index 30cd0231b1..f41d20d294 100644 --- a/platform/reworkd_platform/web/api/agent/agent_service/open_ai_agent_service.py +++ b/platform/reworkd_platform/web/api/agent/agent_service/open_ai_agent_service.py @@ -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 diff --git a/platform/reworkd_platform/web/api/agent/api_utils.py b/platform/reworkd_platform/web/api/agent/api_utils.py index 0f037ea759..47765c929a 100644 --- a/platform/reworkd_platform/web/api/agent/api_utils.py +++ b/platform/reworkd_platform/web/api/agent/api_utils.py @@ -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 diff --git a/platform/reworkd_platform/web/api/agent/dependancies.py b/platform/reworkd_platform/web/api/agent/dependancies.py index fc8f63e1ff..83c897ab7f 100644 --- a/platform/reworkd_platform/web/api/agent/dependancies.py +++ b/platform/reworkd_platform/web/api/agent/dependancies.py @@ -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 diff --git a/platform/reworkd_platform/web/api/agent/helpers.py b/platform/reworkd_platform/web/api/agent/helpers.py index a06cbde299..519af721eb 100644 --- a/platform/reworkd_platform/web/api/agent/helpers.py +++ b/platform/reworkd_platform/web/api/agent/helpers.py @@ -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") diff --git a/platform/reworkd_platform/web/api/agent/model_settings.py b/platform/reworkd_platform/web/api/agent/model_settings.py index 89c3de76d3..878b49d985 100644 --- a/platform/reworkd_platform/web/api/agent/model_settings.py +++ b/platform/reworkd_platform/web/api/agent/model_settings.py @@ -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 diff --git a/platform/reworkd_platform/web/api/agent/views.py b/platform/reworkd_platform/web/api/agent/views.py index d3eca686f5..c8cb666cd2 100644 --- a/platform/reworkd_platform/web/api/agent/views.py +++ b/platform/reworkd_platform/web/api/agent/views.py @@ -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 diff --git a/platform/reworkd_platform/web/api/dependencies.py b/platform/reworkd_platform/web/api/dependencies.py index c0343089d2..99f6d457ba 100644 --- a/platform/reworkd_platform/web/api/dependencies.py +++ b/platform/reworkd_platform/web/api/dependencies.py @@ -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( diff --git a/platform/reworkd_platform/web/api/models/views.py b/platform/reworkd_platform/web/api/models/views.py index 91906c6adc..b2bbc6179f 100644 --- a/platform/reworkd_platform/web/api/models/views.py +++ b/platform/reworkd_platform/web/api/models/views.py @@ -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() diff --git a/platform/reworkd_platform/web/api/workflow/blocks/web/__init__.py b/platform/reworkd_platform/web/api/workflow/blocks/web/__init__.py index a0ffeb2e9c..68cdaf20a8 100644 --- a/platform/reworkd_platform/web/api/workflow/blocks/web/__init__.py +++ b/platform/reworkd_platform/web/api/workflow/blocks/web/__init__.py @@ -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: diff --git a/platform/reworkd_platform/web/api/workflow/views.py b/platform/reworkd_platform/web/api/workflow/views.py index e8ec192580..7e39088bb0 100644 --- a/platform/reworkd_platform/web/api/workflow/views.py +++ b/platform/reworkd_platform/web/api/workflow/views.py @@ -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()