feat: begin work on RAG, refactoring, improved docs and tests (#258)

* feat: started working on RAG (again)

* fix: fixed bugs in rag

* fix: fixed tests after refactor

* build(deps): updated dependencies

* build(deps): updated gptme-rag

* fix: made rag support optional

* test: fixed tests

* docs: fixed docs

* docs: fixed docs for computer tool

* docs: fixed docs, made some funcs private to hide them from autodocs

* test: fixed tests

* Apply suggestions from code review

* fix: fixed when running in a directory without gptme.toml

* fix: fixed lint

* test: fixed tests for rag

* build(deps): updated gptme-rag

* fix: add get_project_dir helper function

* fix: changed rag tool to use functions instead of execute

* Apply suggestions from code review

* fix: more refactoring of rag, made typechecking of imports stricter, fixed tests

* fix: fixed rag tool tests

* ci: fixed typechecking

* fix: hopefully finally fixed rag tests now

* test: made test_chain more reliable for gpt-4o-mini

.github/workflows/build.yml

@@ -182,7 +182,7 @@ jobs:
     - name: Install dependencies
       run: |
         make build
-        poetry install
+        poetry install -E server -E browser
     - name: Typecheck
       run: |
         make typecheck

docs/tools.rst

@@ -32,6 +32,10 @@ The tools can be grouped into the following categories:
 
   - `Chats`_
 
+- Context management
+
+  - `RAG`_
+
 Shell
 -----
 
@@ -112,34 +116,15 @@ Chats
 Computer
 --------
 
+.. include:: computer-use-warning.rst
+
 .. automodule:: gptme.tools.computer
     :members:
     :noindex:
 
-The computer tool provides direct interaction with the desktop environment through X11, allowing for:
-
-- Keyboard input simulation
-- Mouse control (movement, clicks, dragging)
-- Screen capture with automatic scaling
-- Cursor position tracking
-
-To use the computer tool, see the instructions for :doc:`server`.
-
-Example usage::
-
-    # Type text
-    computer(action="type", text="Hello, World!")
+RAG
+---
 
-    # Move mouse and click
-    computer(action="mouse_move", coordinate=(100, 100))
-    computer(action="left_click")
-
-    # Take screenshot
-    computer(action="screenshot")
-
-    # Send keyboard shortcuts
-    computer(action="key", text="Control_L+c")
-
-The tool automatically handles screen resolution scaling to ensure optimal performance with LLM vision capabilities.
-
-.. include:: computer-use-warning.rst
+.. automodule:: gptme.tools.rag
+    :members:
+    :noindex:

gptme/config.py

@@ -1,6 +1,7 @@
 import logging
 import os
 from dataclasses import dataclass, field
+from functools import lru_cache
 from pathlib import Path
 
 import tomlkit
@@ -41,6 +42,7 @@ class ProjectConfig:
     """Project-level configuration, such as which files to include in the context by default."""
 
     files: list[str] = field(default_factory=list)
+    rag: dict = field(default_factory=dict)
 
 
 ABOUT_ACTIVITYWATCH = """ActivityWatch is a free and open-source automated time-tracker that helps you track how you spend your time on your devices."""
@@ -72,12 +74,12 @@ class ProjectConfig:
 def get_config() -> Config:
     global _config
     if _config is None:
-        _config = load_config()
+        _config = _load_config()
     return _config
 
 
-def load_config() -> Config:
-    config = _load_config()
+def _load_config() -> Config:
+    config = _load_config_doc()
     assert "prompt" in config, "prompt key missing in config"
     assert "env" in config, "env key missing in config"
     prompt = config.pop("prompt")
@@ -87,7 +89,7 @@ def load_config() -> Config:
     return Config(prompt=prompt, env=env)
 
 
-def _load_config() -> tomlkit.TOMLDocument:
+def _load_config_doc() -> tomlkit.TOMLDocument:
     # Check if the config file exists
     if not os.path.exists(config_path):
         # If not, create it and write some default settings
@@ -105,7 +107,7 @@ def _load_config() -> tomlkit.TOMLDocument:
 
 
 def set_config_value(key: str, value: str) -> None:  # pragma: no cover
-    doc: TOMLDocument | Container = _load_config()
+    doc: TOMLDocument | Container = _load_config_doc()
 
     # Set the value
     keypath = key.split(".")
@@ -120,9 +122,10 @@ def set_config_value(key: str, value: str) -> None:  # pragma: no cover
 
     # Reload config
     global _config
-    _config = load_config()
+    _config = _load_config()
 
 
+@lru_cache
 def get_project_config(workspace: Path) -> ProjectConfig | None:
     project_config_paths = [
         p

gptme/init.py

@@ -2,8 +2,9 @@
 from typing import cast
 
 from dotenv import load_dotenv
+from rich.logging import RichHandler
 
-from .config import config_path, load_config, set_config_value
+from .config import config_path, get_config, set_config_value
 from .llm import init_llm
 from .models import (
     PROVIDERS,
@@ -30,7 +31,7 @@ def init(model: str | None, interactive: bool, tool_allowlist: list[str] | None)
     logger.debug("Started")
     load_dotenv()
 
-    config = load_config()
+    config = get_config()
 
     # get from config
     if not model:
@@ -77,7 +78,12 @@ def init(model: str | None, interactive: bool, tool_allowlist: list[str] | None)
 
 def init_logging(verbose):
     # log init
-    logging.basicConfig(level=logging.DEBUG if verbose else logging.INFO)
+    logging.basicConfig(
+        level=logging.DEBUG if verbose else logging.INFO,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=[RichHandler()],
+    )
     # set httpx logging to WARNING
     logging.getLogger("httpx").setLevel(logging.WARNING)
 

gptme/interrupt.py

@@ -2,6 +2,7 @@
 Sets up a KeyboardInterrupt handler to handle Ctrl-C during the chat loop.
 """
 
+import os
 import time
 
 from .util import console
@@ -18,7 +19,10 @@ def handle_keyboard_interrupt(signum, frame):  # pragma: no cover
     global last_interrupt_time
     current_time = time.time()
 
-    if interruptible:
+    # if testing with pytest
+    testing = bool(os.getenv("PYTEST_CURRENT_TEST"))
+
+    if interruptible or testing:
         raise KeyboardInterrupt
 
     # if current_time - last_interrupt_time <= timeout:

gptme/logmanager.py

@@ -306,6 +306,13 @@ def to_dict(self, branches=False) -> dict:
 
 def prepare_messages(msgs: list[Message]) -> list[Message]:
     """Prepares the messages before sending to the LLM."""
+    from .tools._rag_context import _HAS_RAG, enhance_messages  # fmt: skip
+
+    # First enhance messages with context
+    if _HAS_RAG:
+        msgs = enhance_messages(msgs)
+
+    # Then reduce and limit as before
     msgs_reduced = list(reduce_log(msgs))
 
     if len_tokens(msgs) != len_tokens(msgs_reduced):

gptme/prompts.py

@@ -7,7 +7,6 @@
 
 import glob
 import logging
-import os
 import platform
 import subprocess
 from collections.abc import Generator, Iterable
@@ -17,7 +16,7 @@
 from .__version__ import __version__
 from .config import get_config, get_project_config
 from .message import Message
-from .util import document_prompt_function
+from .util import document_prompt_function, get_project_dir
 
 PromptType = Literal["full", "short"]
 
@@ -173,19 +172,12 @@ def prompt_project() -> Generator[Message, None, None]:
     """
     Generate the project-specific prompt based on the current Git repository.
     """
-    config_prompt = get_config().prompt
-    try:
-        projectdir = subprocess.run(
-            ["git", "rev-parse", "--show-toplevel"],
-            capture_output=True,
-            text=True,
-            check=True,
-        ).stdout.strip()
-        project = os.path.basename(projectdir)
-    except subprocess.CalledProcessError:
-        logger.debug("Unable to determine Git repository root.")
+    projectdir = get_project_dir()
+    if not projectdir:
         return
 
+    config_prompt = get_config().prompt
+    project = projectdir.name
     project_info = config_prompt.get("project", {}).get(
         project, "No specific project context provided."
     )

gptme/tools/__init__.py

@@ -19,6 +19,7 @@
 from .tmux import tool as tmux_tool
 from .vision import tool as vision_tool
 from .youtube import tool as youtube_tool
+from .rag import tool as rag_tool
 
 logger = logging.getLogger(__name__)
 
@@ -45,6 +46,7 @@
     screenshot_tool,
     vision_tool,
     computer_tool,
+    rag_tool,
     # python tool is loaded last to ensure all functions are registered
     python_tool,
 ]