42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
import json
|
|
import os
|
|
from .base import tool
|
|
|
|
KB_FILE = os.path.expanduser("~/.ollama-cli-kb.json")
|
|
|
|
def load_kb():
|
|
if os.path.exists(KB_FILE):
|
|
with open(KB_FILE, 'r') as f:
|
|
return json.load(f)
|
|
return {}
|
|
|
|
def save_kb(kb):
|
|
with open(KB_FILE, 'w') as f:
|
|
json.dump(kb, f, indent=2)
|
|
|
|
@tool(
|
|
name="kb_add",
|
|
description="Add content to the knowledge base",
|
|
parameters={
|
|
"key": {"type": "string", "description": "Search key"},
|
|
"content": {"type": "string", "description": "Information content"}
|
|
}
|
|
)
|
|
def kb_add(key: str, content: str) -> str:
|
|
kb = load_kb()
|
|
kb[key] = content
|
|
save_kb(kb)
|
|
return f"Added '{key}' to knowledge base"
|
|
|
|
@tool(
|
|
name="kb_search",
|
|
description="Search the knowledge base",
|
|
parameters={
|
|
"query": {"type": "string", "description": "Search query"}
|
|
}
|
|
)
|
|
def kb_search(query: str) -> str:
|
|
kb = load_kb()
|
|
results = [f"{k}: {v[:100]}..." for k, v in kb.items() if query.lower() in k.lower() or query.lower() in v.lower()]
|
|
return "\n".join(results) if results else "No matches found"
|