Add document upload support and update prompt

README.md

@@ -6,7 +6,8 @@ database using Peewee. Histories are persisted per user and session so
 conversations can be resumed with context. One example tool is included:
 
 * **execute_terminal** – Executes a shell command inside a persistent Linux VM
-  with network access. Output from ``stdout`` and ``stderr`` is captured and
+  with network access. Use it to read uploaded documents under ``/data`` or run
+  other commands. Output from ``stdout`` and ``stderr`` is captured and
   returned. The VM is created when a chat session starts and reused for all
   subsequent tool calls.
 
@@ -22,6 +23,14 @@ python run.py
 
 The script will instruct the model to run a simple shell command and print the result. Conversations are automatically persisted to `chat.db` and are now associated with a user and session.
 
+Uploaded files are stored under the `uploads` directory and mounted inside the VM at `/data`. Call ``upload_document`` on the chat session to make a file available to the model:
+
+```python
+async with ChatSession() as chat:
+    path_in_vm = chat.upload_document("path/to/file.pdf")
+    reply = await chat.chat(f"Summarize {path_in_vm}")
+```
+
 ## Docker
 
 A Dockerfile is provided to run the Discord bot along with an Ollama server. The image installs Ollama, pulls the LLM and embedding models, and starts both the server and the bot.

run.py

@@ -7,9 +7,8 @@ from src.chat import ChatSession
 
 async def _main() -> None:
     async with ChatSession(user="demo_user", session="demo_session") as chat:
-        answer = await chat.chat(
-            "check the processes running on the system"
-        )
+        doc_path = chat.upload_document("README.md")
+        answer = await chat.chat(f"List the first three lines of {doc_path}")
         print("\n>>>", answer)
 
 

src/chat.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 from typing import List
 import json
+import shutil
+from pathlib import Path
 
 from ollama import AsyncClient, ChatResponse, Message
 
@@ -12,8 +14,16 @@ from .config import (
     NUM_CTX,
     OLLAMA_HOST,
     SYSTEM_PROMPT,
+    UPLOAD_DIR,
+)
+from .db import (
+    Conversation,
+    Message as DBMessage,
+    User,
+    _db,
+    init_db,
+    add_document,
 )
-from .db import Conversation, Message as DBMessage, User, _db, init_db
 from .log import get_logger
 from .schema import Msg
 from .tools import execute_terminal, set_vm
@@ -53,6 +63,24 @@ class ChatSession:
         if not _db.is_closed():
             _db.close()
 
+    def upload_document(self, file_path: str) -> str:
+        """Save a document for later access inside the VM.
+
+        The file is copied into ``UPLOAD_DIR`` and recorded in the database. The
+        returned path is the location inside the VM (prefixed with ``/data``).
+        """
+
+        src = Path(file_path)
+        if not src.exists():
+            raise FileNotFoundError(file_path)
+
+        dest = Path(UPLOAD_DIR) / self._user.username
+        dest.mkdir(parents=True, exist_ok=True)
+        target = dest / src.name
+        shutil.copy(src, target)
+        add_document(self._user.username, str(target), src.name)
+        return f"/data/{self._user.username}/{src.name}"
+
     def _ensure_system_prompt(self) -> None:
         if any(m.get("role") == "system" for m in self._messages):
             return
@@ -84,9 +112,7 @@ class ChatSession:
         return messages
 
     @staticmethod
-    def _store_assistant_message(
-        conversation: Conversation, message: Message
-    ) -> None:
+    def _store_assistant_message(conversation: Conversation, message: Message) -> None:
         """Persist assistant messages, storing tool calls when present."""
 
         if message.tool_calls:
@@ -135,9 +161,7 @@ class ChatSession:
             )
             nxt = await self.ask(messages, think=True)
             self._store_assistant_message(conversation, nxt.message)
-            return await self._handle_tool_calls(
-                messages, nxt, conversation, depth + 1
-            )
+            return await self._handle_tool_calls(messages, nxt, conversation, depth + 1)
 
         return response
 

src/config.py

@@ -1,18 +1,23 @@
 from __future__ import annotations
 
 import os
+from pathlib import Path
 from typing import Final
 
 MODEL_NAME: Final[str] = os.getenv("OLLAMA_MODEL", "qwen3:1.7b")
-EMBEDDING_MODEL_NAME: Final[str] = os.getenv("OLLAMA_EMBEDDING_MODEL", "snowflake-arctic-embed:137m")   # unused for now
+EMBEDDING_MODEL_NAME: Final[str] = os.getenv(
+    "OLLAMA_EMBEDDING_MODEL", "snowflake-arctic-embed:137m"
+)  # unused for now
 OLLAMA_HOST: Final[str] = os.getenv("OLLAMA_HOST", "http://localhost:11434")
 MAX_TOOL_CALL_DEPTH: Final[int] = 5
 NUM_CTX: Final[int] = int(os.getenv("OLLAMA_NUM_CTX", "16000"))
+UPLOAD_DIR: Final[str] = os.getenv("UPLOAD_DIR", str(Path.cwd() / "uploads"))
 
 SYSTEM_PROMPT: Final[str] = (
     "You are a versatile AI assistant named Starlette able to orchestrate several tools to "
     "complete tasks. Plan your responses carefully and, when needed, call one "
     "or more tools consecutively to gather data, compute answers, or transform "
-    "information. Continue chaining tools until the user's request is fully "
-    "addressed and then deliver a concise, coherent final reply."
+    "information. Uploaded documents are available under /data and can be read "
+    "or modified using the execute_terminal tool. Continue chaining tools until "
+    "the user's request is fully addressed and then deliver a concise, coherent final reply."
 )

src/db.py

@@ -46,12 +46,22 @@ class Message(BaseModel):
     created_at = DateTimeField(default=datetime.utcnow)
 
 
+class Document(BaseModel):
+    id = AutoField()
+    user = ForeignKeyField(User, backref="documents")
+    file_path = CharField()
+    original_name = CharField()
+    created_at = DateTimeField(default=datetime.utcnow)
+
+
 __all__ = [
     "_db",
     "User",
     "Conversation",
     "Message",
+    "Document",
     "reset_history",
+    "add_document",
 ]
 
 
@@ -59,7 +69,7 @@ def init_db() -> None:
     """Initialise the database and create tables if they do not exist."""
     if _db.is_closed():
         _db.connect()
-    _db.create_tables([User, Conversation, Message])
+    _db.create_tables([User, Conversation, Message, Document])
 
 
 def reset_history(username: str, session_name: str) -> int:
@@ -79,3 +89,12 @@ def reset_history(username: str, session_name: str) -> int:
     if not Conversation.select().where(Conversation.user == user).exists():
         user.delete_instance()
     return deleted
+
+
+def add_document(username: str, file_path: str, original_name: str) -> Document:
+    """Record an uploaded document and return the created entry."""
+
+    init_db()
+    user, _ = User.get_or_create(username=username)
+    doc = Document.create(user=user, file_path=file_path, original_name=original_name)
+    return doc

src/tools.py

@@ -20,7 +20,7 @@ def set_vm(vm: LinuxVM | None) -> None:
 def execute_terminal(command: str) -> str:
     """
     Execute a shell command in a Linux terminal.
-    Use this tool to run various commands.
+    Use this tool to inspect uploaded documents under ``/data`` or run other commands.
 
     The command is executed with network access enabled. Output from both
     ``stdout`` and ``stderr`` is captured and returned. Commands are killed if
@@ -29,7 +29,7 @@ def execute_terminal(command: str) -> str:
     timeout = 2
     if not command:
         return "No command provided."
-    
+
     if _VM:
         try:
             return _VM.execute(command, timeout=timeout)

src/vm.py

@@ -2,6 +2,9 @@ from __future__ import annotations
 
 import subprocess
 import uuid
+from pathlib import Path
+
+from .config import UPLOAD_DIR
 
 from .log import get_logger
 
@@ -11,10 +14,14 @@ _LOG = get_logger(__name__)
 class LinuxVM:
     """Manage a lightweight Linux VM using Docker."""
 
-    def __init__(self, image: str = "ubuntu:latest") -> None:
+    def __init__(
+        self, image: str = "ubuntu:latest", host_dir: str = UPLOAD_DIR
+    ) -> None:
         self._image = image
         self._name = f"chat-vm-{uuid.uuid4().hex[:8]}"
         self._running = False
+        self._host_dir = Path(host_dir)
+        self._host_dir.mkdir(parents=True, exist_ok=True)
 
     def start(self) -> None:
         """Start the VM if it is not already running."""
@@ -36,6 +43,8 @@ class LinuxVM:
                     "-d",
                     "--name",
                     self._name,
+                    "-v",
+                    f"{self._host_dir}:/data",
                     self._image,
                     "sleep",
                     "infinity",