Merge pull request #15 from airtai/implement-file-upload

Implement file upload
airtai · Nov 22, 2024 · ace73d9 · ace73d9
2 parents 3f22fa7 + 41a9468
commit ace73d9
Show file tree

Hide file tree

Showing 12 changed files with 286 additions and 43 deletions.
diff --git a/.devcontainer/setup.sh b/.devcontainer/setup.sh
@@ -15,3 +15,9 @@ if [ -z "$OPENAI_API_KEY" ]; then
     echo -e "\033[33mWarning: OPENAI_API_KEY environment variable is not set.\033[0m"
     echo
 fi
+# check MAILCHIMP_API_KEY environment variable is set
+if [ -z "$MAILCHIMP_API_KEY" ]; then
+    echo
+    echo -e "\033[33mWarning: MAILCHIMP_API_KEY environment variable is not set.\033[0m"
+    echo
+fi
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -47,3 +47,4 @@ jobs:
         run: pytest
         env:
           CONTEXT: ${{ runner.os }}-py${{ matrix.python-version }}
+          MAILCHIMP_API_KEY: "test-key" # pragma: allowlist secret
diff --git a/.gitignore b/.gitignore
@@ -18,3 +18,4 @@ token
 .DS_Store
 
 tmp_*
+mailchimp_api/uploaded_files
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -38,4 +38,4 @@ CMD ["/app/run_fastagency.sh"]
 
 # Run the container
 
-# docker run --rm -d --name deploy_fastagency -e OPENAI_API_KEY=$OPENAI_API_KEY  -p 8008:8008 -p 8888:8888  deploy_fastagency
+# docker run --rm -d --name deploy_fastagency -e OPENAI_API_KEY=$OPENAI_API_KEY -e MAILCHIMP_API_KEY=$MAILCHIMP_API_KEY -p 8008:8008 -p 8888:8888  deploy_fastagency
diff --git a/mailchimp_api/constants.py b/mailchimp_api/constants.py
@@ -0,0 +1,3 @@
+from pathlib import Path
+
+UPLOADED_FILES_DIR = Path(__file__).parent / "uploaded_files"
diff --git a/mailchimp_api/deployment/main_1_fastapi.py b/mailchimp_api/deployment/main_1_fastapi.py
@@ -1,8 +1,12 @@
+from pathlib import Path
 from typing import Any
 
+import pandas as pd
 from fastagency.adapters.fastapi import FastAPIAdapter
-from fastapi import FastAPI
+from fastapi import FastAPI, Form, HTTPException, Query, UploadFile, status
+from fastapi.responses import HTMLResponse
 
+from ..constants import UPLOADED_FILES_DIR
 from ..workflow import wf
 
 adapter = FastAPIAdapter(provider=wf)
@@ -17,5 +21,71 @@ def list_workflows() -> dict[str, Any]:
     return {"Workflows": {name: wf.get_description(name) for name in wf.names}}
 
 
+def _save_file(file: UploadFile, timestamp: str) -> Path:
+    UPLOADED_FILES_DIR.mkdir(exist_ok=True)
+    try:
+        contents = file.file.read()
+        file_name = f"uploaded-file-{timestamp}.csv"
+        path = UPLOADED_FILES_DIR / file_name
+        with path.open("wb") as f:
+            f.write(contents)
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="There was an error uploading the file",
+        ) from e
+    finally:
+        file.file.close()
+
+    return path
+
+
+@app.post("/upload")
+def upload(
+    file: UploadFile = UploadFile(...),  # type: ignore[arg-type] # noqa: B008
+    timestamp: str = Form(...),
+) -> dict[str, str]:
+    if not file.size:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Please provide .csv file",
+        )
+    if file.content_type != "text/csv":
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Only CSV files are supported",
+        )
+
+    path = _save_file(file, timestamp)
+    df = pd.read_csv(path)
+    if "email" not in df.columns:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="'email' column not found in CSV file",
+        )
+
+    return {
+        "message": f"Successfully uploaded {file.filename}. Please close the tab and go back to the chat."
+    }
+
+
+@app.get("/upload-file")
+def upload_file(timestamp: str = Query(default="default-timestamp")) -> HTMLResponse:
+    content = f"""<body>
+    <form action='/upload' enctype='multipart/form-data' method='post'>
+        <div style="margin-top: 15px;">
+            <input name='file' type='file'>
+        </div>
+        <!-- Hidden field for timestamp -->
+        <input name='timestamp' type='hidden' value='{timestamp}'>
+        <div style="margin-top: 15px;">
+            <input type='submit'>
+        </div>
+    </form>
+</body>
+"""
+    return HTMLResponse(content=content)
+
+
 # start the adapter with the following command
-# uvicorn mailchimp_api.deployment.main_1_fastapi:app --reload
+# uvicorn mailchimp_api.deployment.main_1_fastapi:app -b 0.0.0.0:8008 --reload
diff --git a/mailchimp_api/workflow.py b/mailchimp_api/workflow.py
@@ -1,47 +1,86 @@
 import os
+import time
 from typing import Any
 
-from autogen.agentchat import ConversableAgent
+import pandas as pd
 from fastagency import UI
 from fastagency.runtimes.autogen import AutoGenWorkflows
 
-llm_config = {
-    "config_list": [
-        {
-            "model": "gpt-4o-mini",
-            "api_key": os.getenv("OPENAI_API_KEY"),
-        }
-    ],
-    "temperature": 0.8,
-}
+from .config import Config
+from .constants import UPLOADED_FILES_DIR
+from .processing.update_tags import update_tags
 
 wf = AutoGenWorkflows()
 
+FASTAPI_URL = os.getenv("FASTAPI_URL", "http://localhost:8008")
 
-@wf.register(name="simple_learning", description="Student and teacher learning chat")  # type: ignore[misc]
-def simple_workflow(ui: UI, params: dict[str, Any]) -> str:
-    initial_message = ui.text_input(
+
+def _get_config() -> Config:
+    api_key = os.getenv("MAILCHIMP_API_KEY")
+    if not api_key:
+        raise ValueError("MAILCHIMP_API_KEY not set")
+
+    config = Config("us14", api_key)
+    return config
+
+
+config = _get_config()
+
+
+def _wait_for_file(timestamp: str) -> pd.DataFrame:
+    file_name = f"uploaded-file-{timestamp}.csv"
+    file_path = UPLOADED_FILES_DIR / file_name
+    while not file_path.exists():
+        time.sleep(2)
+
+    df = pd.read_csv(file_path)
+    file_path.unlink()
+
+    return df
+
+
+@wf.register(name="mailchimp_chat", description="Mailchimp tags update chat")  # type: ignore[misc]
+def mailchimp_chat(ui: UI, params: dict[str, Any]) -> str:
+    timestamp = time.strftime("%Y-%m-%d-%H-%M-%S")
+    body = f"""Please upload **.csv** file with the email addresses for which you want to update the tags.
+
+<a href="{FASTAPI_URL}/upload-file?timestamp={timestamp}" target="_blank">Upload File</a>
+"""
+    ui.text_message(
         sender="Workflow",
         recipient="User",
-        prompt="I can help you learn about mathematics. What subject you would like to explore?",
+        body=body,
     )
 
-    student_agent = ConversableAgent(
-        name="Student_Agent",
-        system_message="You are a student willing to learn.",
-        llm_config=llm_config,
-    )
-    teacher_agent = ConversableAgent(
-        name="Teacher_Agent",
-        system_message="You are a math teacher.",
-        llm_config=llm_config,
+    df = _wait_for_file(timestamp)
+
+    list_name = None
+    while list_name is None:
+        list_name = ui.text_input(
+            sender="Workflow",
+            recipient="User",
+            prompt="Please enter Account Name for which you want to update the tags",
+        )
+
+    add_tag_members, _ = update_tags(
+        crm_df=df, config=config, list_name=list_name.strip()
     )
+    if not add_tag_members:
+        return "No tags added"
 
-    chat_result = student_agent.initiate_chat(
-        teacher_agent,
-        message=initial_message,
-        summary_method="reflection_with_llm",
-        max_turns=3,
+    add_tag_members = dict(sorted(add_tag_members.items()))
+    updates_per_tag = "\n".join(
+        [f"- **{key}**: {len(value)}" for key, value in add_tag_members.items()]
     )
+    body = f"""Number of updates per tag:
 
-    return str(chat_result.summary)
+{updates_per_tag}
+
+(It might take some time for updates to reflect in Mailchimp)
+"""
+    ui.text_message(
+        sender="Workflow",
+        recipient="User",
+        body=body,
+    )
+    return "Task Completed"
diff --git a/pyproject.toml b/pyproject.toml
@@ -9,6 +9,7 @@ name = "mailchimp_api"
 
 dependencies = [
     "fastagency[autogen,mesop,openapi,server,fastapi]>=0.3.0",
+    "python-multipart>=0.0.17",
     "pandas>=2.2.3",
     "tenacity>=9.0.0",
 ]

diff --git a/scripts/deploy_to_fly_io.sh b/scripts/deploy_to_fly_io.sh
@@ -8,3 +8,4 @@ fly launch --config fly.toml --copy-config --yes
 
 echo -e "\033[0;32mSetting secrets\033[0m"
 fly secrets set OPENAI_API_KEY=$OPENAI_API_KEY
+fly secrets set MAILCHIMP_API_KEY=$MAILCHIMP_API_KEY
diff --git a/scripts/run_docker.sh b/scripts/run_docker.sh
@@ -1,3 +1,3 @@
 #!/bin/bash
 
-docker run -it -e OPENAI_API_KEY=$OPENAI_API_KEY   -p 8008:8008 -p 8888:8888  deploy_fastagency
+docker run -it -e OPENAI_API_KEY=$OPENAI_API_KEY -e MAILCHIMP_API_KEY=$MAILCHIMP_API_KEY -p 8008:8008 -p 8888:8888  deploy_fastagency
diff --git a/tests/deployment/test_main_1_fastapi.py b/tests/deployment/test_main_1_fastapi.py
@@ -0,0 +1,87 @@
+from io import BytesIO
+from pathlib import Path
+
+import pandas as pd
+import pytest
+from _pytest.monkeypatch import MonkeyPatch
+from fastapi import UploadFile
+from fastapi.testclient import TestClient
+
+from mailchimp_api.deployment.main_1_fastapi import _save_file, app
+
+
+class TestApp:
+    client = TestClient(app)
+
+    @pytest.fixture(autouse=True)
+    def patch_uploaded_files_dir(
+        self, tmp_path: Path, monkeypatch: MonkeyPatch
+    ) -> Path:
+        uploaded_files_dir = tmp_path / "uploads"
+        uploaded_files_dir.mkdir(exist_ok=True)
+
+        monkeypatch.setattr(
+            "mailchimp_api.deployment.main_1_fastapi.UPLOADED_FILES_DIR",
+            uploaded_files_dir,
+        )
+
+        # Return the temporary directory so it can be used in tests if needed
+        return uploaded_files_dir
+
+    def test_save_file(self) -> None:
+        csv_content = "email\n[email protected]\n[email protected]"
+        csv_file = BytesIO(csv_content.encode("utf-8"))
+        uploaded_file = UploadFile(filename="emails.csv", file=csv_file)
+        path = _save_file(uploaded_file, "22-09-2021")
+        df = pd.read_csv(path)
+
+        expected_df = pd.DataFrame(
+            {"email": ["[email protected]", "[email protected]"]}
+        )
+        assert df.equals(expected_df)
+
+    def test_upload_file_endpoint(self) -> None:
+        timestamp = "22-09-2021"
+        response = self.client.get(f"/upload-file?timestamp={timestamp}")
+        assert response.status_code == 200
+        assert timestamp in response.text
+
+    def test_upload_endpoint(self) -> None:
+        csv_content = "email\n[email protected]\n"
+        csv_file = BytesIO(csv_content.encode("utf-8"))
+
+        response = self.client.post(
+            "/upload",
+            files={"file": ("emails.csv", csv_file)},
+            data={"timestamp": "test-22-09-2021"},
+        )
+        assert response.status_code == 200
+        expected_msg = "Successfully uploaded emails.csv. Please close the tab and go back to the chat."
+        assert expected_msg == response.json()["message"]
+
+    def test_upload_endpoint_raises_400_error_if_file_isnt_provided(self) -> None:
+        response = self.client.post("/upload", data={"timestamp": "test-22-09-2021"})
+        assert response.status_code == 400
+        assert "Please provide .csv file" in response.text
+
+    def test_upload_endpoint_raises_400_error_if_file_is_not_csv(self) -> None:
+        csv_content = "email\n"
+        csv_file = BytesIO(csv_content.encode("utf-8"))
+        response = self.client.post(
+            "/upload",
+            files={"file": ("emails.txt", csv_file)},
+            data={"timestamp": "test-22-09-2021"},
+        )
+        assert response.status_code == 400
+        assert "Only CSV files are supported" in response.text
+
+    def test_upload_endpoint_raises_400_error_if_email_column_not_found(self) -> None:
+        csv_content = "name\n"
+        csv_file = BytesIO(csv_content.encode("utf-8"))
+        response = self.client.post(
+            "/upload",
+            files={"file": ("emails.csv", csv_file)},
+            data={"timestamp": "test-22-09-2021"},
+        )
+        assert response.status_code == 400
+        assert "'email' column not found in CSV file" in response.text
Original file line number	Diff line number	Diff line change
Expand Up		@@ -18,3 +18,4 @@ token
		.DS_Store

		tmp_*
		mailchimp_api/uploaded_files
Original file line number	Diff line number	Diff line change
Expand Up		@@ -38,4 +38,4 @@ CMD ["/app/run_fastagency.sh"]

		# Run the container

		# docker run --rm -d --name deploy_fastagency -e OPENAI_API_KEY=$OPENAI_API_KEY -p 8008:8008 -p 8888:8888 deploy_fastagency
		# docker run --rm -d --name deploy_fastagency -e OPENAI_API_KEY=$OPENAI_API_KEY -e MAILCHIMP_API_KEY=$MAILCHIMP_API_KEY -p 8008:8008 -p 8888:8888 deploy_fastagency
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		from pathlib import Path

		UPLOADED_FILES_DIR = Path(__file__).parent / "uploaded_files"
Original file line number	Diff line number	Diff line change
Expand Up		@@ -8,3 +8,4 @@ fly launch --config fly.toml --copy-config --yes

		echo -e "\033[0;32mSetting secrets\033[0m"
		fly secrets set OPENAI_API_KEY=$OPENAI_API_KEY
		fly secrets set MAILCHIMP_API_KEY=$MAILCHIMP_API_KEY