Update load_book.py

load_book.py

@@ -1,106 +1,76 @@
 import os
-import sys
 import time
-import google.generativeai as genai
 from langchain_community.document_loaders import TextLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_text_splitters import CharacterTextSplitter
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
 from langchain_community.vectorstores import SupabaseVectorStore
 from supabase.client import create_client
 
-class RateLimitedEmbeddings:
-    def __init__(self, api_key, model_name):
-        genai.configure(api_key=api_key)
-        self.model_name = model_name
+# Данные для подключения (берутся из Secrets GitHub)
+supabase_url = os.environ.get("SUPABASE_URL")
+supabase_key = os.environ.get("SUPABASE_SERVICE_KEY")
+google_api_key = os.environ.get("GOOGLE_API_KEY")
 
-    def embed_documents(self, texts):
-        embeddings = []
-        batch_size = 30 
+def load_and_split():
+    print("--- Шаг 1: Загрузка файла book.txt ---")
+    if not os.path.exists("./book.txt"):
+        raise FileNotFoundError("Файл book.txt не найден в корневой директории!")
         
-        for i in range(0, len(texts), batch_size):
-            batch = texts[i : i + batch_size]
-            print(f"--- Обработка батча: {i} из {len(texts)} ---")
-            
-            try:
-                res = genai.embed_content(
-                    model=self.model_name,
-                    content=batch,
-                    task_type="retrieval_document"
-                )
-                embeddings.extend(res['embedding'])
-                time.sleep(5) 
-            except Exception as e:
-                if "429" in str(e):
-                    print("Лимит превышен. Ждем 70 секунд...")
-                    time.sleep(70)
-                    res = genai.embed_content(
-                        model=self.model_name,
-                        content=batch,
-                        task_type="retrieval_document"
-                    )
-                    embeddings.extend(res['embedding'])
-                else:
-                    raise e
-                    
-        return embeddings
+    loader = TextLoader("./book.txt", encoding='utf-8')
+    documents = loader.load()
+    
+    # Настройка нарезки: куски по 1000 символов
+    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+    docs = text_splitter.split_documents(documents)
+    print(f"Успешно: Файл разбит на {len(docs)} фрагментов.")
+    return docs
 
-    def embed_query(self, text):
-        res = genai.embed_content(
-            model=self.model_name,
-            content=text,
-            task_type="retrieval_query"
+def main():
+    try:
+        # 1. Подготовка данных
+        all_docs = load_and_split()
+        
+        # --- МИКРОШАГ: ОГРАНИЧЕНИЕ ДЛЯ ТЕСТА ---
+        # Берем только первые 2 фрагмента для проверки связи и типа данных ID
+        docs = all_docs[:2] 
+        print(f"--- ТЕСТОВЫЙ РЕЖИМ: Работаем только с {len(docs)} фрагментами ---")
+        
+        # 2. Инициализация клиентов
+        print("--- Шаг 2: Инициализация моделей и клиента Supabase ---")
+        supabase = create_client(supabase_url, supabase_key)
+        embeddings = GoogleGenerativeAIEmbeddings(
+            model="models/embedding-001", 
+            google_api_key=google_api_key
         )
-        return res['embedding']
 
-def load_book_to_supabase():
-    supabase_url = os.environ.get("SUPABASE_URL")
-    supabase_key = os.environ.get("SUPABASE_SERVICE_ROLE_KEY")
-    gemini_api_key = os.environ.get("GEMINI_API_KEY")
-
-    if not all([supabase_url, supabase_key, gemini_api_key]):
-        print("ОШИБКА: Проверьте секреты!")
-        sys.exit(1)
-
-    print("--- Шаг 1: Подключение к Supabase ---")
-    supabase = create_client(supabase_url, supabase_key)
-
-    # НОВОЕ: Очистка таблицы перед загрузкой
-    print("--- Шаг 1.1: Очистка старых данных в таблице documents ---")
-    try:
-        # Удаляем все записи, где id не равен 0 (то есть вообще все)
-        supabase.table("documents").delete().neq("id", 0).execute()
+        # 3. Очистка таблицы
+        print("--- Шаг 3: Очистка таблицы documents в базе данных ---")
+        # Мы используем .neq("id", "0"), так как теперь ID — это текст
+        supabase.table("documents").delete().neq("id", "0").execute()
         print("Таблица успешно очищена.")
-    except Exception as e:
-        print(f"Предупреждение при очистке (возможно, таблица уже пуста): {e}")
-
-    print("--- Шаг 2: Загрузка файла книги ---")
-    try:
-        loader = TextLoader("book.txt", encoding="utf-8")
-        documents = loader.load()
-    except Exception as e:
-        print(f"ОШИБКА: {e}")
-        sys.exit(1)
 
-    print("--- Шаг 3: Нарезка текста ---")
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=3000, chunk_overlap=300)
-    docs = text_splitter.split_documents(documents)
-    print(f"Нарезано на {len(docs)} фрагментов.")
+        # 4. Загрузка данных
+        print("--- Шаг 4: Отправка данных в Supabase ---")
+        # Для 2-х фрагментов батчи не нужны, но оставим структуру для будущего
+        batch_size = 30
+        for i in range(0, len(docs), batch_size):
+            batch = docs[i:i + batch_size]
+            print(f"Обработка фрагментов с {i} по {i + len(batch)}...")
+            
+            SupabaseVectorStore.from_documents(
+                batch,
+                embeddings,
+                client=supabase,
+                table_name="documents",
+                query_name="match_documents"
+            )
+            print(f"Успешно записано в базу.")
 
-    print("--- Шаг 4: Инициализация эмбеддингов ---")
-    custom_embeddings = RateLimitedEmbeddings(gemini_api_key, "models/gemini-embedding-001")
+        print("--- ФИНАЛ: Тестовая загрузка завершена успешно! ---")
+        print("Проверьте таблицу 'documents' в интерфейсе Supabase.")
 
-    print("--- Шаг 5: Загрузка в базу ---")
-    try:
-        SupabaseVectorStore.from_documents(
-            docs,
-            custom_embeddings,
-            client=supabase,
-            table_name="documents",
-            query_name="match_documents",
-        )
-        print("--- УСПЕХ: Все данные загружены в Supabase! ---")
     except Exception as e:
-        print(f"КРИТИЧЕСКАЯ ОШИБКА: {e}")
-        sys.exit(1)
+        print(f"!!! КРИТИЧЕСКАЯ ОШИБКА: {e}")
 
 if __name__ == "__main__":
-    load_book_to_supabase()
+    main()