Add user-based session tracking

README.md

@@ -8,4 +8,4 @@ This project provides a simple async interface to interact with an Ollama model
 python run.py
 ```
 
-The script will ask the model to compute an arithmetic expression and print the answer. Conversations are automatically persisted to `chat.db`.
+The script will ask the model to compute an arithmetic expression and print the answer. Conversations are automatically persisted to `chat.db` and are now associated with a user and session.

run.py

@@ -6,7 +6,7 @@ from src.chat import ChatSession
 
 
 async def _main() -> None:
-    async with ChatSession() as chat:
+    async with ChatSession(user="demo_user") as chat:
         answer = await chat.chat("What is 10 + 23?")
         print("\n>>>", answer)
 

src/chat.py

@@ -6,7 +6,7 @@ import json
 from ollama import AsyncClient, ChatResponse
 
 from .config import MAX_TOOL_CALL_DEPTH, MODEL_NAME, OLLAMA_HOST
-from .db import Conversation, Message, _db, init_db
+from .db import Conversation, Message, User, _db, init_db
 from .log import get_logger
 from .schema import Msg
 from .tools import add_two_numbers
@@ -15,10 +15,11 @@ _LOG = get_logger(__name__)
 
 
 class ChatSession:
-    def __init__(self, host: str = OLLAMA_HOST, model: str = MODEL_NAME) -> None:
+    def __init__(self, user: str = "default", host: str = OLLAMA_HOST, model: str = MODEL_NAME) -> None:
         init_db()
         self._client = AsyncClient(host=host)
         self._model = model
+        self._user, _ = User.get_or_create(username=user)
 
     async def __aenter__(self) -> "ChatSession":
         return self
@@ -82,7 +83,7 @@ class ChatSession:
         return response
 
     async def chat(self, prompt: str) -> str:
-        conversation = Conversation.create()
+        conversation = Conversation.create(user=self._user)
         Message.create(conversation=conversation, role="user", content=prompt)
         messages: List[Msg] = [{"role": "user", "content": prompt}]
         response = await self.ask(messages)

src/db.py

@@ -23,8 +23,14 @@ class BaseModel(Model):
         database = _db
 
 
+class User(BaseModel):
+    id = AutoField()
+    username = CharField(unique=True)
+
+
 class Conversation(BaseModel):
     id = AutoField()
+    user = ForeignKeyField(User, backref="conversations")
     started_at = DateTimeField(default=datetime.utcnow)
 
 
@@ -36,11 +42,11 @@ class Message(BaseModel):
     created_at = DateTimeField(default=datetime.utcnow)
 
 
-__all__ = ["_db", "Conversation", "Message"]
+__all__ = ["_db", "User", "Conversation", "Message"]
 
 
 def init_db() -> None:
     """Initialise the database and create tables if they do not exist."""
     if _db.is_closed():
         _db.connect()
-    _db.create_tables([Conversation, Message])
+    _db.create_tables([User, Conversation, Message])