56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
from langchain_community.vectorstores import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain.schema import Document
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
|
|
COLLECTION_NAME = "rag_collection"
|
|
EMBEDDING_MODEL = "nomic-embed-text"
|
|
CHROMA_PERSIST_DIR = "./chroma_db"
|
|
|
|
|
|
class ChromaStore:
|
|
def __init__(self, collection_name: str = COLLECTION_NAME):
|
|
self.client = Chroma(
|
|
embedding_function=OllamaEmbeddings(model=EMBEDDING_MODEL),
|
|
collection_name=collection_name,
|
|
persist_directory=CHROMA_PERSIST_DIR,
|
|
)
|
|
|
|
def add_documents(self, docs: list[Document]) -> None:
|
|
self.client.add_documents(docs)
|
|
|
|
def search(self, query: str, limit: int = 5) -> list[Document]:
|
|
return self.client.similarity_search(query, k=limit)
|
|
|
|
|
|
def add_documents(content: str, title: str) -> int:
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=500,
|
|
chunk_overlap=50,
|
|
separators=["\n\n", "\n", ".", " ", ""],
|
|
)
|
|
chunks = splitter.split_text(content)
|
|
docs = [
|
|
Document(
|
|
page_content=chunk,
|
|
metadata={"title": title, "chunk_index": i, "source": title},
|
|
)
|
|
for i, chunk in enumerate(chunks)
|
|
]
|
|
store = ChromaStore()
|
|
store.add_documents(docs)
|
|
return len(docs)
|
|
|
|
|
|
def search_documents(query: str, max_results: int = 5) -> list[dict]:
|
|
store = ChromaStore()
|
|
results = store.search(query, limit=max_results)
|
|
output = []
|
|
for doc in results:
|
|
output.append(
|
|
{
|
|
"content": doc.page_content,
|
|
"metadata": doc.metadata,
|
|
}
|
|
)
|
|
return output |