Add database support and persist chat history

README.md

@@ -1 +1,11 @@
-# llm-backend
+# llm-backend
+
+This project provides a simple async interface to interact with an Ollama model and demonstrates basic tool usage. Chat histories are stored in a local SQLite database using Peewee.
+
+## Usage
+
+```bash
+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`.

requirements.txt

@@ -1,2 +1,3 @@
 colorlog
-ollama
+ollama
+peewee

run.py

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

src/__init__.py

@@ -1,4 +1,4 @@
 from .chat import ChatSession
 from .tools import add_two_numbers
 
-__all__: list[str] = ["ChatSession", "add_two_numbers"]
+__all__ = ["ChatSession", "add_two_numbers"]

src/chat.py

@@ -5,6 +5,7 @@ from typing import List
 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 .log import get_logger
 from .schema import Msg
 from .tools import add_two_numbers
@@ -14,12 +15,17 @@ _LOG = get_logger(__name__)
 
 class ChatSession:
     def __init__(self, host: str = OLLAMA_HOST, model: str = MODEL_NAME) -> None:
+        init_db()
         self._client = AsyncClient(host=host)
         self._model = model
 
     async def __aenter__(self) -> "ChatSession":
         return self
 
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        if not _db.is_closed():
+            _db.close()
+
     async def ask(self, messages: List[Msg], *, think: bool = True) -> ChatResponse:
         return await self._client.chat(
             self._model,
@@ -32,6 +38,7 @@ class ChatSession:
         self,
         messages: List[Msg],
         response: ChatResponse,
+        conversation: Conversation,
         depth: int = 0,
     ) -> ChatResponse:
         if depth >= MAX_TOOL_CALL_DEPTH or not response.message.tool_calls:
@@ -47,17 +54,43 @@ class ChatSession:
                         "content": str(result),
                     }
                 )
+                Message.create(
+                    conversation=conversation,
+                    role="tool",
+                    content=str(result),
+                )
                 nxt = await self.ask(messages, think=True)
-                return await self._handle_tool_calls(messages, nxt, depth + 1)
+                Message.create(
+                    conversation=conversation,
+                    role="assistant",
+                    content=nxt.message.content,
+                )
+                return await self._handle_tool_calls(
+                    messages, nxt, conversation, depth + 1
+                )
 
         return response
 
     async def chat(self, prompt: str) -> str:
+        conversation = Conversation.create()
+        Message.create(conversation=conversation, role="user", content=prompt)
         messages: List[Msg] = [{"role": "user", "content": prompt}]
         response = await self.ask(messages)
         messages.append(response.message.model_dump())
+        Message.create(
+            conversation=conversation,
+            role="assistant",
+            content=response.message.content,
+        )
 
         _LOG.info("Thinking:\n%s", response.message.thinking or "<no thinking trace>")
 
-        final_resp = await self._handle_tool_calls(messages, response)
+        final_resp = await self._handle_tool_calls(messages, response, conversation)
+        if final_resp is not response:
+            # final response after handling tool calls
+            Message.create(
+                conversation=conversation,
+                role="assistant",
+                content=final_resp.message.content,
+            )
         return final_resp.message.content

src/config.py

@@ -5,4 +5,4 @@ from typing import Final
 
 MODEL_NAME: Final[str] = os.getenv("OLLAMA_MODEL", "qwen3")
 OLLAMA_HOST: Final[str] = os.getenv("OLLAMA_HOST", "http://localhost:11434")
-MAX_TOOL_CALL_DEPTH: Final[int] = 5
+MAX_TOOL_CALL_DEPTH: Final[int] = 5

src/db.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from datetime import datetime
+from pathlib import Path
+
+from peewee import (
+    AutoField,
+    CharField,
+    DateTimeField,
+    ForeignKeyField,
+    Model,
+    SqliteDatabase,
+    TextField,
+)
+
+
+_DB_PATH = Path(__file__).resolve().parent.parent / "chat.db"
+_db = SqliteDatabase(_DB_PATH)
+
+
+class BaseModel(Model):
+    class Meta:
+        database = _db
+
+
+class Conversation(BaseModel):
+    id = AutoField()
+    started_at = DateTimeField(default=datetime.utcnow)
+
+
+class Message(BaseModel):
+    id = AutoField()
+    conversation = ForeignKeyField(Conversation, backref="messages")
+    role = CharField()
+    content = TextField()
+    created_at = DateTimeField(default=datetime.utcnow)
+
+
+__all__ = ["_db", "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])

src/schema.py

@@ -9,4 +9,4 @@ class Msg(TypedDict, total=False):
     role: Literal["user", "assistant", "tool"]
     content: str
     name: Optional[str]
-    tool_calls: Optional[List[Message.ToolCall]]
+    tool_calls: Optional[List[Message.ToolCall]]

src/tools.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-__all__: list[str] = ["add_two_numbers"]
+__all__ = ["add_two_numbers"]
 
 
 def add_two_numbers(a: int, b: int) -> int:  # noqa: D401
@@ -13,4 +13,4 @@ def add_two_numbers(a: int, b: int) -> int:  # noqa: D401
     Returns:
         int: The sum of the two numbers.
     """
-    return a + b
+    return a + b