73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
from langchain_qdrant import QdrantVectorStore
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_core.documents import Document
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, VectorParams
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
import uuid
|
|
|
|
COLLECTION_NAME = "knowledge_base"
|
|
EMBEDDING_MODEL = "nomic-embed-text"
|
|
VECTOR_SIZE = 768
|
|
|
|
|
|
def get_embeddings():
|
|
return OllamaEmbeddings(model=EMBEDDING_MODEL)
|
|
|
|
|
|
def get_qdrant_client():
|
|
return QdrantClient(host="localhost", port=6333)
|
|
|
|
|
|
def init_collection(client: QdrantClient):
|
|
collections = [c.name for c in client.get_collections().collections]
|
|
if COLLECTION_NAME not in collections:
|
|
client.create_collection(
|
|
collection_name=COLLECTION_NAME,
|
|
vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
|
|
)
|
|
|
|
|
|
def get_vector_store() -> QdrantVectorStore:
|
|
client = get_qdrant_client()
|
|
init_collection(client)
|
|
embeddings = get_embeddings()
|
|
return QdrantVectorStore(
|
|
client=client,
|
|
collection_name=COLLECTION_NAME,
|
|
embedding=embeddings,
|
|
)
|
|
|
|
|
|
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 = get_vector_store()
|
|
store.add_documents(docs)
|
|
return len(docs)
|
|
|
|
|
|
def search_documents(query: str, max_results: int = 5) -> list[dict]:
|
|
store = get_vector_store()
|
|
results = store.similarity_search_with_relevance_scores(query, k=max_results)
|
|
output = []
|
|
for doc, score in results:
|
|
output.append(
|
|
{
|
|
"content": doc.page_content,
|
|
"metadata": doc.metadata,
|
|
"score": round(score, 4),
|
|
}
|
|
)
|
|
return output |