Indexe qualquer coisa, pesquise tudo: Vector Search escalável com IA replicada, MongoDB e Hookdeck

1
git clone https://github.com/hookdeck/index-all-the-things.git

1
poetry shell

1
poetry install

1
cp .env-example .env

1
import httpx
2
import re
3
import hashlib
4
import os
5
6
from config import Config
7
8
headers = {
9
    "Authorization": f"Bearer {Config.HOOKDECK_PROJECT_API_KEY}",
10
    "Content-Type": "application/json",
11
}

1
def create_connection(payload):
2
    response = httpx.request(
3
        "PUT",
4
        "https://api.hookdeck.com/latest/connections",
5
        headers=headers,
6
        json=payload,
7
    )
8
    data = response.json()
9
10
    if response.status_code != 200:
11
        raise Exception(f"Failed to create connection: {data}")
12
13
    return data

1
replicate_api_queue_api_key = hashlib.sha256(os.urandom(32)).hexdigest()
2
replicate_api_queue = {
3
    "name": "replicate-api-queue",
4
    "source": {
5
        "name": "replicate-api-queue",
6
        "verification": {
7
            "type": "API_KEY",
8
            "configs": {
9
                "header_key": Config.HOOKDECK_QUEUE_API_KEY_HEADER_NAME,
10
                "api_key": replicate_api_queue_api_key,
11
            },
12
        },
13
    },
14
    "rules": [
15
        {
16
            "type": "retry",
17
            "strategy": "exponential",
18
            "count": 5,
19
            "interval": 30000,
20
            "response_status_codes": ["429", "500"],
21
        }
22
    ],
23
    "destination": {
24
        "name": "replicate-api",
25
        "url": "https://api.replicate.com/v1/",
26
        "auth_method": {
27
            "type": "BEARER_TOKEN",
28
            "config": {
29
                "token": Config.REPLICATE_API_TOKEN,
30
            },
31
        },
32
    },
33
}
34
35
replicate_api_connection = create_connection(replicate_api_queue)

1
replicate_audio = {
2
    "name": "replicate-audio",
3
    "source": {
4
        "name": "replicate-audio",
5
        "verification": {
6
            "type": "REPLICATE",
7
            "configs": {
8
                "webhook_secret_key": Config.REPLICATE_WEBHOOKS_SECRET,
9
            },
10
        },
11
    },
12
    "rules": [
13
        {
14
            "type": "retry",
15
            "count": 5,
16
            "interval": 30000,
17
            "strategy": "exponential",
18
            "response_status_codes": ["!404"],
19
        }
20
    ],
21
    "destination": {
22
        "name": "cli-replicate-audio",
23
        "cli_path": "/webhooks/audio",
24
    },
25
}
26
27
replicate_audio_connection = create_connection(replicate_audio)

1
replicate_embedding = {
2
    "name": "replicate-embedding",
3
    "source": {
4
        "name": "replicate-embedding",
5
        "verification": {
6
            "type": "REPLICATE",
7
            "configs": {
8
                "webhook_secret_key": Config.REPLICATE_WEBHOOKS_SECRET,
9
            },
10
        },
11
    },
12
    "rules": [
13
        {
14
            "type": "retry",
15
            "count": 5,
16
            "interval": 30000,
17
            "strategy": "exponential",
18
            "response_status_codes": ["!200", "!404"],
19
        }
20
    ],
21
    "destination": {
22
        "name": "cli-replicate-embedding",
23
        "cli_path": "/webhooks/embedding",
24
    },
25
}
26
27
replicate_embedding_connection = create_connection(replicate_embedding)

1
# Update .env
2
with open(".env", "r") as file:
3
    env_content = file.read()
4
5
replicate_api_connection_url = replicate_api_connection["source"]["url"]
6
audio_webhook_url = replicate_audio_connection["source"]["url"]
7
embedding_webhook_url = replicate_embedding_connection["source"]["url"]
8
9
# Replace the .env URLs in the .env content
10
env_content = re.sub(
11
    r"HOOKDECK_REPLICATE_API_QUEUE_API_KEY=.*",
12
    f"HOOKDECK_REPLICATE_API_QUEUE_API_KEY={replicate_api_queue_api_key}",
13
    env_content,
14
)
15
env_content = re.sub(
16
    r"HOOKDECK_REPLICATE_API_QUEUE_URL=.*",
17
    f"HOOKDECK_REPLICATE_API_QUEUE_URL={replicate_api_connection_url}",
18
    env_content,
19
)
20
env_content = re.sub(
21
    r"AUDIO_WEBHOOK_URL=.*", f"AUDIO_WEBHOOK_URL={audio_webhook_url}", env_content
22
)
23
env_content = re.sub(
24
    r"EMBEDDINGS_WEBHOOK_URL=.*",
25
    f"EMBEDDINGS_WEBHOOK_URL={embedding_webhook_url}",
26
    env_content,
27
)
28
29
with open(".env", "w") as file:
30
    file.write(env_content)
31
32
print("Connections created successfully!")

1
poetry run python create-hookdeck-connections.py

1
from allthethings.mongo import Database
2
from pymongo.operations import SearchIndexModel
3
4
database = Database()
5
collection = database.get_collection()

1
if collection.name not in collection.database.list_collection_names():
2
    print("Creating empty collection so indexes can be created.")
3
    collection.database.create_collection(collection.name)

1
def create_or_update_search_index(index_name, index_definition, index_type):
2
    indexes = list(collection.list_search_indexes(index_name))
3
4
    if len(indexes) == 0:
5
        print(f'Creating search index: "{index_name}"')
6
        index_model = SearchIndexModel(
7
            definition=index_definition,
8
            name=index_name,
9
            type=index_type,
10
        )
11
        collection.create_search_index(model=index_model)
12
13
    else:
14
        print(f'Search index "{index_name}" already exists. Updating.')
15
        collection.update_search_index(name=index_name, definition=index_definition)

1
vector_result = create_or_update_search_index(
2
    "vector_index",
3
    {
4
        "fields": [
5
            {
6
                "type": "vector",
7
                "path": "embedding",
8
                "numDimensions": 768,
9
                "similarity": "euclidean",
10
            }
11
        ]
12
    },
13
    "vectorSearch",
14
)

1
create_or_update_search_index(
2
    "url_index",
3
    {
4
        "mappings": {
5
            "fields": {
6
                "url": {
7
                    "type": "string",
8
                },
9
            },
10
        }
11
    },
12
    "search",
13
)
14
15
print("Indexes created successfully!")

1
poetry run python create-indexes.py

1
poetry run python -m flask --app app --debug run

1
hookdeck listen 5000 '*'

1
Listening for events on Sources that have Connections with CLI Destinations
2
3
Dashboard
4
👉 Inspect and replay events: https://dashboard.hookdeck.com?team_id=tm_{id}
5
6
Sources
7
🔌 replicate-embedding URL: https://hkdk.events/{id}
8
🔌 replicate-audio URL: https://hkdk.events/{id}
9
10
Connections
11
replicate-embedding -> replicate-embedding forwarding to /webhooks/embedding
12
replicate-audio -> replicate-audio forwarding to /webhooks/audio
13
14
> Ready! (^C to quit)

1
@app.route("/process", methods=["POST"])
2
def process():
3
    url = request.form["url"]
4
5
    parsed_url = urlparse(url)
6
    if not all([parsed_url.scheme, parsed_url.netloc]):
7
        flash("Invalid URL")
8
        return redirect(url_for("index"))

1
    database = Database()
2
    collection = database.get_collection()
3
4
    exists = collection.find_one({"url": url})
5
6
    if exists is not None:
7
        flash("URL has already been indexed")
8
        return redirect(url_for("index"))

1
    req = urllib.request.Request(url, method="HEAD")
2
    fetch = urllib.request.urlopen(req)
3
4
    if fetch.status != 200:
5
        flash("URL is not reachable")
6
        return redirect(url_for("index"))

1
    content_length = fetch.headers["Content-Length"]
2
    content_type = fetch.headers["Content-Type"]

1
    processor = get_asset_processor(content_type)
2
3
    if processor is None:
4
        flash('Unsupported content type "' + content_type + '"')
5
        return redirect(url_for("index"))

1
def get_asset_processor(
2
    content_type,
3
):
4
    if "audio/" in content_type:
5
        return AudioProcessor()
6
    elif "video/" in content_type:
7
        return None
8
    elif "image/" in content_type:
9
        return None
10
    else:
11
        return None

1
    asset = collection.insert_one(
2
        {
3
            "url": url,
4
            "content_type": content_type,
5
            "content_length": content_length,
6
            "status": "SUBMITTED",
7
        }
8
    )

1
    try:
2
        response = processor.process(asset.inserted_id, url)
3
    except Exception as e:
4
        app.logger.error("Error processing asset: %s", e)
5
        collection.update_one(
6
            filter={"url": url},
7
            update={
8
                "$set": {
9
                    "status": "PROCESSING_ERROR",
10
                    "error": str(e),
11
                }
12
            },
13
        )
14
        flash("Error processing asset")
15
        return redirect(url_for("index"))

1
import httpx
2
from config import Config
3
4
...
5
6
class AudioProcessor:
7
    def process(self, id, url):
8
        input = {
9
            "audio": url,
10
            "model": "large-v3",
11
            "language": "auto",
12
            "translate": False,
13
            "temperature": 0,
14
            "transcription": "plain text",
15
            "suppress_tokens": "-1",
16
            "logprob_threshold": -1,
17
            "no_speech_threshold": 0.6,
18
            "condition_on_previous_text": True,
19
            "compression_ratio_threshold": 2.4,
20
            "temperature_increment_on_fallback": 0.2,
21
        }
22
23
        payload = {
24
            "version": "cdd97b257f93cb89dede1c7584e3f3dfc969571b357dbcee08e793740bedd854",
25
            "input": input,
26
            "webhook": f"{Config.AUDIO_WEBHOOK_URL}/{id}",
27
            "webhook_events_filter": ["completed"],
28
        }
29
30
        response = httpx.request(
31
            "POST",
32
            f"{Config.HOOKDECK_REPLICATE_API_QUEUE_URL}/predictions",
33
            headers=Config.HOOKDECK_QUEUE_AUTH_HEADERS,
34
            json=payload,
35
        )
36
37
        return response.json()

1
    collection.update_one(
2
        filter={"url": url},
3
        update={
4
            "$set": {
5
                "status": "PROCESSING",
6
                "processor_response": response,
7
            }
8
        },
9
    )

1
    flash(
2
        message="Processing: " + url + " with content type: " + content_type,
3
        category="success",
4
    )
5
6
    return redirect(url_for("index"))

1
@app.route("/webhooks/audio/<id>", methods=["POST"])
2
def webhook_audio(id):
3
    if not verify_webhook(request):
4
        app.logger.error("Webhook signature verification failed")
5
        return jsonify({"error": "Webhook signature verification failed"}), 401
6
7
    payload = request.json
8
    app.logger.info("Audio payload received for id %s", id)
9
    app.logger.debug(payload)

1
def verify_webhook(request):
2
    if Config.HOOKDECK_WEBHOOK_SECRET is None:
3
        app.logger.error("No HOOKDECK_WEBHOOK_SECRET found.")
4
        return False
5
6
    hmac_header = request.headers.get("x-hookdeck-signature")
7
8
    hash = base64.b64encode(
9
        hmac.new(
10
            Config.HOOKDECK_WEBHOOK_SECRET.encode(), request.data, hashlib.sha256
11
        ).digest()
12
    ).decode()
13
14
    verified = hash == hmac_header
15
    app.logger.debug("Webhook signature verification: %s", verified)
16
    return verified

1
    database = Database()
2
    collection = database.get_collection()
3
4
    status = (
5
        "PROCESSING_ERROR" if "error" in payload and payload["error"] else "PROCESSED"
6
    )

1
    result = collection.find_one_and_update(
2
        filter={"_id": ObjectId(id)},
3
        update={
4
            "$set": {
5
                "status": status,
6
                "text": payload["output"]["transcription"],
7
                "replicate_response": payload,
8
            }
9
        },
10
        return_document=True,
11
    )

1
    if result is None:
2
        app.logger.error(
3
            "No document found for id %s to add audio transcript", payload["id"]
4
        )
5
        return jsonify({"error": "No document found to add audio transcript"}), 404

1
    app.logger.info("Transcription updated")
2
    app.logger.debug(result)
3
4
    request_embeddings(id)
5
6
    return "OK"

1
def request_embeddings(id):
2
    app.logger.info("Requesting embeddings for %s", id)
3
4
    database = Database()
5
    collection = database.get_collection()
6
7
    asset = collection.find_one({"_id": id})
8
9
    if asset is None:
10
        raise RuntimeError("Asset not found")
11
12
    if asset["status"] != "PROCESSED":
13
        raise RuntimeError("Asset has not been processed")

1
    generator = AsyncEmbeddingsGenerator()
2
3
    try:
4
        response = generator.generate(id, asset["text"])
5
    except Exception as e:
6
        app.logger.error("Error generating embeddings for %s: %s", id, e)
7
        raise

1
import httpx
2
from config import Config
3
4
5
class AsyncEmbeddingsGenerator:
6
    def generate(self, id, text):
7
        payload = {
8
            "version": "b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305",
9
            "input": {"text": text},
10
            "webhook": f"{Config.EMBEDDINGS_WEBHOOK_URL}/{id}",
11
            "webhook_events_filter": ["completed"],
12
        }
13
14
        response = httpx.request(
15
            "POST",
16
            f"{Config.HOOKDECK_REPLICATE_API_QUEUE_URL}/predictions",
17
            headers=Config.HOOKDECK_QUEUE_AUTH_HEADERS,
18
            json=payload,
19
        )
20
21
        return response.json()

1
    collection.update_one(
2
        filter={"_id": ObjectId(id)},
3
        update={
4
            "$set": {
5
                "status": "GENERATING_EMBEDDINGS",
6
                "generator_response": response,
7
            }
8
        },
9
    )

1
@app.route("/webhooks/audio/<id>", methods=["POST"])
2
def webhook_audio(id):
3
        if not verify_webhook(request):
4
        app.logger.error("Webhook signature verification failed")
5
        return jsonify({"error": "Webhook signature verification failed"}), 401
6
7
    payload = request.json
8
    app.logger.info("Audio payload received for id %s", id)
9
    app.logger.debug(payload)

1
    status = (
2
        "EMBEDDINGS_ERROR" if "error" in payload and payload["error"] else "SEARCHABLE"
3
    )

1
    embedding = payload["output"][0]["embedding"]
2
3
    database = Database()
4
    collection = database.get_collection()
5
6
    result = collection.update_one(
7
        filter={"_id": ObjectId(id)},
8
        update={
9
            "$set": {
10
                "status": status,
11
                "embedding": embedding,
12
                "replicate_embeddings_response": payload,
13
            }
14
        },
15
    )

1
    if result.matched_count == 0:
2
        app.logger.error(
3
            "No document found for id %s to update embedding", payload["id"]
4
        )
5
        return jsonify({"error": "No document found to update embedding"}), 404
6
7
    return "OK"

1
@app.route("/search", methods=["GET"])
2
def search():
3
    query = request.args.get("query")
4
    if query is None:
5
        return render_template("search.html", results=[])
6
7
    app.logger.info("Query submitted")
8
    app.logger.debug(query)
9
10
    results = query_vector_search(query)
11
12
    results = format_results(results)
13
14
    app.logger.debug("Formatted search results", results)
15
16
    return render_template("search.html", results=results, query=query)

1
def query_vector_search(q):
2
    generator = SyncEmbeddingsGenerator()
3
4
    try:
5
        generator_response = generator.generate(q)
6
        app.logger.debug(generator_response)
7
    except Exception as e:
8
        app.logger.error("Error generating embeddings: %s", e)
9
        return None
10
11
    if generator_response["status"] != "completed":
12
        app.logger.debug("Embeddings generation timed out")
13
        return None
14
15
    query_embedding = generator_response["output"][0]["embedding"]

1
class SyncEmbeddingsGenerator:
2
    def generate(self, text):
3
        payload = {
4
            "version": "b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305",
5
            "input": {"text": text},
6
        }
7
8
        response = httpx.request(
9
            "POST",
10
            "https://api.replicate.com/v1/predictions",
11
            headers={**Config.REPLICATE_API_AUTH_HEADERS, "Prefer": "wait"},
12
            json=payload,
13
            timeout=60,
14
        )
15
16
        return response.json()

1
    ...
2
3
    query_embedding = generate_response[0]["embedding"]
4
5
    vs_query = {
6
        "index": "vector_index",
7
        "path": "embedding",
8
        "queryVector": query_embedding,
9
        "numCandidates": 100,
10
        "limit": 10,
11
    }
12
13
    new_search_query = {"$vectorSearch": vs_query}
14
15
    app.logger.info("Vector search query created")
16
    app.logger.debug(new_search_query)

1
    project = {
2
        "$project": {
3
            "score": {"$meta": "vectorSearchScore"},
4
            "_id": 0,
5
            "url": 1,
6
            "content_type": 1,
7
            "content_length": 1,
8
            "text": 1,
9
        }
10
    }

1
    database = Database()
2
    collection = database.get_collection()
3
4
    app.logger.info("Vector search query without post filter")
5
    res = list(collection.aggregate([new_search_query, project]))
6
7
    app.logger.info("Vector search query run")
8
    app.logger.debug(res)
9
    return res

1
    results = format_results(results)

1
def format_results(results):
2
    formatted_results = []
3
    for _idx, index in enumerate(results):
4
        parse_result = urlparse(index["url"])
5
        parsed_url = {
6
            "netloc": parse_result.netloc,
7
            "path": parse_result.path,
8
            "params": parse_result.params,
9
            "query": parse_result.query,
10
            "fragment": parse_result.fragment,
11
            "hostname": parse_result.hostname,
12
            "last_part": parse_result.path.rstrip("/").split("/")[-1],
13
        }
14
        index["parsed_url"] = parsed_url
15
        formatted_results.append(index)
16
17
    return formatted_results

1
@app.route("/search", methods=["POST"])
2
def search_post():
3
    ...
4
5
    results = format_results(results)
6
7
    return render_template("search.html", results=results, query=query)