Lesson 2: The running project: notes, a store and a model

%pip install -q -U langchain langgraph langchain-groq langchain-google-genai langchain-ollama

import os
from langchain.chat_models import init_chat_model

PROVIDER = "groq"   # "groq" or "gemini" in Colab; "ollama" on your own computer
MODELS = {"groq": "groq:openai/gpt-oss-20b", "gemini": "google_genai:gemini-3.5-flash-lite", "ollama": "ollama:llama3.2:3b"}
KEY = {"groq": "GROQ_API_KEY", "gemini": "GOOGLE_API_KEY", "ollama": None}[PROVIDER]
if KEY and KEY not in os.environ:
    from google.colab import userdata   # keys live in Colab's Secrets panel
    os.environ[KEY] = userdata.get(KEY)
llm = init_chat_model(MODELS[PROVIDER], temperature=0, max_retries=6)

from pathlib import Path

os.environ["LLM_PROVIDER"] = PROVIDER          # the package reads the provider from here
Path("notes_assistant").mkdir(exist_ok=True)
Path("notes_assistant/__init__.py").write_text('"""Nare\'s notes assistant."""\n')
print(sorted(p.name for p in Path("notes_assistant").iterdir()))

from pathlib import Path

NOTES = {
    "coffee-and-focus": """---
title: Caffeine and focus
tags: [caffeine, focus, sleep]
created: 2026-04-01
---
Caffeine blocks adenosine receptors, so sleep pressure is masked rather than removed.
Half-life is about 5 hours, so a 4 pm coffee still leaves half the dose at 9 pm.
Moderate doses (about 100 mg) improve reaction time; above 400 mg per day anxiety rises.
Personal rule: no coffee after 1 pm on days before an experiment.
""",
    "deep-work": """---
title: Deep work blocks
tags: [focus, productivity]
created: 2026-07-08
---
Newport defines deep work as distraction-free concentration that pushes your abilities.
My best writing happens between 8 and 11 am before email.
Tracking: 14 deep hours last week, target is 15.
Phone in another room helps more than any app blocker.
""",
    "lab-sync-2026-09-02": """---
title: Lab sync, 2 September 2026
tags: [meeting, sleep]
created: 2026-09-02
---
Attendees: Priya, Tomas, Nare.
Decision: run the nap study with 24 participants instead of 16, budget approved by Priya.
Action: Tomas books the sleep lab for the weeks of 5 and 12 October.
Action: Nare drafts the consent form by 15 September.
Risk: EEG caps are shared with another group, so booking conflicts are likely.
""",
}

NOTES.update({
    "pomodoro-technique": """---
title: Pomodoro technique
tags: [focus, productivity]
created: 2026-04-15
---
Work for 25 minutes, then take a 5 minute break; after four rounds take a longer 20 minute break.
Works well for writing and reading papers, badly for debugging where context switching is costly.
For deep tasks I now use 50 minute blocks with 10 minute breaks instead.
""",
    "reading-list": """---
title: Reading list
tags: [reading, sleep, memory]
created: 2026-06-20
---
Why We Sleep, Matthew Walker (2017): readable, but some claims are disputed.
Make It Stick, Brown, Roediger and McDaniel (2014): retrieval practice and spacing.
Deep Work, Cal Newport (2016): scheduling focus time.
Pending: a review paper on targeted memory reactivation during sleep.
""",
    "sleep-and-memory": """---
title: Sleep and memory consolidation
tags: [sleep, memory, neuroscience]
created: 2026-03-02
---
Slow-wave sleep replays the day's hippocampal activity and moves it into long-term cortical storage.
Walker (2017) reports that a night without sleep cuts new memory formation by about 40 percent.
A 20 minute nap after learning helps recall of word pairs, but only if the nap reaches stage 2 sleep.
Open question: does REM sleep matter more for procedural skills than for facts?
""",
})

NOTES.update({
    "spaced-repetition": """---
title: Spaced repetition
tags: [memory, learning]
created: 2026-03-10
---
Reviewing material at growing intervals (1 day, 3 days, 7 days, 21 days) beats massed study.
The forgetting curve from Ebbinghaus drops fastest in the first 24 hours after learning.
Anki uses the SM-2 algorithm; FSRS is a newer scheduler that predicts recall probability.
Combine with sleep: a review before bed seems to consolidate better than a morning review.
""",
    "zettelkasten": """---
title: Zettelkasten note-taking
tags: [notes, productivity, learning]
created: 2026-05-05
---
One idea per note, written in your own words, linked to related notes by id.
Links matter more than folders: a note should be findable from several directions.
Luhmann kept about 90,000 index cards and cross-referenced them by hand.
This notes folder follows the same rule: one idea per Markdown file.
""",
})

def write_notes(folder="notes", notes=NOTES):
    """Write every note to <folder>/<note_id>.md and return the folder."""
    root = Path(folder)
    root.mkdir(parents=True, exist_ok=True)
    for note_id, text in notes.items():
        (root / f"{note_id}.md").write_text(text, encoding="utf-8")
    return root

root = write_notes()
print(len(list(root.glob("*.md"))), "notes in", root)
body = NOTES["lab-sync-2026-09-02"].split("---\n")[2].splitlines()
print("body line 2:", body[1])
print(sum(len(t.split("---\n")[2].split()) for t in NOTES.values()), "words of body text")

%%writefile notes_assistant/store.py
"""Nare's notes as a folder of Markdown files. This file knows nothing about MCP."""
from __future__ import annotations

import re
from dataclasses import dataclass
from datetime import date
from pathlib import Path

NOTE_ID_PATTERN = re.compile(r"^[a-z0-9][a-z0-9-]{0,63}$")
WORD_PATTERN = re.compile(r"[a-z0-9]+")
MAX_TITLE_LENGTH = 120
MAX_BODY_LENGTH = 20_000


class NoteError(Exception):
    """Base class for every error the store raises on purpose."""


class NoteNotFound(NoteError):
    """No note has this id."""


class NoteExists(NoteError):
    """A note with this id already exists. The store never overwrites."""


class InvalidNote(NoteError):
    """An id, title, body or query the store refuses."""

%%writefile -a notes_assistant/store.py
@dataclass(frozen=True)
class Note:
    note_id: str
    title: str
    tags: tuple[str, ...]
    created: str
    body: str

    def snippet(self, length: int = 120) -> str:
        """The first line of the body, cut to `length` characters."""
        first = self.body.splitlines()[0] if self.body else ""
        return first if len(first) <= length else first[: length - 3] + "..."


@dataclass(frozen=True)
class SearchHit:
    note_id: str
    title: str
    score: int
    snippet: str
    tags: tuple[str, ...] = ()


def slugify(title: str) -> str:
    """'Nap study: consent form!' -> 'nap-study-consent-form'."""
    return "-".join(WORD_PATTERN.findall(title.lower()))[:64].strip("-")


def validate_note_id(note_id: str) -> str:
    """Refuse any id that is not a simple slug, before a file is touched."""
    if not NOTE_ID_PATTERN.fullmatch(note_id):
        raise InvalidNote(f"Invalid note id {note_id!r}. Use lowercase letters, digits "
                          "and hyphens, for example 'sleep-and-memory'.")
    return note_id

%%writefile -a notes_assistant/store.py
def _parse(note_id: str, text: str) -> Note:
    """Split a file into its frontmatter fields and its body."""
    meta, body = {}, text
    if text.startswith("---\n"):
        header, _, body = text[4:].partition("\n---\n")
        for line in header.splitlines():
            key, _, value = line.partition(":")
            meta[key.strip()] = value.strip()
    tags = tuple(t.strip() for t in meta.get("tags", "").strip("[]").split(",") if t.strip())
    return Note(note_id, meta.get("title", note_id), tags, meta.get("created", ""), body.strip())


class NoteStore:
    """A folder of Markdown notes, one file per note, named <note_id>.md."""

    def __init__(self, root: Path | str) -> None:
        self.root = Path(root)
        if not self.root.is_dir():
            raise NoteError(f"Notes folder not found: {self.root}")

    def _path(self, note_id: str) -> Path:
        return self.root / f"{validate_note_id(note_id)}.md"   # the only place an id becomes a path

    def list_notes(self) -> list[Note]:
        return [self.get(p.stem) for p in sorted(self.root.glob("*.md"))]

    def get(self, note_id: str) -> Note:
        path = self._path(note_id)
        if not path.is_file():
            raise NoteNotFound(f"No note with id {note_id!r}.")
        return _parse(note_id, path.read_text(encoding="utf-8"))

%%writefile -a notes_assistant/store.py
    def search(self, query: str, limit: int = 5, tag: str | None = None) -> list[SearchHit]:
        """Rank notes by query words: 1 per body match, 3 per title match, 2 per tag match."""
        terms = WORD_PATTERN.findall(query.lower())
        if not terms:
            raise InvalidNote("The query must contain at least one letter or digit.")
        hits = []
        for note in self.list_notes():
            if tag and tag not in note.tags:
                continue
            body = WORD_PATTERN.findall(note.body.lower())
            title = WORD_PATTERN.findall(note.title.lower())
            tags = [t.lower() for t in note.tags]
            score = sum(body.count(t) + 3 * title.count(t) + 2 * tags.count(t) for t in terms)
            if score > 0:
                hits.append(SearchHit(note.note_id, note.title, score, note.snippet(), note.tags))
        hits.sort(key=lambda h: (-h.score, h.note_id))   # ties broken by note id
        return hits[:limit]

%%writefile -a notes_assistant/store.py
    def create(self, title: str, body: str, tags=(), *, created: date) -> Note:
        """Write a new note stamped with `created`. Never overwrites an existing note."""
        title = title.strip()
        if not 1 <= len(title) <= MAX_TITLE_LENGTH:
            raise InvalidNote(f"The title must be 1 to {MAX_TITLE_LENGTH} characters.")
        if not body.strip() or len(body) > MAX_BODY_LENGTH:
            raise InvalidNote(f"The body must be 1 to {MAX_BODY_LENGTH:,} characters.")
        note_id = slugify(title)
        if not note_id:
            raise InvalidNote("The title must contain at least one letter or digit.")
        path = self._path(note_id)
        if path.exists():
            raise NoteExists(f"A note with id {note_id!r} already exists.")
        clean_tags = [slugify(t) for t in tags if slugify(t)]
        header = f"---\ntitle: {title}\ntags: [{', '.join(clean_tags)}]\ncreated: {created.isoformat()}\n---\n"
        path.write_text(header + body.strip() + "\n", encoding="utf-8")
        return self.get(note_id)

from datetime import date
from notes_assistant.store import NoteStore, NoteError

store = NoteStore("notes")
for hit in store.search("nap study participants"):
    print(f"{hit.score}  {hit.note_id:<20} {hit.snippet}")
for bad in ["../secrets", "no-such-note"]:
    try:
        store.get(bad)
    except NoteError as err:
        print(type(err).__name__, "|", err)

import shutil
shutil.rmtree("scratch-notes", ignore_errors=True)   # start the copy afresh on every run
scratch = NoteStore(write_notes("scratch-notes"))    # writes go to a copy, never notes/
note = scratch.create("Nap study consent form", "Draft due 15 September.", ["Meeting"], created=date(2026, 9, 15))
print("created", note.note_id, note.tags, note.created)
try:
    scratch.create("Nap study: consent form!", "Second try.", created=date(2026, 9, 15))
except NoteError as err:
    print(type(err).__name__, "|", err)

%%writefile notes_assistant/llm.py
"""chat(): one function for Groq, Gemini and Ollama, built on LangChain."""
from __future__ import annotations

import json
import os
from dataclasses import dataclass, field
from typing import Any

from langchain.chat_models import init_chat_model

MODELS = {"groq": "groq:openai/gpt-oss-20b", "gemini": "google_genai:gemini-3.5-flash-lite", "ollama": "ollama:llama3.2:3b"}
KEYS = {"groq": "GROQ_API_KEY", "gemini": "GOOGLE_API_KEY", "ollama": None}
LLM_PROVIDER = os.environ.get("LLM_PROVIDER", "groq").lower()
LLM_MODEL = os.environ.get("LLM_MODEL") or MODELS.get(LLM_PROVIDER, "")


@dataclass
class ToolCall:
    id: str
    name: str
    arguments: dict[str, Any]


@dataclass
class ChatReply:
    content: str | None
    tool_calls: list[ToolCall] = field(default_factory=list)
    error: str | None = None            # "provider_error: ..." after two refusals

    def as_message(self) -> dict[str, Any]:
        """This reply as an assistant message, to add to the conversation."""
        message = {"role": "assistant", "content": self.content or ""}
        if self.tool_calls:
            message["tool_calls"] = [{"id": c.id, "type": "function", "function": {
                "name": c.name, "arguments": json.dumps(c.arguments)}} for c in self.tool_calls]
        return message

%%writefile -a notes_assistant/llm.py


_CACHE: dict[str, Any] = {}


def _model(provider: str):
    """Build the LangChain model for a provider once, then reuse it."""
    if provider not in MODELS:
        raise ValueError(f"Unknown LLM_PROVIDER {provider!r}. Choose one of {sorted(MODELS)}.")
    key = KEYS[provider]
    if key and not os.environ.get(key):
        raise RuntimeError(f"Set {key} to use {provider}.")
    if provider not in _CACHE:
        name = LLM_MODEL if provider == LLM_PROVIDER else MODELS[provider]
        _CACHE[provider] = init_chat_model(name, temperature=0, max_retries=6)
    return _CACHE[provider]


def chat(messages: list[dict[str, Any]], tools: list[dict] | None = None,
         provider: str | None = None) -> ChatReply:
    """Send messages and optional tools; return text or tool calls. Retries a refusal once."""
    model = _model((provider or LLM_PROVIDER).lower())
    if tools:
        model = model.bind_tools(tools)
    for attempt in range(2):
        try:
            message = model.invoke(messages)
            break
        except Exception as err:
            last = err
    else:
        return ChatReply(content=None, error=f"provider_error: {type(last).__name__}")
    calls = [ToolCall(c["id"] or f"call_{i}", c["name"], c["args"]) for i, c in enumerate(message.tool_calls)]
    return ChatReply(content=message.text or None, tool_calls=calls)

import json
from notes_assistant.llm import ToolCall, chat

SEARCH_TOOL = {
    "type": "function",
    "function": {
        "name": "search_notes",
        "description": "Search the user's research notes by keyword and return the best matching notes.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "Words to look for, for example 'sleep memory'."},
                "limit": {"type": "integer", "description": "Maximum number of notes to return.", "default": 5},
            },
            "required": ["query"],
        },
    },
}

def run_tool(call: ToolCall) -> str:
    """Map the tool the model chose to Python code; return text for the model."""
    if call.name != "search_notes":
        return f"Unknown tool {call.name!r}"
    try:
        hits = store.search(call.arguments["query"], limit=call.arguments.get("limit", 5))
    except NoteError as err:
        return f"Error: {err}"
    return json.dumps([{"note_id": h.note_id, "title": h.title, "snippet": h.snippet} for h in hits])

result = run_tool(ToolCall("demo", "search_notes", {"query": "nap study participants"}))
print(len(result), "characters")
print(json.loads(result)[0])
print("contains 24:", "24" in result)

messages = [{"role": "user", "content": "How many participants will the nap study have?"}]
reply = chat(messages, tools=[SEARCH_TOOL])
if reply.tool_calls:
    messages.append(reply.as_message())
    for call in reply.tool_calls:
        print("model asked for:", call.name, call.arguments)
        messages.append({"role": "tool", "tool_call_id": call.id, "content": run_tool(call)})
    reply = chat(messages, tools=[SEARCH_TOOL])
if reply.error:
    print("no answer:", reply.error)
else:
    print("answer:", reply.content or "(empty reply)")

for hit in store.search("how many participants will the nap study have"):
    print(hit.score, hit.note_id)

import copy

SEARCH_TOOL_V2 = copy.deepcopy(SEARCH_TOOL)
SEARCH_TOOL_V2["function"]["parameters"]["properties"]["tag"] = {
    "type": "string", "description": "Only search notes with this tag, for example 'meeting'."}


def run_tool_v2(call: ToolCall) -> str:
    """Like run_tool, with an optional tag filter."""
    if call.name != "search_notes":
        return f"Unknown tool {call.name!r}"
    args = call.arguments
    try:
        hits = store.search(args["query"], limit=args.get("limit", 5), tag=args.get("tag"))
    except NoteError as err:
        return f"Error: {err}"
    return json.dumps([{"note_id": h.note_id, "title": h.title, "snippet": h.snippet} for h in hits])

found = json.loads(run_tool_v2(ToolCall("ex", "search_notes", {"query": "sleep", "tag": "meeting"})))
assert [h["note_id"] for h in found] == ["lab-sync-2026-09-02"], found
assert "tag" not in SEARCH_TOOL["function"]["parameters"]["properties"]
assert "tag" not in SEARCH_TOOL_V2["function"]["parameters"]["required"]
print("exercise passed:", found[0]["note_id"])