jehoctor-rag-demo 0.1.1.dev1__py3-none-any.whl → 0.2.1__py3-none-any.whl

rag_demo/app_protocol.py

@@ -0,0 +1,101 @@
+"""Interface for the logic to call back into the app code.
+
+This is necessary to make the logic code testable. We don't want to have to run all the app code to test the logic. And,
+we want to have a high degree of confidence when mocking out the app code in logic tests. The basic pattern is that each
+piece of functionality that the logic depends on will have a protocol and an implementation of that protocol using the
+Textual App. In the tests, we create a mock implementation of the same protocol. Correctness of the logic is defined by
+its ability to work correctly with any implementation of the protocol, not just the implementation backed by the app.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable
+
+    from textual.worker import Worker
+
+
+class LoggerProtocol(Protocol):
+    """Protocol that mimics textual.Logger."""
+
+    def __call__(self, *args: object, **kwargs: object) -> None:
+        """Log a message.
+
+        Args:
+            *args (object): Logged directly to the message separated by spaces.
+            **kwargs (object): Logged to the message as f"{key}={value!r}", separated by spaces.
+        """
+
+    def verbosity(self, *, verbose: bool) -> LoggerProtocol:
+        """Get a new logger with selective verbosity.
+
+        Note that unlike when using this method on a Textual logger directly, the type system will enforce that you use
+        `verbose` as a keyword argument (not a positional argument). I made this change to address ruff's FBT001 rule.
+        Put simply, this requirement makes the calling code easier to read.
+        https://docs.astral.sh/ruff/rules/boolean-type-hint-positional-argument/
+
+        Args:
+            verbose: True to use HIGH verbosity, otherwise NORMAL.
+
+        Returns:
+            New logger.
+        """
+
+    @property
+    def verbose(self) -> LoggerProtocol:
+        """A verbose logger."""
+
+    @property
+    def event(self) -> LoggerProtocol:
+        """Logs events."""
+
+    @property
+    def debug(self) -> LoggerProtocol:
+        """Logs debug messages."""
+
+    @property
+    def info(self) -> LoggerProtocol:
+        """Logs information."""
+
+    @property
+    def warning(self) -> LoggerProtocol:
+        """Logs warnings."""
+
+    @property
+    def error(self) -> LoggerProtocol:
+        """Logs errors."""
+
+    @property
+    def system(self) -> LoggerProtocol:
+        """Logs system information."""
+
+    @property
+    def logging(self) -> LoggerProtocol:
+        """Logs from stdlib logging module."""
+
+    @property
+    def worker(self) -> LoggerProtocol:
+        """Logs worker information."""
+
+
+ResultType = TypeVar("ResultType")
+
+
+class AppProtocol(Protocol):
+    """Protocol for the subset of what the main App can do that the runtime needs."""
+
+    def run_worker(self, work: Awaitable[ResultType], *, thread: bool = False) -> Worker[ResultType]:
+        """Run a coroutine in the background.
+
+        See https://textual.textualize.io/guide/workers/.
+
+        Args:
+            work (Awaitable[ResultType]): The coroutine to run.
+            thread (bool): Mark the worker as a thread worker.
+        """
+
+    @property
+    def log(self) -> LoggerProtocol:
+        """Returns the application logger."""

rag_demo/constants.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from enum import StrEnum, auto
+
+
+class LocalProviderType(StrEnum):
+    """Enum of supported local LLM backend provider types."""
+
+    HUGGING_FACE = auto()
+    LLAMA_CPP = auto()
+    OLLAMA = auto()

rag_demo/db.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import aiosqlite
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class AtomicIDManager:
+    """A database manager for managing thread IDs.
+
+    This was written by Claude, and I fixed it up with feedback from Ruff and Flake8.
+    Maybe one day the app logic database will require something fancier, but this gets the job done now.
+
+    As you can see from the conversation with Claude, this was quite a simple task for it:
+    https://claude.ai/share/227d08ff-96a3-495a-9f56-509a1fd528f7
+    """
+
+    def __init__(self, db_path: str | Path) -> None:
+        """Initialize the database manager."""
+        self.db_path = db_path
+
+    async def initialize(self) -> None:
+        """Initialize the database and create the table if it doesn't exist."""
+        async with aiosqlite.connect(self.db_path) as db:
+            # Enable WAL mode for better concurrent access
+            await db.execute("PRAGMA journal_mode=WAL")
+
+            await db.execute("""
+                CREATE TABLE IF NOT EXISTS claimed_ids (
+                    id INTEGER PRIMARY KEY
+                )
+            """)
+            await db.commit()
+
+    async def claim_next_id(self) -> int:
+        """Atomically find the max id, increment it, and claim it. Returns the newly claimed ID.
+
+        This operation is atomic and multiprocess-safe because:
+        1. SQLite serializes writes by default
+        2. We use IMMEDIATE transaction to acquire write lock immediately
+        3. The entire operation happens in a single transaction
+        """
+        async with aiosqlite.connect(self.db_path) as db:
+            # Start an IMMEDIATE transaction to get write lock right away
+            await db.execute("BEGIN IMMEDIATE")
+
+            try:
+                # Find the current max ID
+                async with db.execute("SELECT MAX(id) FROM claimed_ids") as cursor:
+                    row = await cursor.fetchone()
+                    max_id = row[0] if row is not None and row[0] is not None else 0
+
+                # Calculate next ID
+                next_id = max_id + 1
+
+                # Insert the new ID
+                await db.execute("INSERT INTO claimed_ids (id) VALUES (?)", (next_id,))
+
+                # Commit the transaction
+                await db.commit()
+
+            except Exception:
+                await db.rollback()
+                raise
+
+            else:
+                return next_id
+
+    async def get_all_claimed_ids(self) -> list[int]:
+        """Retrieve all claimed IDs."""
+        async with (
+            aiosqlite.connect(self.db_path) as db,
+            db.execute("SELECT id FROM claimed_ids ORDER BY id") as cursor,
+        ):
+            rows = await cursor.fetchall()
+            return [row[0] for row in rows]
+
+    async def get_count(self) -> int:
+        """Get the total number of claimed IDs."""
+        async with aiosqlite.connect(self.db_path) as db, db.execute("SELECT COUNT(*) FROM claimed_ids") as cursor:
+            row = await cursor.fetchone()
+            if row is None:
+                raise ValueError("A SQL COUNT query should always return at least one row")  # noqa: EM101, TRY003
+            return row[0]

rag_demo/dirs.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+
+from platformdirs import PlatformDirs
+
+_appdirs = PlatformDirs(appname="jehoctor-rag-demo", ensure_exists=True)
+
+
+def _ensure(dir_: Path) -> Path:
+    dir_.mkdir(parents=True, exist_ok=True)
+    return dir_
+
+
+DATA_DIR = _appdirs.user_data_path
+CONFIG_DIR = _appdirs.user_config_path

rag_demo/logic.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import time
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, cast
+
+from datasets import Dataset, load_dataset
+from langchain_core.exceptions import LangChainException
+
+from rag_demo import dirs
+from rag_demo.agents import (
+    Agent,
+    AgentProvider,
+    HuggingFaceAgentProvider,
+    LlamaCppAgentProvider,
+    OllamaAgentProvider,
+)
+from rag_demo.db import AtomicIDManager
+from rag_demo.modes.chat import Response, StoppedStreamError
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Sequence
+    from pathlib import Path
+
+    from rag_demo.app_protocol import AppProtocol
+    from rag_demo.constants import LocalProviderType
+    from rag_demo.modes import ChatScreen
+
+
+class UnknownPreferredProviderError(ValueError):
+    """Raised when the preferred provider cannot be checked first due to being unknown."""
+
+    def __init__(self, preferred_provider: LocalProviderType) -> None:  # noqa: D107
+        super().__init__(f"Unknown preferred provider: {preferred_provider}")
+
+
+class NoProviderError(RuntimeError):
+    """Raised when no provider could provide an agent."""
+
+    def __init__(self) -> None:  # noqa: D107
+        super().__init__("No provider could provide an agent.")
+
+
+class Runtime:
+    """The application logic with asynchronously initialized resources."""
+
+    def __init__(
+        self,
+        logic: Logic,
+        app: AppProtocol,
+        agent: Agent,
+        thread_id_manager: AtomicIDManager,
+    ) -> None:
+        """Initialize the runtime.
+
+        Args:
+            logic (Logic): The application logic.
+            app (AppProtocol): The application interface.
+            agent (Agent): The agent to use.
+            thread_id_manager (AtomicIDManager): The thread ID manager.
+        """
+        self.runtime_start_time = time.time()
+        self.logic = logic
+        self.app = app
+        self.agent = agent
+        self.thread_id_manager = thread_id_manager
+
+        self.current_thread: int | None = None
+        self.generating = False
+
+    def _get_rag_datasets(self) -> None:
+        self.qa_test: Dataset = cast(
+            "Dataset",
+            load_dataset("rag-datasets/rag-mini-wikipedia", "question-answer", split="test"),
+        )
+        self.corpus: Dataset = cast(
+            "Dataset",
+            load_dataset("rag-datasets/rag-mini-wikipedia", "text-corpus", split="passages"),
+        )
+
+    async def stream_response(self, response_widget: Response, request_text: str, thread: str) -> None:
+        """Worker method for streaming tokens from the active agent to a response widget.
+
+        Args:
+            response_widget (Response): Target response widget for streamed tokens.
+            request_text (str): Text of the user request.
+            thread (str): ID of the current thread.
+        """
+        self.generating = True
+        async with response_widget.stream_writer() as writer:
+            try:
+                async for message_chunk in self.agent.astream(request_text, thread, self.app):
+                    await writer.write(message_chunk)
+            except StoppedStreamError as e:
+                response_widget.set_shown_object(e)
+            except LangChainException as e:
+                response_widget.set_shown_object(e)
+        self.generating = False
+
+    def new_conversation(self, chat_screen: ChatScreen) -> None:
+        """Clear the screen and start a new conversation with the agent.
+
+        Args:
+            chat_screen (ChatScreen): The chat screen to clear.
+        """
+        self.current_thread = None
+        chat_screen.clear_chats()
+
+    async def submit_request(self, chat_screen: ChatScreen, request_text: str) -> bool:
+        """Submit a new user request in the current conversation.
+
+        Args:
+            chat_screen (ChatScreen): The chat screen in which the request is submitted.
+            request_text (str): The text of the request.
+
+        Returns:
+            bool: True if the request was accepted for immediate processing, False otherwise.
+        """
+        if self.generating:
+            return False
+        self.generating = True
+        if self.current_thread is None:
+            chat_screen.log.info("Starting new thread")
+            self.current_thread = await self.thread_id_manager.claim_next_id()
+            chat_screen.log.info("Claimed thread id", self.current_thread)
+        chat_screen.new_request(request_text)
+        response = chat_screen.new_response()
+        chat_screen.run_worker(self.stream_response(response, request_text, str(self.current_thread)))
+        return True
+
+
+class Logic:
+    """Top-level application logic."""
+
+    def __init__(
+        self,
+        username: str | None = None,
+        preferred_provider_type: LocalProviderType | None = None,
+        application_start_time: float | None = None,
+        checkpoints_sqlite_db: str | Path = dirs.DATA_DIR / "checkpoints.sqlite3",
+        app_sqlite_db: str | Path = dirs.DATA_DIR / "app.sqlite3",
+        agent_providers: Sequence[AgentProvider] = (
+            LlamaCppAgentProvider(),
+            OllamaAgentProvider(),
+            HuggingFaceAgentProvider(),
+        ),
+    ) -> None:
+        """Initialize the application logic.
+
+        Args:
+            username (str | None, optional): The username provided as a command line argument. Defaults to None.
+            preferred_provider_type (LocalProviderType | None, optional): Provider type to prefer. Defaults to None.
+            application_start_time (float | None, optional): The time when the application started. Defaults to None.
+            checkpoints_sqlite_db (str | Path, optional): The connection string for the SQLite database used for
+                Langchain checkpointing. Defaults to (dirs.DATA_DIR / "checkpoints.sqlite3").
+            app_sqlite_db (str | Path, optional): The connection string for the SQLite database used for application
+                state such a thread metadata. Defaults to (dirs.DATA_DIR / "app.sqlite3").
+            agent_providers (Sequence[AgentProvider], optional): Sequence of agent providers in default preference
+                order. If preferred_provider_type is not None, this sequence will be reordered to bring providers of
+                that type to the front, using the original order to break ties. Defaults to (
+                    LlamaCppAgentProvider(),
+                    OllamaAgentProvider(),
+                    HuggingFaceAgentProvider(),
+                ).
+        """
+        self.logic_start_time = time.time()
+        self.username = username
+        self.preferred_provider_type = preferred_provider_type
+        self.application_start_time = application_start_time
+        self.checkpoints_sqlite_db = checkpoints_sqlite_db
+        self.app_sqlite_db = app_sqlite_db
+        self.agent_providers: Sequence[AgentProvider] = agent_providers
+
+    @asynccontextmanager
+    async def runtime(self, app: AppProtocol) -> AsyncIterator[Runtime]:
+        """Returns a runtime context for the application."""
+        thread_id_manager = AtomicIDManager(self.app_sqlite_db)
+        await thread_id_manager.initialize()
+
+        agent_providers: Sequence[AgentProvider] = self.agent_providers
+        if self.preferred_provider_type is not None:
+            preferred_providers: Sequence[AgentProvider] = tuple(
+                ap for ap in agent_providers if ap.type == self.preferred_provider_type
+            )
+            if len(preferred_providers) == 0:
+                raise UnknownPreferredProviderError(self.preferred_provider_type)
+            agent_providers = (
+                *preferred_providers,
+                *(ap for ap in agent_providers if ap.type != self.preferred_provider_type),
+            )
+        for agent_provider in agent_providers:
+            async with agent_provider.get_agent(checkpoints_sqlite_db=self.checkpoints_sqlite_db) as agent:
+                if agent is not None:
+                    yield Runtime(
+                        logic=self,
+                        app=app,
+                        agent=agent,
+                        thread_id_manager=thread_id_manager,
+                    )
+                    return
+        raise NoProviderError

rag_demo/markdown.py

@@ -0,0 +1,17 @@
+from markdown_it import MarkdownIt
+from markdown_it.rules_inline import StateInline
+
+
+def soft2hard_break_plugin(md: MarkdownIt) -> None:
+    md.inline.ruler2.push("soft2hard_break", _soft2hard_break_plugin)
+
+
+def _soft2hard_break_plugin(state: StateInline) -> None:
+    for token in state.tokens:
+        if token.type == "softbreak":
+            token.type = "hardbreak"
+
+
+def parser_factory() -> MarkdownIt:
+    """Modified parser that handles newlines according to LLM conventions."""
+    return MarkdownIt("gfm-like").use(soft2hard_break_plugin)

rag_demo/modes/__init__.py

@@ -0,0 +1,3 @@
+from .chat import ChatScreen
+from .config import ConfigScreen
+from .help import HelpScreen

rag_demo/modes/_logic_provider.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, cast
+
+from textual.screen import Screen
+from textual.widget import Widget
+
+if TYPE_CHECKING:
+    from rag_demo.logic import Logic, Runtime
+
+
+class LogicProvider(Protocol):
+    """Protocol for classes that contain application logic."""
+
+    logic: Logic
+
+    async def runtime(self) -> Runtime:
+        """Returns the application runtime of the parent app."""
+
+
+class LogicProviderScreen(Screen):
+    """A Screen that provides access to the application logic via its parent app."""
+
+    @property
+    def logic(self) -> Logic:
+        """Returns the application logic of the parent app."""
+        return cast("LogicProvider", self.app).logic
+
+    async def runtime(self) -> Runtime:
+        """Returns the application runtime of the parent app."""
+        return await cast("LogicProvider", self.app).runtime()
+
+
+class LogicProviderWidget(Widget):
+    """A Widget that provides access to the application logic via its parent app."""
+
+    @property
+    def logic(self) -> Logic:
+        """Returns the application logic of the parent app."""
+        return cast("LogicProvider", self.app).logic
+
+    async def runtime(self) -> Runtime:
+        """Returns the application runtime of the parent app."""
+        return await cast("LogicProvider", self.app).runtime()