Assistant drive library

Co-authored-by: Brian Krabach <brian@krabach.com>

.devcontainer/devcontainer.json

@@ -57,7 +57,6 @@
         "epivision.vscode-file-header",
         "esbenp.prettier-vscode",
         "github.vscode-github-actions",
-        "matt-rudge.auto-open-preview-panel",
         "ms-azuretools.vscode-docker",
         "ms-python.debugpy",
         "ms-python.python",

.vscode/settings.json

@@ -219,5 +219,6 @@
     "winget",
     "workbenchservice"
   ],
-  "markdown.validate.enabled": true
+  "markdown.validate.enabled": true,
+  "makefile.configureOnOpen": false
 }

assistants/prospector-assistant/assistant/agents/attachment_agent.py

@@ -1,11 +1,12 @@
 import base64
 import io
 import logging
-from pathlib import Path
 from typing import Annotated, Any
 
 import docx2txt
 import pdfplumber
+from assistant_drive import Drive, DriveConfig
+from context import Context
 from openai.types import chat
 from pydantic import BaseModel, Field
 from semantic_workbench_api_model.workbench_model import File
@@ -14,11 +15,6 @@
     FileStorageContext,
 )
 from semantic_workbench_assistant.config import UISchema
-from semantic_workbench_assistant.storage import (
-    read_model,
-    read_models_in_dir,
-    write_model,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -84,15 +80,16 @@ async def create_or_update_attachment_from_file(
         content = await _file_to_str(context, file)
 
         # see if there is already an attachment with this filename
-        attachment = read_model(_get_attachment_storage_path(context, filename), Attachment)
-        if attachment:
+        try:
+            attachment = _get_attachment_drive(context).read_model(Attachment, filename)
             # if there is, update the content
             attachment.content = content
-        else:
+        except FileNotFoundError:
             # if there isn't, create a new attachment
             attachment = Attachment(filename=filename, content=content, metadata=metadata)
 
-        write_model(_get_attachment_storage_path(context, filename), attachment)
+        # write the attachment to the storage
+        _get_attachment_drive(context).write_model(attachment, filename)
 
     @staticmethod
     def delete_attachment_for_file(context: ConversationContext, file: File) -> None:
@@ -101,7 +98,7 @@ def delete_attachment_for_file(context: ConversationContext, file: File) -> None
         """
 
         filename = file.filename
-        _get_attachment_storage_path(context, filename).unlink(missing_ok=True)
+        _get_attachment_drive(context).delete(filename)
 
     @staticmethod
     def generate_attachment_messages(
@@ -120,7 +117,7 @@ def generate_attachment_messages(
         """
 
         # get all attachments and exit early if there are none
-        attachments = read_models_in_dir(_get_attachment_storage_path(context), Attachment)
+        attachments = _get_attachment_drive(context).read_models(Attachment)
         if not attachments:
             return []
 
@@ -233,14 +230,13 @@ def reduce_attachment_payload_from_content(value: Any) -> Any:
 #
 
 
-def _get_attachment_storage_path(context: ConversationContext, filename: str | None = None) -> Path:
+def _get_attachment_drive(context: ConversationContext) -> Drive:
     """
-    Get the path where attachments are stored.
+    Get the Drive instance for the attachments.
     """
-    path = FileStorageContext.get(context).directory / "attachments"
-    if filename:
-        path /= filename
-    return path
+    drive_context = Context(session_id=context.id)
+    drive_root = str(FileStorageContext.get(context).directory / "attachments")
+    return Drive(DriveConfig(context=drive_context, root=drive_root))
 
 
 async def _raw_content_from_file(context: ConversationContext, file: File) -> bytes:

assistants/prospector-assistant/pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "document-skill>=0.1.0",
     "guided-conversation>=0.1.0",
     "openai-client>=0.1.0",
+    "assistant-drive>=0.1.0",
 ]
 
 [tool.uv]
@@ -37,6 +38,7 @@ posix-skill = { path = "../../libraries/python/skills/skills/posix-skill", edita
 prospector-skill = { path = "../../libraries/python/skills/skills/prospector-skill", editable = true }
 document-skill = { path = "../../libraries/python/skills/skills/document-skill", editable = true }
 openai-client = { path = "../../libraries/python/openai-client", editable = true }
+assistant-drive = { path = "../../libraries/python/assistant-drive", editable = true }
 
 [build-system]
 requires = ["hatchling"]

libraries/python/assistant-drive/.gitignore

@@ -0,0 +1,9 @@
+__pycache__
+.pytest_cache
+*.egg*
+.data
+.venv
+venv
+.env
+
+poetry.lock

libraries/python/assistant-drive/.vscode/extensions.json

@@ -0,0 +1,7 @@
+{
+  "recommendations": [
+    "aarontamasfe.even-better-toml",
+    "charliermarsh.ruff",
+    "ms-python.python"
+  ]
+}

libraries/python/assistant-drive/.vscode/settings.json

@@ -0,0 +1,70 @@
+{
+  "editor.bracketPairColorization.enabled": true,
+  "editor.codeActionsOnSave": {
+    "source.fixAll": "always",
+    "source.organizeImports": "always"
+  },
+  "editor.defaultFormatter": "esbenp.prettier-vscode",
+  "editor.formatOnPaste": true,
+  "editor.formatOnSave": true,
+  "editor.formatOnType": true,
+  "editor.guides.bracketPairs": "active",
+  "files.eol": "\n",
+  "files.trimTrailingWhitespace": true,
+  "flake8.ignorePatterns": [
+    "**/*.py"
+  ], // disable flake8 in favor of ruff
+  "jupyter.debugJustMyCode": false,
+  "python.analysis.autoFormatStrings": true,
+  "python.analysis.autoImportCompletions": true,
+  "python.analysis.diagnosticMode": "workspace",
+  "python.analysis.exclude": [
+    "**/.venv/**",
+    "**/.data/**",
+    "**/__pycache__/**"
+  ],
+  "python.analysis.fixAll": [
+    "source.unusedImports"
+  ],
+  "python.analysis.inlayHints.functionReturnTypes": true,
+  "python.analysis.typeCheckingMode": "basic",
+  "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python",
+  "python.testing.cwd": "${workspaceFolder}",
+  "python.testing.pytestArgs": [
+    "--color",
+    "yes"
+  ],
+  "python.testing.pytestEnabled": true,
+  "search.exclude": {
+    "**/.venv": true,
+    "**/data": true
+  },
+  "[json]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode",
+    "editor.formatOnSave": true
+  },
+  "[jsonc]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode",
+    "editor.formatOnSave": true
+  },
+  "[python]": {
+    "editor.defaultFormatter": "charliermarsh.ruff",
+    "editor.formatOnSave": true,
+    "editor.codeActionsOnSave": {
+      "source.fixAll": "explicit",
+      "source.unusedImports": "explicit",
+      "source.organizeImports": "explicit",
+      "source.formatDocument": "explicit"
+    }
+  },
+  "ruff.nativeServer": "on",
+  "cSpell.words": [
+    "dotenv",
+    "httpx",
+    "openai",
+    "pydantic",
+    "pypdf",
+    "runtimes",
+    "tiktoken"
+  ]
+}

libraries/python/assistant-drive/Makefile

@@ -0,0 +1,3 @@
+repo_root = $(shell git rev-parse --show-toplevel)
+include $(repo_root)/tools/makefiles/python.mk
+include $(repo_root)/tools/makefiles/docker-assistant.mk

libraries/python/assistant-drive/README.md

@@ -0,0 +1,3 @@
+# Assistant Drive
+
+These are file storage capabilities.

libraries/python/assistant-drive/assistant_drive/__init__.py

@@ -0,0 +1,6 @@
+from .drive import Drive, DriveConfig
+
+__all__ = [
+    "Drive",
+    "DriveConfig",
+]