2ff4056d89
New optional module (D24): explains the collected findings in plain language,
contacted ONLY on an explicit user action (never automatic).
- core/ai.py: provider chosen explicitly (no default) — ollama (local) or claude
(Anthropic Messages API via stdlib urllib; key in keyring). Grounded prompt;
HTTP error parsing; one-shot (no thinking/caching — snappy).
- core/ai_knowledge.py: curated reference KB (Xid/SMART/Proton/tunables),
exact keyword/code match ("RAG-lite", no embeddings) injected into the prompt —
lifts local models, sharpens Claude. No fine-tuning.
- config: ai_provider/ai_model/ai_endpoint + keyring-backed AI key (generalized
the token keyring helpers).
- GUI: Settings → AI assistant (provider radios, model/endpoint/key, Save/Test);
"Explain with AI" button on the diagnostic dialog (consent prompt for cloud).
- CLI: `rigdoctor ai status|test|explain`.
- Docs: D24, SPEC/MODULES/ROADMAP (Phase 7); tests for providers/grounding/parse.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
102 lines
4.3 KiB
Python
102 lines
4.3 KiB
Python
"""Tests for the M14 AI assistant: provider selection, grounding, parsing (no network)."""
|
|
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from rigdoctor.core import ai, ai_knowledge
|
|
|
|
|
|
class KnowledgeTests(unittest.TestCase):
|
|
def test_matches_xid_and_smart(self):
|
|
facts = ai_knowledge.relevant("Kernel: NVRM: Xid 79: GPU has fallen off the bus")
|
|
self.assertTrue(any("fallen off the bus" in f for f in facts))
|
|
|
|
def test_matches_smart_pending(self):
|
|
facts = ai_knowledge.relevant("SMART 197 Current_Pending_Sector = 8")
|
|
self.assertTrue(any("Pending Sector" in f for f in facts))
|
|
|
|
def test_no_match_returns_empty(self):
|
|
self.assertEqual(ai_knowledge.relevant("everything is fine"), [])
|
|
|
|
|
|
class ConfigStateTests(unittest.TestCase):
|
|
def _cfg(self, **over):
|
|
base = {"ai_provider": "", "ai_model": "", "ai_endpoint": "http://localhost:11434"}
|
|
base.update(over)
|
|
return base
|
|
|
|
def test_unconfigured_by_default(self):
|
|
with mock.patch.object(ai.config, "load_config", return_value=self._cfg()):
|
|
self.assertFalse(ai.is_configured())
|
|
|
|
def test_ollama_needs_model(self):
|
|
with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="ollama")):
|
|
self.assertFalse(ai.is_configured())
|
|
with mock.patch.object(ai.config, "load_config",
|
|
return_value=self._cfg(ai_provider="ollama", ai_model="llama3.1")):
|
|
self.assertTrue(ai.is_configured())
|
|
|
|
def test_claude_needs_key(self):
|
|
with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="claude")), \
|
|
mock.patch.object(ai.config, "load_ai_key", return_value=None):
|
|
self.assertFalse(ai.is_configured())
|
|
with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="claude")), \
|
|
mock.patch.object(ai.config, "load_ai_key", return_value="sk-ant-x"):
|
|
self.assertTrue(ai.is_configured())
|
|
|
|
def test_claude_default_model(self):
|
|
with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="claude")):
|
|
self.assertEqual(ai.model(), ai.CLAUDE_DEFAULT_MODEL)
|
|
|
|
|
|
class PromptTests(unittest.TestCase):
|
|
def test_build_prompt_includes_facts_and_findings(self):
|
|
prompt = ai.build_prompt("Xid 79: GPU has fallen off the bus")
|
|
self.assertIn("Reference facts", prompt)
|
|
self.assertIn("Collected findings", prompt)
|
|
self.assertIn("fallen off the bus", prompt)
|
|
|
|
def test_format_findings(self):
|
|
class F:
|
|
severity, category, title, detail = "warn", "GPU", "Hot", "92C"
|
|
text = ai.format_findings([F()])
|
|
self.assertIn("[WARN] GPU: Hot — 92C", text)
|
|
|
|
|
|
class ExplainTests(unittest.TestCase):
|
|
def _cfg(self, **over):
|
|
base = {"ai_provider": "", "ai_model": "", "ai_endpoint": "http://localhost:11434"}
|
|
base.update(over)
|
|
return base
|
|
|
|
def test_no_provider(self):
|
|
with mock.patch.object(ai.config, "load_config", return_value=self._cfg()):
|
|
ok, msg = ai.explain("x")
|
|
self.assertFalse(ok)
|
|
self.assertIn("No AI provider", msg)
|
|
|
|
def test_ollama_parses_response(self):
|
|
with mock.patch.object(ai.config, "load_config",
|
|
return_value=self._cfg(ai_provider="ollama", ai_model="llama3.1")), \
|
|
mock.patch.object(ai, "_post", return_value={"response": "It's the PSU."}) as post:
|
|
ok, msg = ai.explain("Xid 79")
|
|
self.assertTrue(ok)
|
|
self.assertEqual(msg, "It's the PSU.")
|
|
self.assertIn("/api/generate", post.call_args[0][0])
|
|
|
|
def test_claude_parses_content_blocks(self):
|
|
with mock.patch.object(ai.config, "load_config", return_value=self._cfg(ai_provider="claude")), \
|
|
mock.patch.object(ai.config, "load_ai_key", return_value="sk-ant-x"), \
|
|
mock.patch.object(ai, "_post", return_value={"content": [
|
|
{"type": "text", "text": "Likely a failing disk."}]}) as post:
|
|
ok, msg = ai.explain("SMART 197")
|
|
self.assertTrue(ok)
|
|
self.assertEqual(msg, "Likely a failing disk.")
|
|
headers = post.call_args[0][2]
|
|
self.assertEqual(headers["anthropic-version"], ai.ANTHROPIC_VERSION)
|
|
self.assertEqual(headers["x-api-key"], "sk-ant-x")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|