embed.py

import os
import json
from openai import OpenAI
import chromadb

embedding_client = OpenAI(
        base_url="https://ai.hackclub.com/proxy/v1",
        api_key="sk-hc-v1-1eba2acb06d4436f9ffba79f9449e6400856e2ddb89f45dda3e36de5f18c3f39"#os.environ.get("HACKCLUB_AI_KEY")
)

chroma_client = chromadb.PersistentClient(path="./db")
collection = chroma_client.get_or_create_collection(name="messages")
with open("messages.json", "r") as f:
    messages = json.load(f)
print(f"found {len(messages)} in messages.json")
newly_embedded = 0
for message in messages:
    msg_id = message["ts"]
    existing = collection.get(ids=[msg_id])
    if existing["ids"]:
        continue 
    response = embedding_client.embeddings.create(
        model="openai/text-embedding-3-small",
        input=message["text"]
        )
    vector = response.data[0].embedding
    collection.add(
        ids=[msg_id],
        embeddings=[vector],
        documents=[message["text"]],
        metadatas=[{
            "user": message["user"],
            "channel": message["channel"],
            "ts": message["ts"]
        }]
        
    )
    newly_embedded += 1
    print(f"Embedded message {newly_embedded}: {message['text'][:50]}...")
print(f"Done! Embedded {newly_embedded} new messages. Total in DB: {collection.count()}")