chore: import upstream snapshot with attribution

2026-07-13 12:37:47 +08:00
commit 7653f56fed
1422 changed files with 359026 additions and 0 deletions
@@ -0,0 +1,69 @@
+from config import vector_collection, voyage_client, VOYAGE_MODEL
+from pymongo.operations import SearchIndexModel
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+import time
+
+def get_embedding(data, input_type = "document"):
+  embeddings = voyage_client.embed(
+      data, model = VOYAGE_MODEL, input_type = input_type
+  ).embeddings
+  return embeddings[0]
+
+def ingest_data():
+    loader = PyPDFLoader("https://investors.mongodb.com/node/13176/pdf")
+    data = loader.load()
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=20)
+    documents = text_splitter.split_documents(data)
+    print(f"Successfully split PDF into {len(documents)} chunks.")
+
+    print("Generating embeddings and ingesting documents...")
+    docs_to_insert = []
+    for i, doc in enumerate(documents):
+        embedding = get_embedding(doc.page_content)
+        if embedding:
+            docs_to_insert.append({
+                "text": doc.page_content,
+                "embedding": embedding
+            })
+
+    if docs_to_insert:
+        result = vector_collection.insert_many(docs_to_insert)
+        print(f"Inserted {len(result.inserted_ids)} documents into the collection.")
+    else:
+        print("No documents were inserted. Check embedding generation process.")
+
+    index_name = "vector_index"
+    
+    search_index_model = SearchIndexModel(
+        definition = {
+            "fields": [
+                {
+                    "type": "vector",
+                    "numDimensions": 1024, 
+                    "path": "embedding",
+                    "similarity": "cosine"
+                }
+            ]
+        },
+        name=index_name,
+        type="vectorSearch"
+    )
+    try:
+        vector_collection.create_search_index(model=search_index_model)
+        print(f"Search index '{index_name}' creation initiated.")
+    except Exception as e:
+        print(f"Error creating search index: {e}")
+        return
+
+    print("Polling to check if the index is ready. This may take up to a minute.")
+    predicate=None
+    if predicate is None:
+       predicate = lambda index: index.get("queryable") is True
+
+    while True:
+       indices = list(vector_collection.list_search_indexes(index_name))
+       if len(indices) and predicate(indices[0]):
+          break
+       time.sleep(5)
+    print(index_name + " is ready for querying.")