From 846cab6b09e33b8a8bed0f4f5b8f795a98e4439c Mon Sep 17 00:00:00 2001
From: Mike Smith <mic-smith@users.noreply.github.com>
Date: Thu, 21 Nov 2024 16:47:41 +0000
Subject: [PATCH] PR comments - move directors

---
 backend/src/api/app.py                                        | 4 ++--
 backend/src/{director.py => directors/chat_director.py}       | 2 +-
 backend/src/{ => directors}/report_director.py                | 4 +++-
 .../src/prompts/templates/{director.j2 => chat_director.j2}   | 0
 backend/tests/api/app_test.py                                 | 2 +-
 backend/tests/{ => directors}/report_director_test.py         | 4 ++--
 6 files changed, 9 insertions(+), 7 deletions(-)
 rename backend/src/{director.py => directors/chat_director.py} (96%)
 rename backend/src/{ => directors}/report_director.py (85%)
 rename backend/src/prompts/templates/{director.j2 => chat_director.j2} (100%)
 rename backend/tests/{ => directors}/report_director_test.py (80%)

diff --git a/backend/src/api/app.py b/backend/src/api/app.py
index e2f4eab7..d01c0d4c 100644
--- a/backend/src/api/app.py
+++ b/backend/src/api/app.py
@@ -8,12 +8,12 @@
 from fastapi import FastAPI, HTTPException, Response, WebSocket, UploadFile
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
-from src.report_director import report_on_file_upload
+from src.directors.report_director import report_on_file_upload
 from src.session.file_uploads import clear_session_file_uploads
 from src.session.redis_session_middleware import reset_session
 from src.utils.graph_db_utils import populate_db
 from src.utils import Config, test_connection
-from src.director import question
+from src.directors.chat_director import question
 from src.websockets.connection_manager import connection_manager, parse_message
 from src.session import RedisSessionMiddleware
 from src.utils.cyper_import_data_from_csv import import_data_from_csv_script
diff --git a/backend/src/director.py b/backend/src/directors/chat_director.py
similarity index 96%
rename from backend/src/director.py
rename to backend/src/directors/chat_director.py
index c65be9ee..f3092a48 100644
--- a/backend/src/director.py
+++ b/backend/src/directors/chat_director.py
@@ -11,7 +11,7 @@
 logger = logging.getLogger(__name__)
 config = Config()
 engine = PromptEngine()
-director_prompt = engine.load_prompt("director")
+director_prompt = engine.load_prompt("chat_director")
 
 
 async def question(question: str) -> str:
diff --git a/backend/src/report_director.py b/backend/src/directors/report_director.py
similarity index 85%
rename from backend/src/report_director.py
rename to backend/src/directors/report_director.py
index 494dcc25..1d177b32 100644
--- a/backend/src/report_director.py
+++ b/backend/src/directors/report_director.py
@@ -2,7 +2,7 @@
 from typing import TypedDict
 from fastapi import UploadFile
 
-from src.utils.scratchpad import update_scratchpad
+from src.utils.scratchpad import clear_scratchpad, update_scratchpad
 from src.utils.file_utils import handle_file_upload
 
 class FileUploadReport(TypedDict):
@@ -18,4 +18,6 @@ async def report_on_file_upload(upload:UploadFile) -> FileUploadReport:
 
     report = "#Report on upload as markdown" # await report_agent.invoke(file["content"])
 
+    clear_scratchpad()
+
     return {"filename": file["filename"], "id": file["uploadId"], "report": report}
diff --git a/backend/src/prompts/templates/director.j2 b/backend/src/prompts/templates/chat_director.j2
similarity index 100%
rename from backend/src/prompts/templates/director.j2
rename to backend/src/prompts/templates/chat_director.j2
diff --git a/backend/tests/api/app_test.py b/backend/tests/api/app_test.py
index 92d88865..7207bcc0 100644
--- a/backend/tests/api/app_test.py
+++ b/backend/tests/api/app_test.py
@@ -1,6 +1,6 @@
 from fastapi.testclient import TestClient
 import pytest
-from src.report_director import FileUploadReport
+from src.directors.report_director import FileUploadReport
 from src.api import app, healthy_response, unhealthy_neo4j_response, chat_fail_response
 
 client = TestClient(app)
diff --git a/backend/tests/report_director_test.py b/backend/tests/directors/report_director_test.py
similarity index 80%
rename from backend/tests/report_director_test.py
rename to backend/tests/directors/report_director_test.py
index 269aee72..32f90770 100644
--- a/backend/tests/report_director_test.py
+++ b/backend/tests/directors/report_director_test.py
@@ -4,14 +4,14 @@
 import pytest
 
 from src.session.file_uploads import FileUpload
-from src.report_director import report_on_file_upload
+from src.directors.report_director import report_on_file_upload
 
 @pytest.mark.asyncio
 async def test_report_on_file_upload(mocker):
 
     file_upload = FileUpload(uploadId="1", filename="test.txt", content="test", contentType="text/plain", size=4)
 
-    mock_handle_file_upload = mocker.patch("src.report_director.handle_file_upload", return_value=file_upload)
+    mock_handle_file_upload = mocker.patch("src.directors.report_director.handle_file_upload", return_value=file_upload)
 
     headers = Headers({"content-type": "text/plain"})
     file = BytesIO(b"test content")