diff --git a/rag_store.py b/rag_store.py new file mode 100644 index 0000000..36b31f5 --- /dev/null +++ b/rag_store.py @@ -0,0 +1,23 @@ +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) \ No newline at end of file