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)