Multimodal Knowledge Base

from openai import OpenAI

from mixpeek import Mixpeek

client = Mixpeek(api_key="YOUR_API_KEY", namespace="knowledge-base")
openai = OpenAI(api_key="YOUR_OPENAI_KEY")

# 1. PDFs become layout blocks, and videos become transcribed passages
docs_bucket = client.buckets.create(
    bucket_name="knowledge-docs",
    bucket_schema={"properties": {"document": {"type": "pdf"}}},
)
docs = client.collections.create(
    collection_name="knowledge-docs",
    source={"type": "bucket", "bucket_ids": [docs_bucket["bucket_id"]]},
    feature_extractor={
        "feature_extractor_name": "document_graph_extractor",
        "version": "v1",
    },
)
videos_bucket = client.buckets.create(
    bucket_name="knowledge-videos",
    bucket_schema={"properties": {"video": {"type": "video"}}},
)
videos = client.collections.create(
    collection_name="knowledge-videos",
    source={"type": "bucket", "bucket_ids": [videos_bucket["bucket_id"]]},
    feature_extractor={
        "feature_extractor_name": "multimodal_extractor",
        "version": "v1",
        "parameters": {
            "split_method": "silence",
            "run_transcription": True,
            "run_transcription_embedding": True,
        },
    },
)
client.buckets.upload(
    docs_bucket["bucket_id"],
    blobs=[{"property": "document", "type": "pdf", "data": "s3://your-bucket/handbook/remote-work-policy.pdf"}],
)
client.buckets.upload(
    videos_bucket["bucket_id"],
    blobs=[{"property": "video", "type": "video", "data": "s3://your-bucket/all-hands/2026-08.mp4"}],
)
client.collections.trigger(docs["collection_id"])
client.collections.trigger(videos["collection_id"])

# 2. One stage searches both collections: each search reads the index its collection writes
retriever = client.retrievers.create(
    retriever_name="knowledge-base",
    collection_identifiers=["knowledge-docs", "knowledge-videos"],
    input_schema={"query": {"type": "text", "required": True}},
    stages=[
        {
            "stage_name": "search",
            "stage_id": "feature_search",
            "parameters": {
                "searches": [
                    {
                        "feature_uri": "mixpeek://document_graph_extractor@v1/intfloat__multilingual_e5_large_instruct",
                        "query": {"input_mode": "text", "value": "{{INPUT.query}}"},
                        "top_k": 30,
                    },
                    {
                        "feature_uri": "mixpeek://multimodal_extractor@v1/multilingual_e5_large_instruct_v1",
                        "query": {"input_mode": "text", "value": "{{INPUT.query}}"},
                        "top_k": 30,
                    },
                ],
                "fusion": "rrf",
                "final_top_k": 12,
            },
        },
    ],
)

question = "What is our company policy on remote work?"
results = client.retrievers.execute(retriever["retriever_id"], inputs={"query": question})

# 3. PDF blocks keep their text in text_raw and video passages in transcription, so the
#    numbered context is assembled here with a citation for each kind
passages = []
for i, doc in enumerate(results["documents"], 1):
    if doc.get("text_raw"):
        passages.append(f"[{i}] {doc['text_raw']} (PDF {doc.get('root_object_id')}, page {doc.get('page_number')})")
    else:
        passages.append(f"[{i}] {doc.get('transcription')} (video {doc.get('root_object_id')} at {doc.get('start_time')}s)")

response = openai.chat.completions.create(
    model="gpt-4o",
    messages=[
        {"role": "system", "content": "Answer from this knowledge base and cite the passage numbers. " + " ".join(passages)},
        {"role": "user", "content": question},
    ],
)
print(response.choices[0].message.content)