fix: fixes after RAG improvement PR

ErikBjare · Feb 5, 2025 · 486d985 · 486d985
1 parent f4b6f02
commit 486d985
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/docs/conf.py b/docs/conf.py
@@ -119,6 +119,7 @@ def setup(app):
     ("py:class", "gptme.tools.python.T"),
     ("py:class", "threading.Thread"),
     ("py:class", "gptme.tools.computer.ScalingSource"),
+    ("py:class", "gptme.config.RagConfig"),
 ]
 
 # -- Options for HTML output -------------------------------------------------

diff --git a/tests/test_tools_rag.py b/tests/test_tools_rag.py
@@ -3,6 +3,7 @@
 from unittest.mock import patch
 
 import pytest
+from gptme.config import RagConfig
 from gptme.message import Message
 from gptme.tools.rag import _has_gptme_rag, rag_enhance_messages
 
@@ -42,7 +43,7 @@ def test_enhance_messages_disabled():
         patch("subprocess.run", return_value=type("Proc", (), {"returncode": 0})),
         patch("gptme.tools.rag.get_project_config") as mock_config,
     ):
-        mock_config.return_value.rag = {"enabled": False}
+        mock_config.return_value.rag = RagConfig(enabled=False)
         messages = [
             Message("user", "Tell me about Python"),
             Message("assistant", "Python is a programming language"),