89 lines
2.6 KiB
Python
89 lines
2.6 KiB
Python
"""
|
|
Knowledge base implementation using Ollama embeddings.
|
|
Provides tools for searching and adding documents.
|
|
"""
|
|
|
|
import numpy as np
|
|
from typing import List
|
|
from langchain.docstore.document import Document
|
|
from embeddings import get_embedding_model
|
|
|
|
class KnowledgeBase:
|
|
"""
|
|
In-memory knowledge base that stores documents and their embeddings.
|
|
"""
|
|
|
|
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:
|
|
"""
|
|
Adds a new document to the knowledge base.
|
|
|
|
Args:
|
|
content (str): The text content to add.
|
|
|
|
Returns:
|
|
str: Confirmation message.
|
|
"""
|
|
doc = Document(page_content=content)
|
|
embedding = self.embedding_model.embed_query(content)
|
|
embedding = np.array(embedding).reshape(1, -1)
|
|
|
|
self.documents.append(doc)
|
|
if self.embeddings.size == 0:
|
|
self.embeddings = embedding
|
|
else:
|
|
self.embeddings = np.vstack([self.embeddings, embedding])
|
|
|
|
return f"Document added. Total documents: {len(self.documents)}."
|
|
|
|
def search_knowledge_base(self, query: str, k: int = 3) -> List[Document]:
|
|
"""
|
|
Searches the knowledge base for the most relevant documents.
|
|
|
|
Args:
|
|
query (str): The search query.
|
|
k (int): Number of top documents to return.
|
|
|
|
Returns:
|
|
List[Document]: List of top matching documents.
|
|
"""
|
|
if not self.documents:
|
|
return []
|
|
|
|
query_embedding = self.embedding_model.embed_query(query)
|
|
query_embedding = np.array(query_embedding).reshape(1, -1)
|
|
|
|
similarities = np.dot(self.embeddings, query_embedding.T).flatten()
|
|
top_indices = similarities.argsort()[-k:][::-1]
|
|
return [self.documents[i] for i in top_indices]
|
|
|
|
# Global knowledge base instance
|
|
kb = KnowledgeBase()
|
|
|
|
def add_to_knowledge_base(content: str) -> str:
|
|
"""
|
|
Tool wrapper for adding content to the knowledge base.
|
|
|
|
Args:
|
|
content (str): Text to add.
|
|
|
|
Returns:
|
|
str: Confirmation message.
|
|
"""
|
|
return kb.add_to_knowledge_base(content)
|
|
|
|
def search_knowledge_base(query: str) -> List[Document]:
|
|
"""
|
|
Tool wrapper for searching the knowledge base.
|
|
|
|
Args:
|
|
query (str): Search query.
|
|
|
|
Returns:
|
|
List[Document]: Matching documents.
|
|
"""
|
|
return kb.search_knowledge_base(query) |