KorChat / database.py
Jacek Zadro偶ny
Pe艂ny dost臋p do PDF i stron PFRON, embedowanie wsadowe i kilka innych zmian.
4e4c288
import os
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document
# Sprawdzenie, czy klucz API jest ustawiony
if not os.getenv("OPENAI_API_KEY"):
raise ValueError("Klucz OPENAI_API_KEY nie jest ustawiony w zmiennych 艣rodowiskowych. Ustaw go, aby kontynuowa膰.")
# Model do tworzenia wektor贸w (embeddings) - ten sam, co w hr_assistant.py
EMBEDDINGS = OpenAIEmbeddings(model="text-embedding-3-small")
class FaissCollectionWrapper:
"""
Klasa-adapter do pracy z baz膮 FAISS w pami臋ci.
"""
def __init__(self, vector_store=None):
if vector_store is None:
# Utw贸rz pust膮 baz臋 FAISS z minimaln膮 zawarto艣ci膮
self._vector_store = FAISS.from_texts(["placeholder"], EMBEDDINGS)
else:
self._vector_store = vector_store
def add(self, documents, metadatas, ids):
"""
Dodaje dokumenty do bazy FAISS (tylko w pami臋ci, bez zapisu na dysk).
"""
docs_to_add = []
for i, content in enumerate(documents):
docs_to_add.append(Document(page_content=content, metadata=metadatas[i]))
if docs_to_add:
new_docs_vectorstore = FAISS.from_documents(docs_to_add, EMBEDDINGS)
self._vector_store.merge_from(new_docs_vectorstore)
print(f"Dodano {len(docs_to_add)} dokument贸w do bazy w pami臋ci.")
def get_collection():
"""
Tworzy now膮, pust膮 baz臋 FAISS w pami臋ci.
"""
print("Tworzenie nowej bazy danych FAISS w pami臋ci...")
# Tworzymy nowy wrapper, kt贸ry automatycznie utworzy pust膮 baz臋
return FaissCollectionWrapper()
if __name__ == '__main__':
print("Testowanie modu艂u database.py...")
collection = get_collection()
print("Pomy艣lnie zainicjalizowano baz臋 danych FAISS.")