Haystack

pip install mixedbread-ai-haystack

MXBAI_API_KEY=your_api_key_here

from mixedbread_ai_haystack import MixedbreadTextEmbedder

print("=== Basic Text Embedding Usage ===")

text_embedder = MixedbreadTextEmbedder(
    model="mixedbread-ai/mxbai-embed-large-v1",
    # api_key="your-api-key"  # or set MXBAI_API_KEY env var
)

query = "Which country has the biggest population?"
result = text_embedder.run(text=query)

print(f"Query: {query}")
print(f"Embedding dimension: {len(result['embedding'])}")
print(f"First 5 values: {result['embedding'][:5]}")

from haystack import Document
from mixedbread_ai_haystack import MixedbreadDocumentEmbedder

print("=== Basic Document Embedding Usage ===")

doc_embedder = MixedbreadDocumentEmbedder(
    model="mixedbread-ai/mxbai-embed-large-v1",
    # api_key="your-api-key"  # or set MXBAI_API_KEY env var
)

documents = [
    Document(content="China is the most populous country in the world."),
    Document(content="India is the second most populous country in the world."),
    Document(content="Machine learning is a subset of artificial intelligence."),
]

result = doc_embedder.run(documents=documents)

print(f"Embedded {len(result['documents'])} documents")
for i, doc in enumerate(result["documents"]):
    print(f"Doc {i+1}: {doc.content[:50]}...")
    print(f"  First 5 values: {doc.embedding[:5]}")
    print(f"  Embedding dimension: {len(doc.embedding)}")

from mixedbread_ai_haystack import MixedbreadDocumentParser

print("=== Basic Document Parsing ===")

parser = MixedbreadDocumentParser(
    # api_key="your-api-key"  # or set MXBAI_API_KEY env var
)

try:
    result = parser.run(sources=["path/to/your/document.pdf"])

    print(f"Parsed {len(result['documents'])} document chunks")

    for i, doc in enumerate(result["documents"]):
        print(f"Chunk {i + 1}:")
        print(f"  Content: {doc.content[:100]}...")

except Exception as e:
    print(f"Error parsing document: {e}")

from haystack import Document
from mixedbread_ai_haystack import MixedbreadReranker

print("=== Basic Reranking Usage ===")

reranker = MixedbreadReranker(
    model="mixedbread-ai/mxbai-rerank-large-v2",
    top_k=3,
    # api_key="your-api-key"  # or set MXBAI_API_KEY env var
)

documents = [
    Document(content="Machine learning is a subset of artificial intelligence that focuses on algorithms."),
    Document(content="Deep learning uses neural networks with multiple layers for complex pattern recognition."),
    Document(content="Natural language processing helps computers understand and generate human language."),
    Document(content="Computer vision enables machines to interpret and analyze visual information."),
    Document(content="Reinforcement learning trains agents through trial and error with rewards."),
]

query = "What is deep learning and neural networks?"
result = reranker.run(documents=documents, query=query)

print(f"Query: {query}")
print(f"Original documents: {len(documents)}")
print(f"Reranked documents: {len(result['documents'])}")

for i, doc in enumerate(result["documents"], 1):
    print(f"Rank {i}:")
    print(f"  Content: {doc.content}")
    print(f"  Score: {doc.meta['rerank_score']}")

from mixedbread_ai_haystack import MixedbreadVectorStoreRetriever

print("=== Chunk Search Example ===")

retriever = MixedbreadVectorStoreRetriever(
    vector_store_identifiers=["your-vector-store-id"],
    search_type="chunk",
    top_k=5,
    # api_key="your-api-key"  # or set MXBAI_API_KEY env var
)

query = "Your question here"
result = retriever.run(query=query)

print(f"Query: {query}")
print(f"Found {len(result['documents'])} relevant chunks:")

for i, doc in enumerate(result["documents"]):
    print(f"Chunk {i+1}:")
    print(f"  Content: {doc.content[:100]}...")
    print(f"  Source: {doc.meta.get('filename', 'unknown')}")
    print(f"  Score: {doc.meta.get('retrieval_score', 'N/A')}")

from dotenv import load_dotenv
from haystack import Pipeline
from haystack.components.generators import OpenAIGenerator
from haystack.components.builders import PromptBuilder
from mixedbread_ai_haystack import MixedbreadVectorStoreRetriever

load_dotenv()

question = "Your question here"

template = """
Answer the question based on the provided context.

Context:
{% for document in documents %}
{{ document.content }}
{% endfor %}

Question: {{ question }}
Answer:
"""

rag_pipeline = Pipeline()
rag_pipeline.add_component(
    "retriever",
    MixedbreadVectorStoreRetriever(
        vector_store_identifiers=["your-vector-store-id"],
        search_type="chunk",
        top_k=3,
    ),
)
rag_pipeline.add_component(
    "prompt_builder",
    PromptBuilder(template=template, required_variables=["documents", "question"]),
)
rag_pipeline.add_component("llm", OpenAIGenerator(model="gpt-4o"))

rag_pipeline.connect("retriever.documents", "prompt_builder.documents")
rag_pipeline.connect("prompt_builder", "llm")

result = rag_pipeline.run(
    {"retriever": {"query": question}, "prompt_builder": {"question": question}}
)
print(result["llm"]["replies"][0])