Skip to content

Commit dbf2486

Browse files
kacperlukawskikacperlukawski
authored
feat: add Prior Labs integration documentation (#430)
* feat: add Prior Labs integration documentation * docs: enhance tool descriptions
1 parent 7c3c655 commit dbf2486

2 files changed

Lines changed: 177 additions & 0 deletions

File tree

‎integrations/prior-labs.md‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
---
2+
layout: integration
3+
name: Prior Labs
4+
description: Tabular data science via MCP - predict missing values, classify, and run regression on tabular datasets using Prior Labs' foundation model.
5+
authors:
6+
- name: Prior Labs
7+
socials:
8+
github: priorlabs
9+
pypi: https://pypi.org/project/mcp-haystack/
10+
repo: https://github.com/priorlabs
11+
type: Tool Integration
12+
report_issue: https://github.com/priorlabs
13+
logo: /logos/prior-labs.svg
14+
version: Haystack 2.0
15+
toc: true
16+
mcp: true
17+
---
18+
19+
**Table of Contents**
20+
21+
- [Overview](#overview)
22+
- [Prerequisites](#prerequisites)
23+
- [Installation](#installation)
24+
- [Usage](#usage)
25+
- [Connecting to Prior Labs](#connecting-to-prior-labs)
26+
- [Predicting Tabular Data with an Agent](#predicting-tabular-data-with-an-agent)
27+
- [License](#license)
28+
29+
## Overview
30+
31+
[Prior Labs](https://priorlabs.ai) is the team behind [TabPFN](https://github.com/PriorLabs/TabPFN), a foundation model for tabular data. They expose TabPFN as a cloud service via the [Model Context Protocol (MCP)](https://modelcontextprotocol.io), enabling Haystack agents to run tabular machine learning without writing any ML code.
32+
33+
The MCP server exposes five tools (as of March 2026):
34+
35+
| Tool | Description |
36+
|------|-------------|
37+
| `upload_dataset` | Upload a CSV file (train or test) and receive a `dataset_id` for use in subsequent calls |
38+
| `fit_and_predict_from_dataset` | Train a TabPFN model on an uploaded training file and predict on an uploaded test file |
39+
| `predict_from_dataset` | Run prediction with a previously trained model on an uploaded test file |
40+
| `fit_and_predict_inline` | Train and predict on small arrays already present in the conversation |
41+
| `predict_inline` | Predict with a previously trained model on inline arrays |
42+
43+
Both classification and regression tasks are supported. For small datasets shared directly in the conversation the agent uses the inline tools, while for larger file-based datasets it calls `upload_dataset` first.
44+
45+
This integration does not require a separate Prior Labs Python package. Instead, it uses the existing `mcp-haystack` package to connect to the Prior Labs MCP server and discover its tools automatically.
46+
47+
## Prerequisites
48+
49+
1. Sign up for a Prior Labs account at [priorlabs.ai](https://priorlabs.ai) and obtain an API token.
50+
2. Set the required environment variable:
51+
52+
```bash
53+
export PRIOR_LABS_MCP_TOKEN="your-prior-labs-token"
54+
```
55+
56+
## Installation
57+
58+
```bash
59+
pip install mcp-haystack
60+
```
61+
62+
## Usage
63+
64+
### Connecting to Prior Labs
65+
66+
Use `StreamableHttpServerInfo` to point to the Prior Labs MCP server and `MCPToolset` to connect and fetch all available tools:
67+
68+
```python
69+
from haystack_integrations.tools.mcp import MCPToolset, StreamableHttpServerInfo
70+
from haystack.utils import Secret
71+
72+
server_info = StreamableHttpServerInfo(
73+
url="https://api.priorlabs.ai/mcp/server",
74+
token=Secret.from_env_var("PRIOR_LABS_MCP_TOKEN"),
75+
)
76+
toolset = MCPToolset(server_info=server_info, eager_connect=True)
77+
```
78+
79+
`eager_connect=True` causes the toolset to immediately connect and fetch the tool definitions from the server.
80+
81+
### Predicting Tabular Data with an Agent
82+
83+
The following example sets up a Haystack `Agent` that uses Prior Labs tools to predict missing values in a CSV-style dataset:
84+
85+
```python
86+
from haystack_integrations.tools.mcp import MCPToolset, StreamableHttpServerInfo
87+
from haystack_integrations.components.generators.anthropic import AnthropicChatGenerator
88+
from haystack.utils import Secret
89+
from haystack.components.agents import Agent
90+
from haystack.dataclasses import ChatMessage
91+
92+
server_info = StreamableHttpServerInfo(
93+
url="https://api.priorlabs.ai/mcp/server",
94+
token=Secret.from_env_var("PRIOR_LABS_MCP_TOKEN"),
95+
)
96+
toolset = MCPToolset(server_info=server_info, eager_connect=True)
97+
98+
llm = AnthropicChatGenerator(model="claude-opus-4-6", tools=toolset)
99+
100+
agent = Agent(
101+
chat_generator=llm,
102+
tools=toolset,
103+
system_prompt=(
104+
"You are a data scientist working on tabular data. "
105+
"Use provided tools to solve all the given tasks. "
106+
"Always use the tools when you can, and only answer directly if you are sure "
107+
"you have the correct answer. If you are not sure, use the tools to find the answer. "
108+
"Always think step by step."
109+
),
110+
max_agent_steps=10,
111+
)
112+
113+
input_table = """age,income,purchased
114+
22,32000,0
115+
25,45000,0
116+
28,51000,0
117+
31,58000,1
118+
34,62000,1
119+
37,75000,1
120+
40,88000,1
121+
43,95000,1
122+
26,48000,
123+
33,61000,
124+
29,39000,
125+
41,82000,
126+
"""
127+
128+
output = agent.run(messages=[
129+
ChatMessage.from_user(
130+
"I have a dataset with columns 'age', 'income', and 'purchased' "
131+
"(0 = did not purchase, 1 = purchased). Some rows are missing the 'purchased' value. "
132+
"Can you predict the missing values based on the known examples?"
133+
),
134+
ChatMessage.from_user(input_table),
135+
])
136+
print(output["last_message"].text)
137+
```
138+
139+
Because the dataset is small and provided inline, the agent calls `fit_and_predict_inline` with the known rows as training data and the incomplete rows as test data, using `task_type="classification"`. The result is the predicted `purchased` values (0 or 1) for each incomplete row.
140+
141+
## License
142+
143+
The `mcp-haystack` package is distributed under the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license. Use of the Prior Labs API is subject to [Prior Labs' terms of service](https://priorlabs.ai).

‎logos/prior-labs.svg‎

Lines changed: 34 additions & 0 deletions
Loading

0 commit comments

Comments
 (0)