diff --git a/.env.example b/.env.example
index 636aa4a..2e49080 100644
--- a/.env.example
+++ b/.env.example
@@ -1 +1,31 @@
-ANTHROPIC_API_KEY=your_api_key_here
\ No newline at end of file
+ANTHROPIC_API_KEY=your_api_key_here
+
+# Azure Databricks
+DATABRICKS_HOST=adb-xxxxx.azuredatabricks.net
+DATABRICKS_TOKEN=dapi...
+
+# Snowflake
+SNOWFLAKE_ACCOUNT=xy12345.eu-west-1
+
+# PostgreSQL
+POSTGRES_HOST=localhost
+
+# MySQL
+MYSQL_HOST=localhost
+
+# MongoDB
+MONGODB_URI=mongodb://localhost:27017
+MONGODB_DATABASE=your_database_name
+
+# Amazon Redshift
+REDSHIFT_HOST=cluster.abc123.eu-west-1.redshift.amazonaws.com
+
+# Elasticsearch
+ELASTICSEARCH_HOST=localhost
+ELASTICSEARCH_PORT=9200
+
+# DuckDB
+DUCKDB_FILE_PATH=/path/to/my.duckdb
+
+# Microsoft Fabric
+FABRIC_SERVER=xyz.datawarehouse.fabric.microsoft.com
\ No newline at end of file
diff --git a/README.md b/README.md
index 11caab8..66161de 100644
--- a/README.md
+++ b/README.md
@@ -217,7 +217,7 @@ Connect directly to your database — agents fetch a table and run the full pipe
-5 database connectors — Azure Databricks, Snowflake, PostgreSQL, MySQL, BigQuery
+10 database connectors — Azure Databricks, Snowflake, PostgreSQL, MySQL, BigQuery, MongoDB, Redshift, DuckDB, Microsoft Fabric, Elasticsearch
---
@@ -352,7 +352,12 @@ multi-agent-data-pipeline/
│ │ ├── snowflake_conn.py # Snowflake
│ │ ├── postgres.py # PostgreSQL
│ │ ├── mysql.py # MySQL
-│ │ └── bigquery.py # BigQuery
+│ │ ├── bigquery.py # BigQuery
+│ │ ├── mongodb.py # MongoDB
+│ │ ├── redshift.py # Amazon Redshift
+│ │ ├── duckdb_conn.py # DuckDB
+│ │ ├── fabric.py # Microsoft Fabric
+│ │ └── elasticsearch_conn.py # Elasticsearch
│ ├── models.py # Pydantic schemas
│ └── pipeline.py # Orchestrator
├── demo/
@@ -517,6 +522,78 @@ df = fetch_table(
)
```
+#### MongoDB
+
+```python
+from src.connectors.mongodb import fetch_collection
+
+df = fetch_collection(
+ uri="mongodb://localhost:27017",
+ database="my_database",
+ collection="my_collection",
+ limit=1000
+)
+```
+
+#### Amazon Redshift
+
+```python
+from src.connectors.redshift import fetch_table
+
+df = fetch_table(
+ host="cluster.abc123.eu-west-1.redshift.amazonaws.com",
+ port=5439,
+ database="dev",
+ user="awsuser",
+ password="my_password",
+ table="my_table"
+)
+```
+
+#### DuckDB
+
+```python
+from src.connectors.duckdb_conn import fetch_table
+
+df = fetch_table(
+ filepath="/path/to/my.duckdb",
+ table="my_table",
+ limit=1000
+)
+```
+
+#### Microsoft Fabric
+
+> Requires [ODBC Driver 18 for SQL Server](https://learn.microsoft.com/en-us/sql/connect/odbc/download-odbc-driver-for-sql-server) installed at the OS level.
+
+```python
+from src.connectors.fabric import fetch_table
+
+df = fetch_table(
+ server="xyz.datawarehouse.fabric.microsoft.com",
+ database="my_warehouse",
+ user="user@org.com",
+ password="my_password",
+ table="my_table"
+)
+```
+
+#### Elasticsearch
+
+```python
+from src.connectors.elasticsearch_conn import fetch_index
+
+df = fetch_index(
+ host="localhost",
+ port=9200,
+ index="my_index",
+ username="elastic", # optional
+ password="my_password", # optional
+ use_ssl=False,
+ limit=1000
+)
+```
+
---
### Connector Status
@@ -528,10 +605,11 @@ df = fetch_table(
| PostgreSQL | User/Pass | ✅ | ✅ | Stable |
| MySQL | User/Pass | ✅ | ✅ | Stable |
| BigQuery | Service Account JSON | ✅ | ✅ | Stable |
-| MongoDB | — | 🔜 | 🔜 | Planned |
-| Redshift | — | 🔜 | 🔜 | Planned |
-| DuckDB | — | 🔜 | 🔜 | Planned |
-| Microsoft Fabric | — | 🔜 | 🔜 | Planned |
+| MongoDB | URI | ✅ | ✅ | Stable |
+| Amazon Redshift | User/Pass | ✅ | ✅ | Stable |
+| DuckDB | File path | ✅ | ✅ | Stable |
+| Microsoft Fabric | User/Pass | ✅ | ✅ | Requires ODBC Driver 18 |
+| Elasticsearch | Optional User/Pass | ✅ | ✅ | Stable |
> Want to add a connector? See [Contributing](#contributing)
@@ -693,11 +771,14 @@ We want to support every major database. Next targets:
| Database | Difficulty | Issue |
|----------|-----------|-------|
-| MongoDB | Medium | #1 |
-| Redshift | Easy | #2 |
-| DuckDB | Easy | #3 |
-| Microsoft Fabric | Medium | #4 |
-| Elasticsearch | Hard | #5 |
+| MongoDB | Medium | #1 ✅ |
+| Redshift | Easy | #2 ✅ |
+| DuckDB | Easy | #3 ✅ |
+| Microsoft Fabric | Medium | #4 ✅ |
+| Elasticsearch | Hard | #5 ✅ |
+| Oracle DB | Medium | #6 |
+| CockroachDB | Easy | #7 |
+| ClickHouse | Medium | #8 |
#### ☁️ Cloud Implementations
Deploy this on your cloud and contribute the implementation:
@@ -802,13 +883,17 @@ Follow this pattern — every agent has the same structure:
```python
# src/agents/your_agent.py
+from pydantic import BaseModel, Field
+from typing import List
+
SYSTEM_PROMPT = """You are a [role] agent.
Respond ONLY with valid JSON. No markdown. No explanation.
JSON format: { ... }"""
-class YourAgentResult:
- def __init__(self, **kwargs): ...
- def model_dump(self): return self.__dict__
+class YourAgentResult(BaseModel):
+ some_field: str = "default"
+ some_list: List[str] = Field(default_factory=list)
+ some_count: int = 0
def run(data: str, context: int) -> YourAgentResult:
response = client.messages.create(...)
@@ -867,7 +952,7 @@ tests/test_pipeline.py::TestCSVLoading::test_demo_csv_has_rows PASSED
| Validation | Pydantic v2 |
| CLI | Typer + Rich |
| UI | Streamlit |
-| Connectors | Databricks SDK, Snowflake, psycopg2, mysql-connector, BigQuery |
+| Connectors | Databricks SDK, Snowflake, psycopg2, mysql-connector, BigQuery, pymongo, redshift-connector, duckdb, elasticsearch, pyodbc |
| Testing | pytest |
| Packaging | pyproject.toml |
@@ -877,15 +962,16 @@ tests/test_pipeline.py::TestCSVLoading::test_demo_csv_has_rows PASSED
- [x] CSV pipeline — 5 agents
- [x] PDF intelligence — 5 agents
-- [x] Database connectors — 5 databases
+- [x] Database connectors — 10 databases
- [x] Streamlit UI — dark theme
- [x] CLI entrypoint
- [x] JSON export
+- [x] MongoDB connector
+- [x] Redshift connector
+- [x] DuckDB connector
+- [x] Microsoft Fabric connector
+- [x] Elasticsearch connector
- [ ] pip package — `pip install multi-agent-data-pipeline`
-- [ ] MongoDB connector
-- [ ] Redshift connector
-- [ ] DuckDB connector
-- [ ] Microsoft Fabric connector
- [ ] Async parallel agent execution
- [ ] Agent memory — learn from past runs
- [ ] Webhook support — trigger via HTTP
diff --git a/app.py b/app.py
index c6e46aa..ac41c95 100644
--- a/app.py
+++ b/app.py
@@ -630,7 +630,7 @@ def run_pipeline_ui(df):
db_type = st.selectbox(
"Database",
- ["Azure Databricks", "Snowflake", "PostgreSQL", "MySQL", "BigQuery"],
+ ["Azure Databricks", "Snowflake", "PostgreSQL", "MySQL", "BigQuery", "MongoDB", "Microsoft Fabric", "Amazon Redshift", "DuckDB", "Elasticsearch"],
label_visibility="collapsed"
)
@@ -733,6 +733,121 @@ def run_pipeline_ui(df):
else:
st.warning("Please fill all fields")
+ elif db_type == "MongoDB":
+ col1, col2 = st.columns(2)
+ with col1:
+ uri = st.text_input("Connection URI", placeholder="mongodb://localhost:27017")
+ collection = st.text_input("Collection", placeholder="my_collection")
+ with col2:
+ database = st.text_input("Database", placeholder="my_database")
+ limit = st.number_input("Row limit", min_value=1, max_value=100000, value=1000)
+
+ if st.button("🔌 Connect & Fetch Collection"):
+ if uri and database and collection:
+ try:
+ from src.connectors.mongodb import fetch_collection
+ with st.spinner("Connecting to MongoDB..."):
+ df = fetch_collection(uri, database, collection, int(limit))
+ st.success(f"Connected — {len(df)} documents fetched from {database}.{collection}")
+ st.dataframe(df, use_container_width=True, height=240)
+ except Exception as e:
+ st.error(f"Connection failed: {e}")
+ else:
+ st.warning("Please fill all fields")
+
+ elif db_type == "Microsoft Fabric":
+ col1, col2 = st.columns(2)
+ with col1:
+ server = st.text_input("Server", placeholder="xyz.datawarehouse.fabric.microsoft.com")
+ database = st.text_input("Database", placeholder="my_warehouse")
+ table = st.text_input("Table", placeholder="my_table")
+ with col2:
+ user = st.text_input("Username", placeholder="user@org.com")
+ password = st.text_input("Password", type="password")
+
+ if st.button("🔌 Connect & Fetch Table"):
+ if server and database and user and password and table:
+ try:
+ from src.connectors.fabric import fetch_table
+ with st.spinner("Connecting to Microsoft Fabric..."):
+ df = fetch_table(server, database, user, password, table)
+ st.success(f"Connected — {len(df)} rows fetched from {table}")
+ st.dataframe(df, use_container_width=True, height=240)
+ except Exception as e:
+ st.error(f"Connection failed: {e}")
+ else:
+ st.warning("Please fill all fields")
+
+ elif db_type == "Amazon Redshift":
+ col1, col2 = st.columns(2)
+ with col1:
+ host = st.text_input("Host", placeholder="cluster.abc123.eu-west-1.redshift.amazonaws.com")
+ database = st.text_input("Database", placeholder="dev")
+ table = st.text_input("Table", placeholder="my_table")
+ with col2:
+ port = st.text_input("Port", value="5439")
+ user = st.text_input("Username", placeholder="awsuser")
+ password = st.text_input("Password", type="password")
+
+ if st.button("🔌 Connect & Fetch Table"):
+ if host and database and user and password and table:
+ try:
+ from src.connectors.redshift import fetch_table
+ with st.spinner("Connecting to Redshift..."):
+ df = fetch_table(host, int(port), database, user, password, table)
+ st.success(f"Connected — {len(df)} rows fetched from {table}")
+ st.dataframe(df, use_container_width=True, height=240)
+ except Exception as e:
+ st.error(f"Connection failed: {e}")
+ else:
+ st.warning("Please fill all fields")
+
+ elif db_type == "DuckDB":
+ col1, col2 = st.columns(2)
+ with col1:
+ filepath = st.text_input("Database file path", placeholder="/path/to/my.duckdb")
+ table = st.text_input("Table", placeholder="my_table")
+ with col2:
+ limit = st.number_input("Row limit", min_value=1, max_value=100000, value=1000)
+
+ if st.button("🔌 Connect & Fetch Table"):
+ if filepath and table:
+ try:
+ from src.connectors.duckdb_conn import fetch_table
+ with st.spinner("Connecting to DuckDB..."):
+ df = fetch_table(filepath, table, int(limit))
+ st.success(f"Connected — {len(df)} rows fetched from {table}")
+ st.dataframe(df, use_container_width=True, height=240)
+ except Exception as e:
+ st.error(f"Connection failed: {e}")
+ else:
+ st.warning("Please fill all fields")
+
+ elif db_type == "Elasticsearch":
+ col1, col2 = st.columns(2)
+ with col1:
+ es_host = st.text_input("Host", placeholder="localhost")
+ es_index = st.text_input("Index", placeholder="my_index")
+ es_limit = st.number_input("Row limit", min_value=1, max_value=10000, value=1000)
+ with col2:
+ es_port = st.text_input("Port", value="9200")
+ es_user = st.text_input("Username (optional)", placeholder="elastic")
+ es_password = st.text_input("Password (optional)", type="password")
+ es_ssl = st.checkbox("Use SSL (HTTPS)")
+
+ if st.button("🔌 Connect & Fetch Index"):
+ if es_host and es_index:
+ try:
+ from src.connectors.elasticsearch_conn import fetch_index
+ with st.spinner("Connecting to Elasticsearch..."):
+ df = fetch_index(es_host, int(es_port), es_index, es_user, es_password, es_ssl, int(es_limit))
+ st.success(f"Connected — {len(df)} documents fetched from {es_index}")
+ st.dataframe(df, use_container_width=True, height=240)
+ except Exception as e:
+ st.error(f"Connection failed: {e}")
+ else:
+ st.warning("Please fill Host and Index")
+
elif db_type == "BigQuery":
col1, col2 = st.columns(2)
with col1:
diff --git a/docs/images/arch.png:Zone.Identifier b/docs/images/arch.png:Zone.Identifier
deleted file mode 100644
index 2d1b061..0000000
Binary files a/docs/images/arch.png:Zone.Identifier and /dev/null differ
diff --git a/logo.png b/logo.png
new file mode 100644
index 0000000..e69de29
diff --git a/logo.png:Zone.Identifier b/logo.png:Zone.Identifier
deleted file mode 100644
index 6a2d663..0000000
Binary files a/logo.png:Zone.Identifier and /dev/null differ
diff --git a/main.py b/main.py
index 25b292e..4feea65 100644
--- a/main.py
+++ b/main.py
@@ -3,7 +3,7 @@
import os
from rich.console import Console
from rich.panel import Panel
-from src.pipeline import run_pipeline
+from src.pipeline import run_pipeline, run_pipeline_mongo
app = typer.Typer()
console = Console()
@@ -36,5 +36,28 @@ def run(
json.dump(result.model_dump(), f, indent=2)
console.print(f"[green]→ Results saved to {output}[/green]")
+@app.command("run-mongo")
+def run_mongo(
+ uri: str = typer.Option(..., "--uri", help="MongoDB connection URI (e.g. mongodb://localhost:27017)"),
+ database: str = typer.Option(..., "--database", "-d", help="Database name"),
+ collection: str = typer.Option(..., "--collection", "-c", help="Collection name"),
+ limit: int = typer.Option(1000, "--limit", "-l", help="Max documents to fetch"),
+ output: str = typer.Option(None, "--output", "-o", help="Save results to JSON file"),
+):
+ """
+ Run the pipeline on a MongoDB collection.
+
+ Example:
+ python main.py run-mongo --uri mongodb://localhost:27017 --database mydb --collection orders
+ """
+ result = run_pipeline_mongo(uri, database, collection, limit)
+
+ if output:
+ import json
+ with open(output, "w") as f:
+ json.dump(result.model_dump(), f, indent=2)
+ console.print(f"[green]→ Results saved to {output}[/green]")
+
+
if __name__ == "__main__":
app()
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index f91dc07..e7762b3 100644
Binary files a/requirements.txt and b/requirements.txt differ
diff --git a/src/agents/action_extractor.py b/src/agents/action_extractor.py
index 61bb024..ce833b6 100644
--- a/src/agents/action_extractor.py
+++ b/src/agents/action_extractor.py
@@ -1,64 +1,59 @@
-import os
-import json
-from anthropic import Anthropic
-from dotenv import load_dotenv
-
-load_dotenv()
-
-client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
-
-SYSTEM_PROMPT = """You are an action item extraction agent.
-Your job is to extract all action items, decisions, deadlines, and follow-ups from document text.
-You must respond ONLY with valid JSON. No explanation, no markdown, no code fences.
-JSON format:
-{
- "action_items": ["action1", "action2"],
- "decisions_made": ["decision1", "decision2"],
- "deadlines": ["deadline1 - date", "deadline2 - date"],
- "follow_ups": ["follow up1", "follow up2"],
- "owners": ["person/team responsible1", "person/team responsible2"],
- "priority_actions": ["most urgent action1", "most urgent action2"],
- "total_actions": 8
-}"""
-
-class ActionExtractorResult:
- def __init__(self, **kwargs):
- self.action_items = kwargs.get("action_items", [])
- self.decisions_made = kwargs.get("decisions_made", [])
- self.deadlines = kwargs.get("deadlines", [])
- self.follow_ups = kwargs.get("follow_ups", [])
- self.owners = kwargs.get("owners", [])
- self.priority_actions = kwargs.get("priority_actions", [])
- self.total_actions = kwargs.get("total_actions", 0)
-
- def model_dump(self):
- return self.__dict__
-
-def run(text_preview: str, total_pages: int) -> ActionExtractorResult:
- print("[Action Extractor Agent] Starting...")
-
- response = client.messages.create(
- model="claude-sonnet-4-5",
- max_tokens=1000,
- system=SYSTEM_PROMPT,
- messages=[
- {
- "role": "user",
- "content": f"Extract all action items from this document ({total_pages} pages):\n\n{text_preview}"
- }
- ]
- )
-
- raw = response.content[0].text.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
-
- try:
- data = json.loads(raw)
- result = ActionExtractorResult(**data)
- print(f"[Action Extractor Agent] Done — {result.total_actions} actions found")
- return result
- except Exception as e:
- print(f"[Action Extractor Agent] Error: {e}")
- return ActionExtractorResult(
- action_items=["Could not parse response"],
- total_actions=0
- )
\ No newline at end of file
+import os
+import json
+from anthropic import Anthropic
+from dotenv import load_dotenv
+from pydantic import BaseModel, Field
+from typing import List
+
+load_dotenv()
+
+client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
+
+SYSTEM_PROMPT = """You are an action item extraction agent.
+Your job is to extract all action items, decisions, deadlines, and follow-ups from document text.
+You must respond ONLY with valid JSON. No explanation, no markdown, no code fences.
+JSON format:
+{
+ "action_items": ["action1", "action2"],
+ "decisions_made": ["decision1", "decision2"],
+ "deadlines": ["deadline1 - date", "deadline2 - date"],
+ "follow_ups": ["follow up1", "follow up2"],
+ "owners": ["person/team responsible1", "person/team responsible2"],
+ "priority_actions": ["most urgent action1", "most urgent action2"],
+ "total_actions": 8
+}"""
+
+class ActionExtractorResult(BaseModel):
+ action_items: List[str] = Field(default_factory=list)
+ decisions_made: List[str] = Field(default_factory=list)
+ deadlines: List[str] = Field(default_factory=list)
+ follow_ups: List[str] = Field(default_factory=list)
+ owners: List[str] = Field(default_factory=list)
+ priority_actions: List[str] = Field(default_factory=list)
+ total_actions: int = 0
+
+def run(text_preview: str, total_pages: int) -> ActionExtractorResult:
+ print("[Action Extractor Agent] Starting...")
+
+ response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1000,
+ system=SYSTEM_PROMPT,
+ messages=[
+ {
+ "role": "user",
+ "content": f"Extract all action items from this document ({total_pages} pages):\n\n{text_preview}"
+ }
+ ]
+ )
+
+ raw = response.content[0].text.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
+
+ try:
+ data = json.loads(raw)
+ result = ActionExtractorResult(**data)
+ print(f"[Action Extractor Agent] Done — {result.total_actions} actions found")
+ return result
+ except Exception as e:
+ print(f"[Action Extractor Agent] Error: {e}")
+ return ActionExtractorResult(action_items=["Could not parse response"])
diff --git a/src/agents/entity_extractor.py b/src/agents/entity_extractor.py
index d918bc4..a755e4d 100644
--- a/src/agents/entity_extractor.py
+++ b/src/agents/entity_extractor.py
@@ -2,6 +2,8 @@
import json
from anthropic import Anthropic
from dotenv import load_dotenv
+from pydantic import BaseModel, Field
+from typing import List
load_dotenv()
@@ -21,18 +23,14 @@
"total_entities": 15
}"""
-class EntityExtractorResult:
- def __init__(self, **kwargs):
- self.people = kwargs.get("people", [])
- self.organisations = kwargs.get("organisations", [])
- self.locations = kwargs.get("locations", [])
- self.dates = kwargs.get("dates", [])
- self.amounts = kwargs.get("amounts", [])
- self.emails = kwargs.get("emails", [])
- self.total_entities = kwargs.get("total_entities", 0)
-
- def model_dump(self):
- return self.__dict__
+class EntityExtractorResult(BaseModel):
+ people: List[str] = Field(default_factory=list)
+ organisations: List[str] = Field(default_factory=list)
+ locations: List[str] = Field(default_factory=list)
+ dates: List[str] = Field(default_factory=list)
+ amounts: List[str] = Field(default_factory=list)
+ emails: List[str] = Field(default_factory=list)
+ total_entities: int = 0
def run(text_preview: str, total_pages: int) -> EntityExtractorResult:
print("[Entity Extractor Agent] Starting...")
@@ -58,12 +56,4 @@ def run(text_preview: str, total_pages: int) -> EntityExtractorResult:
return result
except Exception as e:
print(f"[Entity Extractor Agent] Error: {e}")
- return EntityExtractorResult(
- people=[],
- organisations=[],
- locations=[],
- dates=[],
- amounts=[],
- emails=[],
- total_entities=0
- )
+ return EntityExtractorResult()
diff --git a/src/agents/pdf_parser.py b/src/agents/pdf_parser.py
index b3acd3c..80c5a7d 100644
--- a/src/agents/pdf_parser.py
+++ b/src/agents/pdf_parser.py
@@ -1,66 +1,61 @@
-import os
-import json
-from anthropic import Anthropic
-from dotenv import load_dotenv
-
-load_dotenv()
-
-client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
-
-SYSTEM_PROMPT = """You are a document parsing agent.
-Your job is to analyse extracted PDF text and identify document structure, metadata and content type.
-You must respond ONLY with valid JSON. No explanation, no markdown, no code fences.
-JSON format:
-{
- "document_type": "invoice/contract/report/letter/other",
- "language": "English",
- "total_sections": 5,
- "has_tables": true,
- "has_numbers": true,
- "key_topics": ["topic1", "topic2", "topic3"],
- "document_quality": "good/fair/poor",
- "parsing_notes": ["note1", "note2"]
-}"""
-
-class PDFParserResult:
- def __init__(self, **kwargs):
- self.document_type = kwargs.get("document_type", "unknown")
- self.language = kwargs.get("language", "English")
- self.total_sections = kwargs.get("total_sections", 0)
- self.has_tables = kwargs.get("has_tables", False)
- self.has_numbers = kwargs.get("has_numbers", False)
- self.key_topics = kwargs.get("key_topics", [])
- self.document_quality = kwargs.get("document_quality", "unknown")
- self.parsing_notes = kwargs.get("parsing_notes", [])
-
- def model_dump(self):
- return self.__dict__
-
-def run(text_preview: str, total_pages: int) -> PDFParserResult:
- print("[PDF Parser Agent] Starting...")
-
- response = client.messages.create(
- model="claude-sonnet-4-5",
- max_tokens=1000,
- system=SYSTEM_PROMPT,
- messages=[
- {
- "role": "user",
- "content": f"Parse this PDF document ({total_pages} pages):\n\n{text_preview}"
- }
- ]
- )
-
- raw = response.content[0].text.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
-
- try:
- data = json.loads(raw)
- result = PDFParserResult(**data)
- print(f"[PDF Parser Agent] Done — {result.document_type} document detected")
- return result
- except Exception as e:
- print(f"[PDF Parser Agent] Error: {e}")
- return PDFParserResult(
- document_type="unknown",
- parsing_notes=["Could not parse response"]
- )
\ No newline at end of file
+import os
+import json
+from anthropic import Anthropic
+from dotenv import load_dotenv
+from pydantic import BaseModel, Field
+from typing import List
+
+load_dotenv()
+
+client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
+
+SYSTEM_PROMPT = """You are a document parsing agent.
+Your job is to analyse extracted PDF text and identify document structure, metadata and content type.
+You must respond ONLY with valid JSON. No explanation, no markdown, no code fences.
+JSON format:
+{
+ "document_type": "invoice/contract/report/letter/other",
+ "language": "English",
+ "total_sections": 5,
+ "has_tables": true,
+ "has_numbers": true,
+ "key_topics": ["topic1", "topic2", "topic3"],
+ "document_quality": "good/fair/poor",
+ "parsing_notes": ["note1", "note2"]
+}"""
+
+class PDFParserResult(BaseModel):
+ document_type: str = "unknown"
+ language: str = "English"
+ total_sections: int = 0
+ has_tables: bool = False
+ has_numbers: bool = False
+ key_topics: List[str] = Field(default_factory=list)
+ document_quality: str = "unknown"
+ parsing_notes: List[str] = Field(default_factory=list)
+
+def run(text_preview: str, total_pages: int) -> PDFParserResult:
+ print("[PDF Parser Agent] Starting...")
+
+ response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1000,
+ system=SYSTEM_PROMPT,
+ messages=[
+ {
+ "role": "user",
+ "content": f"Parse this PDF document ({total_pages} pages):\n\n{text_preview}"
+ }
+ ]
+ )
+
+ raw = response.content[0].text.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
+
+ try:
+ data = json.loads(raw)
+ result = PDFParserResult(**data)
+ print(f"[PDF Parser Agent] Done — {result.document_type} document detected")
+ return result
+ except Exception as e:
+ print(f"[PDF Parser Agent] Error: {e}")
+ return PDFParserResult(parsing_notes=["Could not parse response"])
diff --git a/src/agents/risk_detector.py b/src/agents/risk_detector.py
index 6e37f3f..bd30978 100644
--- a/src/agents/risk_detector.py
+++ b/src/agents/risk_detector.py
@@ -1,68 +1,61 @@
-import os
-import json
-from anthropic import Anthropic
-from dotenv import load_dotenv
-
-load_dotenv()
-
-client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
-
-SYSTEM_PROMPT = """You are a document risk detection agent.
-Your job is to identify risks, sensitive data, PII, compliance issues and red flags in documents.
-You must respond ONLY with valid JSON. No explanation, no markdown, no code fences.
-JSON format:
-{
- "pii_detected": true,
- "pii_types": ["names", "emails", "phone numbers"],
- "compliance_risks": ["GDPR risk - personal data present", "risk2"],
- "legal_risks": ["unsigned contract clause", "risk2"],
- "financial_risks": ["large payment terms", "risk2"],
- "overall_risk_score": 7.5,
- "risk_level": "high",
- "recommendations": ["recommendation1", "recommendation2"]
-}"""
-
-class RiskDetectorResult:
- def __init__(self, **kwargs):
- self.pii_detected = kwargs.get("pii_detected", False)
- self.pii_types = kwargs.get("pii_types", [])
- self.compliance_risks = kwargs.get("compliance_risks", [])
- self.legal_risks = kwargs.get("legal_risks", [])
- self.financial_risks = kwargs.get("financial_risks", [])
- self.overall_risk_score = kwargs.get("overall_risk_score", 0.0)
- self.risk_level = kwargs.get("risk_level", "low")
- self.recommendations = kwargs.get("recommendations", [])
-
- def model_dump(self):
- return self.__dict__
-
-def run(text_preview: str, total_pages: int) -> RiskDetectorResult:
- print("[Risk Detector Agent] Starting...")
-
- response = client.messages.create(
- model="claude-sonnet-4-5",
- max_tokens=1000,
- system=SYSTEM_PROMPT,
- messages=[
- {
- "role": "user",
- "content": f"Detect risks in this document ({total_pages} pages):\n\n{text_preview}"
- }
- ]
- )
-
- raw = response.content[0].text.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
-
- try:
- data = json.loads(raw)
- result = RiskDetectorResult(**data)
- print(f"[Risk Detector Agent] Done — risk level: {result.risk_level}, score: {result.overall_risk_score}/10")
- return result
- except Exception as e:
- print(f"[Risk Detector Agent] Error: {e}")
- return RiskDetectorResult(
- pii_detected=False,
- overall_risk_score=0.0,
- risk_level="unknown",
- recommendations=["Could not parse response"]
- )
\ No newline at end of file
+import os
+import json
+from anthropic import Anthropic
+from dotenv import load_dotenv
+from pydantic import BaseModel, Field
+from typing import List
+
+load_dotenv()
+
+client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
+
+SYSTEM_PROMPT = """You are a document risk detection agent.
+Your job is to identify risks, sensitive data, PII, compliance issues and red flags in documents.
+You must respond ONLY with valid JSON. No explanation, no markdown, no code fences.
+JSON format:
+{
+ "pii_detected": true,
+ "pii_types": ["names", "emails", "phone numbers"],
+ "compliance_risks": ["GDPR risk - personal data present", "risk2"],
+ "legal_risks": ["unsigned contract clause", "risk2"],
+ "financial_risks": ["large payment terms", "risk2"],
+ "overall_risk_score": 7.5,
+ "risk_level": "high",
+ "recommendations": ["recommendation1", "recommendation2"]
+}"""
+
+class RiskDetectorResult(BaseModel):
+ pii_detected: bool = False
+ pii_types: List[str] = Field(default_factory=list)
+ compliance_risks: List[str] = Field(default_factory=list)
+ legal_risks: List[str] = Field(default_factory=list)
+ financial_risks: List[str] = Field(default_factory=list)
+ overall_risk_score: float = 0.0
+ risk_level: str = "low"
+ recommendations: List[str] = Field(default_factory=list)
+
+def run(text_preview: str, total_pages: int) -> RiskDetectorResult:
+ print("[Risk Detector Agent] Starting...")
+
+ response = client.messages.create(
+ model="claude-sonnet-4-5",
+ max_tokens=1000,
+ system=SYSTEM_PROMPT,
+ messages=[
+ {
+ "role": "user",
+ "content": f"Detect risks in this document ({total_pages} pages):\n\n{text_preview}"
+ }
+ ]
+ )
+
+ raw = response.content[0].text.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
+
+ try:
+ data = json.loads(raw)
+ result = RiskDetectorResult(**data)
+ print(f"[Risk Detector Agent] Done — risk level: {result.risk_level}, score: {result.overall_risk_score}/10")
+ return result
+ except Exception as e:
+ print(f"[Risk Detector Agent] Error: {e}")
+ return RiskDetectorResult(recommendations=["Could not parse response"])
diff --git a/src/connectors/__init__.py b/src/connectors/__init__.py
index e69de29..d2743c0 100644
--- a/src/connectors/__init__.py
+++ b/src/connectors/__init__.py
@@ -0,0 +1,10 @@
+from . import mongodb
+from . import postgres
+from . import mysql
+from . import databricks
+from . import snowflake_conn
+from . import bigquery
+from . import fabric
+from . import redshift
+from . import duckdb_conn
+from . import elasticsearch_conn
diff --git a/src/connectors/bigquery.py b/src/connectors/bigquery.py
index c3cc394..32cf11a 100644
--- a/src/connectors/bigquery.py
+++ b/src/connectors/bigquery.py
@@ -1,8 +1,14 @@
+import re
import pandas as pd
from google.cloud import bigquery
from google.oauth2 import service_account
import json
+def _validate_identifier(name: str) -> str:
+ if not re.match(r'^[A-Za-z0-9_\-\.]+$', name):
+ raise ValueError(f"Invalid identifier: {name!r}")
+ return name
+
def connect(project_id: str, credentials_json: dict):
credentials = service_account.Credentials.from_service_account_info(
credentials_json,
@@ -26,6 +32,6 @@ def list_datasets(project_id: str, credentials_json: dict) -> list:
def fetch_table(project_id: str, credentials_json: dict, dataset: str, table: str, limit: int = 1000) -> pd.DataFrame:
client = connect(project_id, credentials_json)
- query = f"SELECT * FROM `{project_id}.{dataset}.{table}` LIMIT {limit}"
+ query = f"SELECT * FROM `{_validate_identifier(project_id)}.{_validate_identifier(dataset)}.{_validate_identifier(table)}` LIMIT {limit}"
df = client.query(query).to_dataframe()
return df
\ No newline at end of file
diff --git a/src/connectors/databricks.py b/src/connectors/databricks.py
index 25c100a..357b633 100644
--- a/src/connectors/databricks.py
+++ b/src/connectors/databricks.py
@@ -1,7 +1,13 @@
+import re
import pandas as pd
from databricks import sql
from databricks.sdk import WorkspaceClient
+def _validate_identifier(name: str) -> str:
+ if not re.match(r'^[A-Za-z0-9_\.]+$', name):
+ raise ValueError(f"Invalid identifier: {name!r}")
+ return name
+
def connect(host: str, token: str, http_path: str):
conn = sql.connect(
server_hostname=host,
@@ -22,7 +28,7 @@ def list_tables(host: str, token: str, http_path: str) -> list:
def fetch_table(host: str, token: str, http_path: str, table: str, limit: int = 1000) -> pd.DataFrame:
conn = connect(host, token, http_path)
cursor = conn.cursor()
- cursor.execute(f"SELECT * FROM {table} LIMIT {limit}")
+ cursor.execute(f"SELECT * FROM {_validate_identifier(table)} LIMIT {limit}")
columns = [desc[0] for desc in cursor.description]
rows = cursor.fetchall()
cursor.close()
diff --git a/src/connectors/duckdb_conn.py b/src/connectors/duckdb_conn.py
new file mode 100644
index 0000000..388a69a
--- /dev/null
+++ b/src/connectors/duckdb_conn.py
@@ -0,0 +1,23 @@
+import re
+import pandas as pd
+import duckdb
+
+def _validate_identifier(name: str) -> str:
+ if not re.match(r'^[A-Za-z0-9_\.]+$', name):
+ raise ValueError(f"Invalid identifier: {name!r}")
+ return name
+
+def connect(filepath: str):
+ return duckdb.connect(filepath)
+
+def list_tables(filepath: str) -> list:
+ conn = connect(filepath)
+ result = conn.execute("SHOW TABLES").fetchall()
+ conn.close()
+ return [row[0] for row in result]
+
+def fetch_table(filepath: str, table: str, limit: int = 1000) -> pd.DataFrame:
+ conn = connect(filepath)
+ df = conn.execute(f"SELECT * FROM {_validate_identifier(table)} LIMIT {limit}").df()
+ conn.close()
+ return df
diff --git a/src/connectors/elasticsearch_conn.py b/src/connectors/elasticsearch_conn.py
new file mode 100644
index 0000000..57f16c4
--- /dev/null
+++ b/src/connectors/elasticsearch_conn.py
@@ -0,0 +1,24 @@
+import pandas as pd
+from elasticsearch import Elasticsearch
+
+def connect(host: str, port: int, username: str = "", password: str = "", use_ssl: bool = False) -> Elasticsearch:
+ scheme = "https" if use_ssl else "http"
+ if username and password:
+ return Elasticsearch(f"{scheme}://{username}:{password}@{host}:{port}", verify_certs=False)
+ return Elasticsearch(f"{scheme}://{host}:{port}")
+
+def list_indices(host: str, port: int, username: str = "", password: str = "", use_ssl: bool = False) -> list:
+ es = connect(host, port, username, password, use_ssl)
+ indices = sorted(i for i in es.indices.get_alias().keys() if not i.startswith("."))
+ es.close()
+ return indices
+
+def fetch_index(host: str, port: int, index: str, username: str = "", password: str = "", use_ssl: bool = False, limit: int = 1000) -> pd.DataFrame:
+ es = connect(host, port, username, password, use_ssl)
+ response = es.search(index=index, query={"match_all": {}}, size=min(limit, 10000))
+ hits = response["hits"]["hits"]
+ es.close()
+ if not hits:
+ return pd.DataFrame()
+ rows = [{"_id": h["_id"], **h["_source"]} for h in hits]
+ return pd.DataFrame(rows)
diff --git a/src/connectors/fabric.py b/src/connectors/fabric.py
new file mode 100644
index 0000000..49a6014
--- /dev/null
+++ b/src/connectors/fabric.py
@@ -0,0 +1,38 @@
+import re
+import pandas as pd
+import pyodbc
+
+def _validate_identifier(name: str) -> str:
+ if not re.match(r'^[A-Za-z0-9_\.]+$', name):
+ raise ValueError(f"Invalid identifier: {name!r}")
+ return name
+
+def connect(server: str, database: str, user: str, password: str):
+ conn_str = (
+ "DRIVER={ODBC Driver 18 for SQL Server};"
+ f"SERVER={server};"
+ f"DATABASE={database};"
+ f"UID={user};"
+ f"PWD={password};"
+ "Encrypt=yes;TrustServerCertificate=no;"
+ )
+ return pyodbc.connect(conn_str)
+
+def list_tables(server: str, database: str, user: str, password: str) -> list:
+ conn = connect(server, database, user, password)
+ cursor = conn.cursor()
+ cursor.execute("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE='BASE TABLE' ORDER BY TABLE_NAME")
+ tables = [row[0] for row in cursor.fetchall()]
+ cursor.close()
+ conn.close()
+ return tables
+
+def fetch_table(server: str, database: str, user: str, password: str, table: str, limit: int = 1000) -> pd.DataFrame:
+ conn = connect(server, database, user, password)
+ cursor = conn.cursor()
+ cursor.execute(f"SELECT TOP {limit} * FROM {_validate_identifier(table)}")
+ columns = [desc[0] for desc in cursor.description]
+ rows = cursor.fetchall()
+ cursor.close()
+ conn.close()
+ return pd.DataFrame(rows, columns=columns)
diff --git a/src/connectors/mongodb.py b/src/connectors/mongodb.py
new file mode 100644
index 0000000..2e93969
--- /dev/null
+++ b/src/connectors/mongodb.py
@@ -0,0 +1,22 @@
+import pandas as pd
+from pymongo import MongoClient
+
+def connect(uri: str, database: str) -> tuple:
+ client = MongoClient(uri)
+ db = client[database]
+ return client, db
+
+def list_collections(uri: str, database: str) -> list:
+ client, db = connect(uri, database)
+ try:
+ return db.list_collection_names()
+ finally:
+ client.close()
+
+def fetch_collection(uri: str, database: str, collection: str, limit: int = 1000) -> pd.DataFrame:
+ client, db = connect(uri, database)
+ try:
+ cursor = db[collection].find({}, {"_id": 0}, limit=limit)
+ return pd.DataFrame(list(cursor))
+ finally:
+ client.close()
diff --git a/src/connectors/mysql.py b/src/connectors/mysql.py
index 35fb89d..81b020c 100644
--- a/src/connectors/mysql.py
+++ b/src/connectors/mysql.py
@@ -1,6 +1,12 @@
+import re
import pandas as pd
import mysql.connector
+def _validate_identifier(name: str) -> str:
+ if not re.match(r'^[A-Za-z0-9_\.]+$', name):
+ raise ValueError(f"Invalid identifier: {name!r}")
+ return name
+
def connect(host: str, port: int, database: str, user: str, password: str):
conn = mysql.connector.connect(
host=host,
@@ -23,7 +29,7 @@ def list_tables(host: str, port: int, database: str, user: str, password: str) -
def fetch_table(host: str, port: int, database: str, user: str, password: str, table: str, limit: int = 1000) -> pd.DataFrame:
conn = connect(host, port, database, user, password)
cursor = conn.cursor()
- cursor.execute(f"SELECT * FROM {table} LIMIT {limit}")
+ cursor.execute(f"SELECT * FROM {_validate_identifier(table)} LIMIT {limit}")
columns = [desc[0] for desc in cursor.description]
rows = cursor.fetchall()
cursor.close()
diff --git a/src/connectors/postgres.py b/src/connectors/postgres.py
index 47a8bba..8dc68cf 100644
--- a/src/connectors/postgres.py
+++ b/src/connectors/postgres.py
@@ -1,6 +1,12 @@
+import re
import pandas as pd
import psycopg2
+def _validate_identifier(name: str) -> str:
+ if not re.match(r'^[A-Za-z0-9_\.]+$', name):
+ raise ValueError(f"Invalid identifier: {name!r}")
+ return name
+
def connect(host: str, port: int, database: str, user: str, password: str):
conn = psycopg2.connect(
host=host,
@@ -28,7 +34,7 @@ def list_tables(host: str, port: int, database: str, user: str, password: str) -
def fetch_table(host: str, port: int, database: str, user: str, password: str, table: str, limit: int = 1000) -> pd.DataFrame:
conn = connect(host, port, database, user, password)
cursor = conn.cursor()
- cursor.execute(f"SELECT * FROM {table} LIMIT {limit}")
+ cursor.execute(f"SELECT * FROM {_validate_identifier(table)} LIMIT {limit}")
columns = [desc[0] for desc in cursor.description]
rows = cursor.fetchall()
cursor.close()
diff --git a/src/connectors/redshift.py b/src/connectors/redshift.py
new file mode 100644
index 0000000..73f1a3a
--- /dev/null
+++ b/src/connectors/redshift.py
@@ -0,0 +1,36 @@
+import re
+import pandas as pd
+import redshift_connector
+
+def _validate_identifier(name: str) -> str:
+ if not re.match(r'^[A-Za-z0-9_\.]+$', name):
+ raise ValueError(f"Invalid identifier: {name!r}")
+ return name
+
+def connect(host: str, port: int, database: str, user: str, password: str):
+ return redshift_connector.connect(
+ host=host,
+ port=port,
+ database=database,
+ user=user,
+ password=password
+ )
+
+def list_tables(host: str, port: int, database: str, user: str, password: str) -> list:
+ conn = connect(host, port, database, user, password)
+ cursor = conn.cursor()
+ cursor.execute("SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname='public' ORDER BY tablename")
+ tables = [row[0] for row in cursor.fetchall()]
+ cursor.close()
+ conn.close()
+ return tables
+
+def fetch_table(host: str, port: int, database: str, user: str, password: str, table: str, limit: int = 1000) -> pd.DataFrame:
+ conn = connect(host, port, database, user, password)
+ cursor = conn.cursor()
+ cursor.execute(f"SELECT * FROM {_validate_identifier(table)} LIMIT {limit}")
+ columns = [desc[0] for desc in cursor.description]
+ rows = cursor.fetchall()
+ cursor.close()
+ conn.close()
+ return pd.DataFrame(rows, columns=columns)
diff --git a/src/connectors/snowflake_conn.py b/src/connectors/snowflake_conn.py
index ce74ff1..05fa10c 100644
--- a/src/connectors/snowflake_conn.py
+++ b/src/connectors/snowflake_conn.py
@@ -1,6 +1,12 @@
+import re
import pandas as pd
import snowflake.connector
+def _validate_identifier(name: str) -> str:
+ if not re.match(r'^[A-Za-z0-9_\.]+$', name):
+ raise ValueError(f"Invalid identifier: {name!r}")
+ return name
+
def connect(account: str, user: str, password: str, database: str, schema: str):
conn = snowflake.connector.connect(
account=account,
@@ -23,7 +29,7 @@ def list_tables(account: str, user: str, password: str, database: str, schema: s
def fetch_table(account: str, user: str, password: str, database: str, schema: str, table: str, limit: int = 1000) -> pd.DataFrame:
conn = connect(account, user, password, database, schema)
cursor = conn.cursor()
- cursor.execute(f"SELECT * FROM {database}.{schema}.{table} LIMIT {limit}")
+ cursor.execute(f"SELECT * FROM {_validate_identifier(database)}.{_validate_identifier(schema)}.{_validate_identifier(table)} LIMIT {limit}")
columns = [desc[0] for desc in cursor.description]
rows = cursor.fetchall()
cursor.close()
diff --git a/src/pipeline.py b/src/pipeline.py
index bfa2f9e..d7226cd 100644
--- a/src/pipeline.py
+++ b/src/pipeline.py
@@ -1,110 +1,106 @@
-import os
-import json
-import pandas as pd
-from rich.console import Console
-from rich.progress import Progress, SpinnerColumn, TextColumn
-from rich.table import Table
-from rich.panel import Panel
-from src.models import PipelineResult
-from src.agents import cleaner, validator, transformer, anomaly, summariser
-
-console = Console()
-
-def load_csv(file_path: str) -> tuple[pd.DataFrame, str, int]:
- df = pd.read_csv(file_path)
- total_rows = len(df)
- preview = df.head(20).to_csv(index=False)
- return df, preview, total_rows
-
-def run_pipeline(file_path: str) -> PipelineResult:
- file_name = os.path.basename(file_path)
-
- console.print(Panel.fit(
- f"[bold green]MULTI-AGENT DATA PIPELINE[/bold green]\n[dim]Processing: {file_name}[/dim]",
- border_style="green"
- ))
-
- df, preview, total_rows = load_csv(file_path)
- console.print(f"\n[cyan]→ Loaded {total_rows} rows from {file_name}[/cyan]\n")
-
- with Progress(
- SpinnerColumn(),
- TextColumn("[progress.description]{task.description}"),
- console=console
- ) as progress:
-
- task1 = progress.add_task("[green]Agent 1/5 — Cleaner...", total=None)
- cleaner_result = cleaner.run(preview, total_rows)
- progress.update(task1, description="[green]Agent 1/5 — Cleaner ✓")
- progress.stop_task(task1)
-
- task2 = progress.add_task("[blue]Agent 2/5 — Validator...", total=None)
- validator_result = validator.run(preview, total_rows)
- progress.update(task2, description="[blue]Agent 2/5 — Validator ✓")
- progress.stop_task(task2)
-
- task3 = progress.add_task("[yellow]Agent 3/5 — Transformer...", total=None)
- transformer_result = transformer.run(preview, total_rows)
- progress.update(task3, description="[yellow]Agent 3/5 — Transformer ✓")
- progress.stop_task(task3)
-
- task4 = progress.add_task("[red]Agent 4/5 — Anomaly Detector...", total=None)
- anomaly_result = anomaly.run(preview, total_rows)
- progress.update(task4, description="[red]Agent 4/5 — Anomaly Detector ✓")
- progress.stop_task(task4)
-
- context = f"""
- Cleaner found {len(cleaner_result.issues_fixed)} issues affecting {cleaner_result.rows_affected} rows.
- Validator score: {validator_result.completeness_score}% completeness, {len(validator_result.violations)} violations.
- Transformer applied {len(transformer_result.transformations_applied)} transformations.
- Anomaly detector found {anomaly_result.anomaly_count} anomalies with risk score {anomaly_result.anomaly_score}/10.
- """
-
- task5 = progress.add_task("[magenta]Agent 5/5 — Summariser...", total=None)
- summariser_result = summariser.run(preview, total_rows, context)
- progress.update(task5, description="[magenta]Agent 5/5 — Summariser ✓")
- progress.stop_task(task5)
-
- result = PipelineResult(
- file_name=file_name,
- total_rows=total_rows,
- cleaner=cleaner_result,
- validator=validator_result,
- transformer=transformer_result,
- anomaly=anomaly_result,
- summariser=summariser_result,
- status="complete"
- )
-
- _print_summary(result)
- return result
-
-def _print_summary(result: PipelineResult):
- console.print("\n")
- table = Table(title="Pipeline Results", border_style="green")
- table.add_column("Agent", style="cyan")
- table.add_column("Result", style="white")
-
- table.add_row(
- "Cleaner",
- f"{result.cleaner.rows_affected} rows fixed, {len(result.cleaner.issues_fixed)} issues"
- )
- table.add_row(
- "Validator",
- f"{result.validator.completeness_score}% complete, {len(result.validator.violations)} violations"
- )
- table.add_row(
- "Transformer",
- f"{result.transformer.rows_transformed} rows transformed, {len(result.transformer.new_columns)} new columns"
- )
- table.add_row(
- "Anomaly",
- f"{result.anomaly.anomaly_count} anomalies, risk score {result.anomaly.anomaly_score}/10"
- )
- table.add_row(
- "Summariser",
- f"{len(result.summariser.recommendations)} recommendations"
- )
-
- console.print(table)
- console.print(f"\n[bold green]✓ Pipeline complete — {result.total_rows} rows processed[/bold green]\n")
\ No newline at end of file
+import os
+import pandas as pd
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich.table import Table
+from rich.panel import Panel
+from src.models import PipelineResult
+from src.agents import cleaner, validator, transformer, anomaly, summariser
+from src.connectors import mongodb
+
+console = Console()
+
+def load_csv(file_path: str) -> tuple[pd.DataFrame, str, int]:
+ df = pd.read_csv(file_path)
+ total_rows = len(df)
+ preview = df.head(20).to_csv(index=False)
+ return df, preview, total_rows
+
+def load_mongodb(uri: str, database: str, collection: str, limit: int = 1000) -> tuple[pd.DataFrame, str, int]:
+ df = mongodb.fetch_collection(uri, database, collection, limit)
+ total_rows = len(df)
+ preview = df.head(20).to_csv(index=False)
+ return df, preview, total_rows
+
+def run_pipeline(file_path: str) -> PipelineResult:
+ file_name = os.path.basename(file_path)
+ df, preview, total_rows = load_csv(file_path)
+ return _run_pipeline_from_data(file_name, preview, total_rows)
+
+def run_pipeline_mongo(uri: str, database: str, collection: str, limit: int = 1000) -> PipelineResult:
+ df, preview, total_rows = load_mongodb(uri, database, collection, limit)
+ return _run_pipeline_from_data(f"{database}.{collection}", preview, total_rows)
+
+def _run_pipeline_from_data(source_label: str, preview: str, total_rows: int) -> PipelineResult:
+ console.print(Panel.fit(
+ f"[bold green]MULTI-AGENT DATA PIPELINE[/bold green]\n[dim]Processing: {source_label}[/dim]",
+ border_style="green"
+ ))
+ console.print(f"\n[cyan]→ Loaded {total_rows} rows from {source_label}[/cyan]\n")
+
+ with Progress(
+ SpinnerColumn(),
+ TextColumn("[progress.description]{task.description}"),
+ console=console
+ ) as progress:
+
+ task1 = progress.add_task("[green]Agent 1/5 — Cleaner...", total=None)
+ cleaner_result = cleaner.run(preview, total_rows)
+ progress.update(task1, description="[green]Agent 1/5 — Cleaner ✓")
+ progress.stop_task(task1)
+
+ task2 = progress.add_task("[blue]Agent 2/5 — Validator...", total=None)
+ validator_result = validator.run(preview, total_rows)
+ progress.update(task2, description="[blue]Agent 2/5 — Validator ✓")
+ progress.stop_task(task2)
+
+ task3 = progress.add_task("[yellow]Agent 3/5 — Transformer...", total=None)
+ transformer_result = transformer.run(preview, total_rows)
+ progress.update(task3, description="[yellow]Agent 3/5 — Transformer ✓")
+ progress.stop_task(task3)
+
+ task4 = progress.add_task("[red]Agent 4/5 — Anomaly Detector...", total=None)
+ anomaly_result = anomaly.run(preview, total_rows)
+ progress.update(task4, description="[red]Agent 4/5 — Anomaly Detector ✓")
+ progress.stop_task(task4)
+
+ context = f"""
+ Cleaner found {len(cleaner_result.issues_fixed)} issues affecting {cleaner_result.rows_affected} rows.
+ Validator score: {validator_result.completeness_score}% completeness, {len(validator_result.violations)} violations.
+ Transformer applied {len(transformer_result.transformations_applied)} transformations.
+ Anomaly detector found {anomaly_result.anomaly_count} anomalies with risk score {anomaly_result.anomaly_score}/10.
+ """
+
+ task5 = progress.add_task("[magenta]Agent 5/5 — Summariser...", total=None)
+ summariser_result = summariser.run(preview, total_rows, context)
+ progress.update(task5, description="[magenta]Agent 5/5 — Summariser ✓")
+ progress.stop_task(task5)
+
+ result = PipelineResult(
+ file_name=source_label,
+ total_rows=total_rows,
+ cleaner=cleaner_result,
+ validator=validator_result,
+ transformer=transformer_result,
+ anomaly=anomaly_result,
+ summariser=summariser_result,
+ status="complete"
+ )
+
+ _print_summary(result)
+ return result
+
+def _print_summary(result: PipelineResult):
+ console.print("\n")
+ table = Table(title="Pipeline Results", border_style="green")
+ table.add_column("Agent", style="cyan")
+ table.add_column("Result", style="white")
+
+ table.add_row("Cleaner", f"{result.cleaner.rows_affected} rows fixed, {len(result.cleaner.issues_fixed)} issues")
+ table.add_row("Validator", f"{result.validator.completeness_score}% complete, {len(result.validator.violations)} violations")
+ table.add_row("Transformer", f"{result.transformer.rows_transformed} rows transformed, {len(result.transformer.new_columns)} new columns")
+ table.add_row("Anomaly", f"{result.anomaly.anomaly_count} anomalies, risk score {result.anomaly.anomaly_score}/10")
+ table.add_row("Summariser", f"{len(result.summariser.recommendations)} recommendations")
+
+ console.print(table)
+ console.print(f"\n[bold green]✓ Pipeline complete — {result.total_rows} rows processed[/bold green]\n")