Refactors custom apps (#18)

Author: leokster

README.rst

@@ -64,7 +64,7 @@ Once you have Docker, Docker Compose, and Python3 installed, you can download an
 
     git clone https://github.com/modulos/data_copilot.git
     cd data_copilot
-    pip install -r requirements_dev.txt
+    pip install ".[dev]"
     make setup
 
 **Open Data Copilot in your browser: http://localhost:80**
@@ -210,18 +210,14 @@ This command will start the Data Copilot service in development mode. Now, whene
 
 Data Copilot is not just a standalone application, but also a framework that you can use to build your own data processing and analysis tools. Here are the steps to get started:
 
-1. **Worker Logic:** The worker logic can be found in the `celery_app/apps` directory. You can modify the logic here to suit your specific needs.
-
-2. **Getting Started Example:** For a basic understanding of the worker logic, you can refer to the `celery_app/apps/getting_started_example.py` file. This file provides a simple example that can serve as a starting point for your custom logic.
-
-3. **Executor Logic:** The executor logic is contained in the `celery_app/executors/getting_started_executor.py` file. You can modify this file to customize how tasks are executed.
-
-4. **Supported File Types:** If you want to change the supported file types (e.g., extend support to PDF), you will need to configure this on the backend side in the `backend/config/config.py` file. Additionally, you need to implement the logic for handling the new file type in the `backend/routers/artifacts.py` file.
-
-5. **File Type Interaction:** Once you've configured the backend to support the new file type, you'll need to implement the specific logic for interacting with that file type on the worker side.
-
-6. **Return Types:** Currently, Data Copilot is configured to only return tables to the user. However, the framework supports other return types such as heatmaps, histograms, and barplots. You can see the implementation details for these types in the `getting_started_executor.py` file.
+1. **Logic:** All the logic of your app should be in the `data_copilot/execution_apps/apps` directory. You can modify the logic here to suit your specific needs. You can inherit from the `data_copilot.exection_apps.base.DataCopilotApp` class. You need to implement at least the three static methods:
+   - `supported_file_types`: This method should return a dict of the supported file types. The keys should be the identifier and the value the content-type of the file type.
+   - `process_data_upload`: This method gets a list of the FastAPI `UploadFile` objects and should return a list of dict where the key is the file name and the value the content of the file as BufferedIOBase
+   - `execute_message`: This method contnains the execution logic of your app, which gets executed on the worker.
 
+2. **Message** The execution_message should return a `data_copilot.execution_apps.helpers.Message` object. 
+   
+  
 With these steps, you can customize Data Copilot to handle your specific data processing and analysis tasks. Remember to thoroughly test your changes to ensure they work as expected.
 
 

data_copilot/backend/config/config.py

@@ -19,13 +19,11 @@ class Config(BaseSettings):
         "application/vnd.ms-excel": "xls",
         "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx",
     }
+    COMPUTE_BACKEND: str = Field(..., env="COMPUTE_BACKEND")
 
     if os.getenv("ENVIRONMENT") != "TEST":
         DB_CONNECTION_STRING: str = Field(..., env="DB_CONNECTION_STRING")
         SECRET_KEY: str = Field(..., env="JWT_SECRET_KEY")
-        AZURE_STORAGE_ACCOUNT_NAME: str = Field(..., env="AZURE_STORAGE_ACCOUNT_NAME")
-        AZURE_STORAGE_ACCOUNT_KEY: str = Field(..., env="AZURE_STORAGE_ACCOUNT_KEY")
-        CONTAINER_NAME: str = Field(..., env="CONTAINER_NAME")
         STORAGE_BACKEND: str = Field(..., env="STORAGE_BACKEND")
         BACKEND_HOST: str = Field(..., env="BACKEND_HOST")
 
@@ -35,9 +33,14 @@ class Config(BaseSettings):
         secret_key = "c86d3444a380bb36cca7abe1b6dcc8caaee0ecf5bbe254c5473783d147ebc12e"
         SECRET_KEY: str = secret_key
         BACKEND_HOST: str = "http://localhost:8000/api"
-        STORAGE_BACKEND: str = "./artifacts/"
+        STORAGE_BACKEND: str = "file://./artifacts/"
         CELERY_BROKER_URL: str = ""
 
+    if "dfs.core.windows.net" in os.getenv("STORAGE_BACKEND"):
+        AZURE_STORAGE_ACCOUNT_NAME: str = Field(..., env="AZURE_STORAGE_ACCOUNT_NAME")
+        AZURE_STORAGE_ACCOUNT_KEY: str = Field(..., env="AZURE_STORAGE_ACCOUNT_KEY")
+        CONTAINER_NAME: str = Field(..., env="CONTAINER_NAME")
+
     @property
     def POSTGRES_CONNECTION(self):
         return self.DB_CONNECTION_STRING

data_copilot/backend/routers/artifacts.py

@@ -1,9 +1,6 @@
-import json
-import logging
 import os
 from typing import Optional
 
-import pandas as pd
 from fastapi import Depends, HTTPException, UploadFile, routing
 
 from data_copilot.backend.artifacts.artifact import CreateArtifactVersionCM
@@ -37,6 +34,8 @@
 )
 from data_copilot.backend.schemas.authentication import User
 from data_copilot.storage_handler import get_signed_download_url, list_files
+from data_copilot.execution_apps import get_app
+
 
 artifacts_router = routing.APIRouter(prefix="/artifacts")
 
@@ -321,45 +320,15 @@ async def post_artifacts_id_artifactid_versions(
             status_code=400, detail="Only dataset artifact versions can be uploaded"
         )
 
-    if file.content_type not in CONFIG.ALLOWED_ARTIFACTS_CONTENT_TYPES.keys():
+    if file.content_type not in get_app().supported_file_types.keys():
         raise routing.HTTPException(
             status_code=415, detail=f"File format '{file.content_type}' not allowed"
         )
 
-    try:
-        file_type = CONFIG.ALLOWED_ARTIFACTS_CONTENT_TYPES[file.content_type]
-        match file_type:
-            case "csv":
-                data_frame = pd.read_csv(file.file, sep=None, encoding="utf-8-sig")
-            case "xls" | "xlsx":
-                data_frame = pd.read_excel(await file.read(), dtype={"dteday": str})
-
-        if data_frame.empty:
-            raise routing.HTTPException(
-                status_code=400, detail=f"Wrong '{file.file}'content"
-            )
-
-    except Exception as e:
-        logging.error(e)
-        raise routing.HTTPException(
-            status_code=500, detail=f"Loading '{file.content_type}' failed"
-        )
-    schema = {col: str(data_frame[col].dtype) for col in data_frame.columns}
-    file_config = {
-        "file_name": file.filename,
-        "file_type": file_type,
-        "file_schema": schema,
-        "rows": data_frame.shape[0],
-    }
-    file.file.seek(0)
     with CreateArtifactVersionCM(artifact.id, db) as cm:
-        cm.write(file.filename, file.file)
-        artifact_version_config = {
-            "artifact_id": str(artifact.id),
-            "artifact_version_id": str(cm.uuid),
-            "files": [file_config],
-        }
-        cm.write("config.json", json.dumps(artifact_version_config, indent=4))
+        files = get_app().process_data_upload([file], artifact, cm)
+        for file in files:
+            cm.write(*file)
         uuid = cm.uuid
 
     return crud_get_artifact_version(db, uuid)

data_copilot/backend/routers/chats.py

@@ -1,6 +1,4 @@
 import asyncio
-import json
-import os
 import time
 
 from fastapi import Depends, HTTPException, routing
@@ -46,11 +44,6 @@
     UpdateChat,
     parse_chat_message,
 )
-from data_copilot.storage_handler.functions import (
-    exists,
-    get_signed_download_url,
-    read_file,
-)
 
 CONFIG = Config()
 
@@ -293,7 +286,6 @@ async def post_chats_chatid_messages_messageid(
         message (Message): The message to be executed
 
     """
-    artifact_config = {}
     if message.artifact_version_id:
         artifact_version: ArtifactVersion = await get_artifact_version_dependency(
             message.artifact_version_id, db=db
@@ -304,39 +296,30 @@ async def post_chats_chatid_messages_messageid(
         await check_if_user_has_access_to_artifact(artifact, current_user)
         await check_if_artifact_is_active(artifact)
 
-        if not exists(os.path.join(artifact_version.artifact_uri, "config.json")):
-            raise HTTPException(
-                status_code=400,
-                detail="The artifact version must contain a config.json file",
-            )
+    artifact_version_uri = artifact_version.artifact_uri if artifact_version else None
 
-        artifact_config = json.load(
-            read_file(os.path.join(artifact_version.artifact_uri, "config.json"))
-        )
-        file_name = artifact_config.get("files", [dict()])[0].get("file_name", None)
-        if not exists(os.path.join(artifact_version.artifact_uri, file_name)):
-            raise HTTPException(
-                status_code=400,
-                detail=f"The artifact version must contain a file named {file_name}",
-            )
+    previous_messages = crud_get_messages_by_chat_id_sorted_desc(
+        db,
+        chat_id=message.chat_id,
+        limit=10,
+        offset=0,
+    )
 
-        sas_url = get_signed_download_url(
-            os.path.join(artifact_version.artifact_uri, file_name)
-        )
-        artifact_version_id = artifact_version.id
-    else:
-        sas_url = None
-        artifact_version_id = None
+    # make json from previous messages
+    previous_messages = [
+        parse_chat_message(Message.from_orm(message)).dict()
+        for message in previous_messages
+    ]
 
     execution_app.send_task(
         "execute_user_message",
         args=(
             message.content,
             message.chat_id,
             message.id,
-            artifact_version_id,
-            sas_url,
-            artifact_config,
+            artifact_version.id,
+            artifact_version_uri,
+            previous_messages,
         ),
     )