23 lines
855 B
Python
23 lines
855 B
Python
import os
|
|
from langchain_openai import OpenAIEmbeddings
|
|
from langchain_qdrant import Qdrant
|
|
from langchain_core.documents import Document
|
|
|
|
class KnowledgeBase:
|
|
def __init__(self, collection_name: str = "knowledge", host: str = "localhost", port: int = 6333):
|
|
self.embeddings = OpenAIEmbeddings(
|
|
model="text-embedding-3-small",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
api_key=os.getenv("OPENAI_API_KEY"),
|
|
)
|
|
self.vector_store = Qdrant(
|
|
collection_name=collection_name,
|
|
url=f"http://{host}:{port}",
|
|
embedding_function=self.embeddings,
|
|
)
|
|
|
|
def add_documents(self, documents):
|
|
self.vector_store.add_documents(documents)
|
|
|
|
def similarity_search(self, query: str, k: int = 3):
|
|
return self.vector_store.similarity_search(query, k=k) |