This commit is contained in:
+64
-106
@@ -1,131 +1,89 @@
|
||||
"""
|
||||
Knowledge Base Tool for the Agent.
|
||||
|
||||
This module implements a simple file‑based knowledge base that can be
|
||||
used by the agent and accessed via the CLI. The knowledge base is
|
||||
stored as a JSON file (`knowledge_base.json`) in the same directory
|
||||
as this module. Each entry is a key/value pair where the key is a
|
||||
string and the value is any JSON‑serialisable object.
|
||||
|
||||
The class provides three public methods:
|
||||
|
||||
* add_entry(key, value) – Add or update an entry.
|
||||
* query_entry(key) – Retrieve the value for a key.
|
||||
* delete_entry(key) – Remove an entry.
|
||||
|
||||
The tool is intentionally lightweight and does not depend on any
|
||||
external libraries beyond the Python standard library.
|
||||
Knowledge base implementation using Ollama embeddings.
|
||||
Provides tools for searching and adding documents.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
import numpy as np
|
||||
from typing import List
|
||||
from langchain.docstore.document import Document
|
||||
from embeddings import get_embedding_model
|
||||
|
||||
|
||||
class KnowledgeBaseTool:
|
||||
class KnowledgeBase:
|
||||
"""
|
||||
A simple file‑based knowledge base tool.
|
||||
In-memory knowledge base that stores documents and their embeddings.
|
||||
"""
|
||||
|
||||
def __init__(self, storage_path: Optional[Path] = None) -> None:
|
||||
def __init__(self):
|
||||
self.embedding_model = get_embedding_model()
|
||||
self.documents: List[Document] = []
|
||||
self.embeddings: np.ndarray = np.empty((0, self.embedding_model.get_sentence_embedding_dimension()))
|
||||
|
||||
def add_to_knowledge_base(self, content: str) -> str:
|
||||
"""
|
||||
Initialise the knowledge base.
|
||||
Adds a new document to the knowledge base.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
storage_path : Optional[Path]
|
||||
Path to the JSON file used for storage. If not provided,
|
||||
a file named ``knowledge_base.json`` in the same directory
|
||||
as this module is used.
|
||||
Args:
|
||||
content (str): The text content to add.
|
||||
|
||||
Returns:
|
||||
str: Confirmation message.
|
||||
"""
|
||||
if storage_path is None:
|
||||
storage_path = Path(__file__).parent / "knowledge_base.json"
|
||||
self.storage_path = storage_path
|
||||
# Ensure the storage file exists
|
||||
if not self.storage_path.exists():
|
||||
self.storage_path.write_text("{}")
|
||||
doc = Document(page_content=content)
|
||||
embedding = self.embedding_model.embed_query(content)
|
||||
embedding = np.array(embedding).reshape(1, -1)
|
||||
|
||||
def _load(self) -> Dict[str, Any]:
|
||||
"""Load the knowledge base from disk."""
|
||||
try:
|
||||
data = json.loads(self.storage_path.read_text())
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Knowledge base file corrupted: not a dict")
|
||||
return data
|
||||
except json.JSONDecodeError:
|
||||
raise ValueError("Knowledge base file corrupted: invalid JSON")
|
||||
self.documents.append(doc)
|
||||
if self.embeddings.size == 0:
|
||||
self.embeddings = embedding
|
||||
else:
|
||||
self.embeddings = np.vstack([self.embeddings, embedding])
|
||||
|
||||
def _save(self, data: Dict[str, Any]) -> None:
|
||||
"""Persist the knowledge base to disk."""
|
||||
self.storage_path.write_text(json.dumps(data, indent=2, ensure_ascii=False))
|
||||
return f"Document added. Total documents: {len(self.documents)}."
|
||||
|
||||
def add_entry(self, key: str, value: Any) -> None:
|
||||
def search_knowledge_base(self, query: str, k: int = 3) -> List[Document]:
|
||||
"""
|
||||
Add or update an entry in the knowledge base.
|
||||
Searches the knowledge base for the most relevant documents.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
key : str
|
||||
The key under which the value will be stored.
|
||||
value : Any
|
||||
The value to store. Must be JSON‑serialisable.
|
||||
Args:
|
||||
query (str): The search query.
|
||||
k (int): Number of top documents to return.
|
||||
|
||||
Returns:
|
||||
List[Document]: List of top matching documents.
|
||||
"""
|
||||
data = self._load()
|
||||
data[key] = value
|
||||
self._save(data)
|
||||
if not self.documents:
|
||||
return []
|
||||
|
||||
def query_entry(self, key: str) -> Any:
|
||||
"""
|
||||
Retrieve the value for a given key.
|
||||
query_embedding = self.embedding_model.embed_query(query)
|
||||
query_embedding = np.array(query_embedding).reshape(1, -1)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
key : str
|
||||
The key to look up.
|
||||
similarities = np.dot(self.embeddings, query_embedding.T).flatten()
|
||||
top_indices = similarities.argsort()[-k:][::-1]
|
||||
return [self.documents[i] for i in top_indices]
|
||||
|
||||
Returns
|
||||
-------
|
||||
Any
|
||||
The stored value.
|
||||
# Global knowledge base instance
|
||||
kb = KnowledgeBase()
|
||||
|
||||
Raises
|
||||
------
|
||||
KeyError
|
||||
If the key does not exist.
|
||||
"""
|
||||
data = self._load()
|
||||
if key not in data:
|
||||
raise KeyError(f"Key '{key}' not found in knowledge base.")
|
||||
return data[key]
|
||||
def add_to_knowledge_base(content: str) -> str:
|
||||
"""
|
||||
Tool wrapper for adding content to the knowledge base.
|
||||
|
||||
def delete_entry(self, key: str) -> None:
|
||||
"""
|
||||
Delete an entry from the knowledge base.
|
||||
Args:
|
||||
content (str): Text to add.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
key : str
|
||||
The key to delete.
|
||||
Returns:
|
||||
str: Confirmation message.
|
||||
"""
|
||||
return kb.add_to_knowledge_base(content)
|
||||
|
||||
Raises
|
||||
------
|
||||
KeyError
|
||||
If the key does not exist.
|
||||
"""
|
||||
data = self._load()
|
||||
if key not in data:
|
||||
raise KeyError(f"Key '{key}' not found in knowledge base.")
|
||||
del data[key]
|
||||
self._save(data)
|
||||
def search_knowledge_base(query: str) -> List[Document]:
|
||||
"""
|
||||
Tool wrapper for searching the knowledge base.
|
||||
|
||||
def list_entries(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Return a copy of all entries in the knowledge base.
|
||||
Args:
|
||||
query (str): Search query.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Dict[str, Any]
|
||||
All key/value pairs.
|
||||
"""
|
||||
return self._load()
|
||||
Returns:
|
||||
List[Document]: Matching documents.
|
||||
"""
|
||||
return kb.search_knowledge_base(query)
|
||||
Reference in New Issue
Block a user