aiforge-kit

AIForge — Plug-and-Play AI Toolkit for Product Teams

AIForge is a modular, provider-agnostic Python SDK that gives product teams production-ready AI capabilities with a single pip install and a config file. Each module — summarization, document Q&A, extraction, classification, semantic search, content generation, conversational agents, content moderation, translation, and vision — is self-contained, composable, and ships with a built-in Streamlit UI widget for instant prototyping.

Enable any developer to add production-quality AI capabilities to their product in minutes, not weeks.

---

Key Features

---

Technology Stack

---

Quick Start

Prerequisites

• Python 3.12+
• An API key for at least one LLM provider (OpenAI, Anthropic, Gemini, Mistral, AWS credentials for Bedrock, or Ollama for local)

1. Clone & Set Up Environment

git clone <repo-url>
cd aiforge

# Create and activate a Python virtual environment
python -m venv venv
source venv/bin/activate        # macOS / Linux
# venv\Scripts\activate         # Windows

# Install dependencies
pip install -r requirements.txt

2. Configure

# Copy example config files
cp .env.example .env
cp aiforge.yaml.example aiforge.yaml

# Edit .env and set your API key
# AIFORGE_OPENAI_API_KEY=sk-...

3. Use in Code

from aiforge import Summarizer

s = Summarizer()
result = s.run("AIForge is a plug-and-play AI toolkit that gives teams...")
print(result.summary)
print(f"Cost: ${result.cost:.4f}")

4. Launch the Playground UI

aiforge playground
# Opens Streamlit app at http://localhost:8501

5. Start the API Server

aiforge serve
# REST API at http://localhost:8000
# Swagger docs at http://localhost:8000/docs

---

Installation Options

# From PyPI (the distribution is aiforge-kit; the import name is aiforge)
pip install aiforge-kit                          # Core + OpenAI provider + ChromaDB
pip install "aiforge-kit[anthropic,mcp]"         # Pick extras as needed

# From a clone
pip install -r requirements.txt
pip install -e ".[all]"                          # Editable install for development

# Extras (combine with commas, e.g. "aiforge-kit[gemini,server]")
pip install "aiforge-kit[anthropic]"      # Anthropic + Voyage AI embeddings
pip install "aiforge-kit[gemini]"         # Google Gemini
pip install "aiforge-kit[bedrock]"        # Amazon Bedrock
pip install "aiforge-kit[mistral]"        # Mistral AI
pip install "aiforge-kit[local]"          # Ollama (local models)
pip install "aiforge-kit[pinecone]"       # Pinecone vector store
pip install "aiforge-kit[qdrant]"         # Qdrant vector store
pip install "aiforge-kit[pgvector]"       # PostgreSQL + pgvector
pip install "aiforge-kit[server]"         # FastAPI REST server
pip install "aiforge-kit[playground]"     # Streamlit playground
pip install "aiforge-kit[agentic]"        # LangGraph agentic RAG / chat agent
pip install "aiforge-kit[observability]"  # OpenTelemetry + LangSmith
pip install "aiforge-kit[docs]"           # PDF and DOCX document loaders
pip install "aiforge-kit[mcp]"            # MCP client (stdio + Streamable HTTP)
pip install "aiforge-kit[all]"            # Everything

For development, uv sync --all-extras installs the exact versions CI tests against from uv.lock. CI also runs the suite on the lowest versions pyproject.toml allows, and Dependabot proposes lockfile updates weekly.

---

Lint and Type checking

# Lint checker
ruff check .
# Output the lint findings to a file
ruff check . > ruff_output.txt

# Type checker
pyright .

---

Modules

Summarize — Text Summarization

from aiforge import Summarizer

s = Summarizer(style="concise", max_length=200)
result = s.run("Long article text here...")
print(result.summary)

# Long documents auto-handled via map-reduce
result = s.run_file("report.pdf")

DocQA — Document Q&A (RAG)

from aiforge import DocQA

qa = DocQA(collection="my_docs")
qa.ingest("./docs/")                              # ingest a directory
answer = qa.ask("How does authentication work?")
print(answer.text)
print(answer.sources)                              # source citations

Hybrid retrieval, reranking and heading-aware chunking are opt-in:

qa = DocQA(
    collection="my_docs",
    chunking="markdown",        # split on headings ("recursive" | "markdown" | "token")
    retrieval_mode="hybrid",    # BM25 keyword + vector, fused ("vector" | "hybrid" | "keyword")
    rerank=True,                # LLM reranker over the fused candidates
    rerank_top_n=4,
)

SemanticSearch takes the same retrieval_mode, rerank and chunking options.

Extract — Structured Data Extraction

from aiforge import Extractor
from pydantic import BaseModel

class Invoice(BaseModel):
    vendor: str
    total: float
    currency: str

e = Extractor(output_schema=Invoice)
result = e.run("Invoice from Acme Corp. Total: $86.00 USD.")
print(result.data)  # Invoice(vendor="Acme Corp", total=86.0, currency="USD")

Classify — Text Classification

from aiforge import Classifier

c = Classifier(labels=["bug", "feature", "question"])
result = c.run("The login button is broken on mobile")
print(result.label)       # "bug"
print(result.confidence)  # 0.94

Search — Semantic Search

from aiforge import SemanticSearch

ss = SemanticSearch(collection="products")
ss.index([
    {"id": "1", "text": "Wireless headphones", "metadata": {"price": 299}},
    {"id": "2", "text": "Bluetooth speaker", "metadata": {"price": 79}},
])
results = ss.query("something to listen to music", top_k=5)

Generate — Template-Based Content Generation

from aiforge import Generator

g = Generator(
    template="Write a {tone} product description for: {product_name}",
    guardrails=["no_pii", "professional_tone"],
)
result = g.run(product_name="AIForge", tone="professional")

Chat — Conversational Agent

from aiforge import ChatAgent

agent = ChatAgent(
    system_prompt="You are a helpful assistant.",
    memory="buffer",
)
response = agent.chat("What's the return policy?")
response = agent.chat("What about international orders?")  # remembers context

Moderate — Content Safety & PII Detection

from aiforge import Moderator

m = Moderator(checks=["pii", "toxicity", "prompt_injection"])
result = m.check("My email is john@example.com")
print(result.flagged)  # True

cleaned = m.redact("Call me at 555-0123")
print(cleaned.text)    # "Call me at [PHONE]"

Guardrails — Input & Output Checks

from aiforge import Summarizer
from aiforge.guardrails import Guard, PIIGuard, PromptInjectionGuard

guard = Guard(
    input=[PIIGuard(), PromptInjectionGuard()],
    output=[PIIGuard()],
    on_fail="redact",            # "block" (default) | "redact" | "warn"
)
safe_summarize = guard.wrap(Summarizer())
result = safe_summarize.run("Email jane@example.com about the outage...")

See Guardrails for all checks and redaction styles.

Evals — Regression Testing for Prompts and Models

from aiforge import Classifier
from aiforge.evals import evaluate, label_match

report = evaluate(
    Classifier(labels=["bug", "feature", "question"]),
    "evals/tickets.jsonl",       # or a list of {"input": ..., "expected": ...}
    scorers=[label_match()],
    min_pass_rate=0.9,
)
report.print()
assert report.passed

Or from CI: aiforge eval run evals/tickets.jsonl -m classify -s label_match --min-pass-rate 0.9. See Evaluations.

---

Pipeline Composition

Chain modules into workflows:

from aiforge import Pipeline, Extractor, Classifier, Summarizer

pipe = Pipeline([
    ("extract", Extractor(output_schema=TicketSchema)),
    ("classify", Classifier(labels=["critical", "high", "medium", "low"])),
    ("summarize", Summarizer(style="concise")),
])

result = pipe.run("Long customer complaint email...")
print(result["classify"].label)       # "critical"
print(result["summarize"].summary)    # One-line summary

Conditional routing:

pipe = Pipeline([
    ("classify", Classifier(labels=["support", "sales", "spam"])),
    ("route", {
        "support": Summarizer(style="bullet_points"),
        "sales": Generator(template="..."),
        "spam": None,  # drop
    }),
])

---

Configuration

AIForge uses a layered configuration system:

1. Defaults (sensible built-in values)
        ↓ overridden by
2. aiforge.yaml (project-level config)
        ↓ overridden by
3. .env file / environment variables (AIFORGE_*)
        ↓ overridden by
4. Constructor arguments (code-level)

Key Environment Variables

See .env.example for the full list.

Swapping Providers

# aiforge.yaml — change these 2 lines
provider: anthropic
model: claude-sonnet-5-5

Or at runtime:

s = Summarizer(provider="anthropic", model="claude-sonnet-5-5")

---

REST API

Start with aiforge serve — all modules are exposed as HTTP endpoints:

Example:

curl -X POST http://localhost:8000/classify \
  -H "Content-Type: application/json" \
  -d '{"text": "The login button is broken", "labels": ["bug", "feature", "question"]}'

---

Playground UI

Interactive Streamlit app with tabs for each module:

aiforge playground
# http://localhost:8501

Tabs: Summarize | Q&A | Extract | Classify | Search | Generate | Chat | Moderate | Translate | Vision | Pipeline | Settings

---

Docker

# Build and run
docker compose up --build

# Services:
#   api:        http://localhost:8000  (REST API)
#   playground:  http://localhost:8501  (Streamlit UI)

---

CLI Reference

aiforge init          # Scaffold aiforge.yaml + .env.example
aiforge serve         # Start FastAPI server (default: 0.0.0.0:8000)
aiforge playground    # Start Streamlit playground (default: port 8501)
aiforge migrate       # Migrate vector store collections

aiforge summarize --text "Long text..."      # Every module has a command
aiforge qa ingest ./docs && aiforge qa ask "How does auth work?"
aiforge eval run evals/cases.jsonl -m classify -s label_match --min-pass-rate 0.9
aiforge eval compare report.json main-report.json  # exit 3 on regression
aiforge models list                          # Models, pricing, cost estimates

See the CLI reference for every command and flag.

---

Project Structure

aiforge/
├── pyproject.toml                     # Package config, dependencies, extras
├── requirements.txt                   # Pinned dependencies for venv development
├── Makefile                           # Dev commands
├── Dockerfile                         # Multi-stage production build
├── docker-compose.yml                 # API + Playground services
├── aiforge.yaml.example              # Example configuration
├── .env.example                       # Example environment variables
│
├── aiforge/                           # SDK source code
│   ├── __init__.py                    # Public API exports
│   ├── __version__.py                 # Version string
│   ├── registry.py                    # Plugin registry (modules, providers, vector stores)
│   ├── core/                          # Provider-agnostic abstractions
│   │   ├── types.py                   # BaseModule, CompletionResult, error hierarchy
│   │   ├── provider.py               # LLMProvider, EmbeddingProvider ABCs
│   │   ├── vectorstore.py            # VectorStoreAdapter + ChromaDB default
│   │   ├── vectorstore_pinecone.py    # PineconeVectorStore adapter
│   │   ├── vectorstore_qdrant.py      # QdrantVectorStore adapter
│   │   ├── vectorstore_pgvector.py    # PgVectorStore adapter (PostgreSQL + pgvector)
│   │   ├── vectorstore_factory.py     # Config-driven vector store factory
│   │   ├── doc_loader.py             # Document loading (files, directories, git)
│   │   ├── text_splitter.py          # Recursive, markdown and token chunking
│   │   ├── retrieval.py              # BM25 index + hybrid retrieval (RRF)
│   │   ├── rerank.py                 # LLM and score-function rerankers
│   │   └── pricing.py                # Token pricing tables
│   ├── infra/                         # Cross-cutting infrastructure
│   │   ├── config.py                  # AIForgeSettings (Pydantic Settings)
│   │   ├── retry.py                   # Tenacity retry wrappers
│   │   ├── rate_limiter.py            # Token-bucket rate limiter
│   │   ├── cost_tracker.py            # Session-level cost aggregation
│   │   ├── tracing.py                 # LangSmith / OpenTelemetry setup
│   │   ├── logging.py                 # Structured logging (structlog)
│   │   └── streaming.py              # Async streaming helpers
│   ├── providers/                     # Concrete provider implementations
│   │   ├── _langchain.py              # Shared LangChain provider base classes
│   │   ├── openai.py                  # OpenAI (GPT-4o, text-embedding-3-small)
│   │   ├── anthropic.py              # Anthropic (Claude) + Voyage embeddings
│   │   ├── gemini.py                  # Google Gemini
│   │   ├── bedrock.py                 # Amazon Bedrock (Converse API)
│   │   ├── mistral.py                 # Mistral AI
│   │   ├── ollama.py                  # Ollama (Llama 3, Mistral — local)
│   │   └── azure.py                   # Azure OpenAI
│   ├── modules/                       # 10 AI modules
│   │   ├── summarize/                 # Text summarization
│   │   ├── qa/                        # Document Q&A (RAG)
│   │   ├── extract/                   # Structured data extraction
│   │   ├── classify/                  # Text classification
│   │   ├── search/                    # Semantic search
│   │   ├── generate/                  # Content generation
│   │   ├── chat/                      # Conversational agent (LangGraph)
│   │   ├── moderate/                  # Content moderation & PII detection
│   │   ├── translate/                 # Text translation
│   │   └── vision/                    # Vision / image understanding
│   ├── evals/                         # Datasets, scorers, eval runner, reports
│   ├── guardrails/                    # Guard, input/output checks, PII redaction
│   ├── mcp/                           # MCP client + tool bridge
│   ├── pipeline/                      # Pipeline composition engine
│   ├── server/                        # FastAPI REST API
│   ├── playground/                    # Streamlit UI
│   └── cli/                           # CLI (Typer)
│
├── tests/                             # Test suite
│   ├── conftest.py                    # Shared fixtures
│   ├── unit/                          # Unit tests per module
│   ├── integration/                   # Integration tests (real API calls)
│   └── fixtures/                      # Sample texts, documents, recorded responses
│
└── docs/                              # Documentation
    ├── quickstart.md
    ├── configuration.md
    ├── providers.md
    ├── pipelines.md
    ├── guardrails.md
    ├── evals.md
    ├── cost-tracking.md
    ├── cli.md
    └── modules/                       # Per-module docs

---

Development

Setup

# Clone the repo
git clone <repo-url>
cd aiforge

# Create Python virtual environment
python -m venv venv
source venv/bin/activate

# Install all dependencies (including dev + all extras)
pip install -e ".[all]"

# Copy config templates
cp .env.example .env
cp aiforge.yaml.example aiforge.yaml

Running Tests

# Run all tests
pytest

# Run with coverage
pytest --cov=aiforge --cov-report=term-missing

# Run specific test file
pytest tests/unit/test_summarize.py

# Run specific test
pytest tests/unit/test_config.py::test_yaml_loading

Code Quality

# Linting
ruff check aiforge/ tests/

# Type checking
mypy aiforge/

# Format
ruff format aiforge/ tests/

Make Commands

make install         # Install dependencies
make test            # Run tests
make lint            # Run linter
make serve           # Start API server
make playground      # Start Streamlit UI
make docker          # Build Docker image

---

Supported Providers

See Providers for setup, default models and install extras.

---

Extending AIForge

AIForge ships with a plugin system powered by a unified registry. Register custom modules, providers, and vector stores — then use them interchangeably with built-ins.

Custom Module

from aiforge.core.types import BaseModule, ModuleResponse
from aiforge.registry import register_module

class SentimentAnalyzer(BaseModule):
    def run(self, text: str, **kwargs) -> ModuleResponse:
        result = self._call_llm(f"Classify sentiment: {text}")
        return ModuleResponse(text=result.text.strip(), cost=result.cost)

    async def arun(self, text: str, **kwargs) -> ModuleResponse:
        result = await self._acall_llm(f"Classify sentiment: {text}")
        return ModuleResponse(text=result.text.strip(), cost=result.cost)

register_module("sentiment", SentimentAnalyzer)

Custom Provider

from aiforge.core.provider import LLMProvider, EmbeddingProvider
from aiforge.registry import register_provider

class MyProvider(LLMProvider):
    def complete(self, messages, **kwargs): ...
    async def complete_async(self, messages, **kwargs): ...
    async def stream(self, messages, **kwargs): ...

class MyEmbeddingProvider(EmbeddingProvider):
    def embed(self, texts): ...
    def embed_query(self, text): ...

register_provider("my_provider", MyProvider, MyEmbeddingProvider)

Custom Vector Store

from aiforge.core.vectorstore import VectorStoreAdapter
from aiforge.registry import register_vectorstore

class WeaviateStore(VectorStoreAdapter):
    def add_documents(self, texts, embeddings, metadatas, ids): ...
    def query(self, embedding, top_k, filter, include_embeddings): ...
    def delete_collection(self, name): ...
    def list_collections(self): ...
    def collection_stats(self, name): ...

register_vectorstore("weaviate", WeaviateStore)

Registry Discovery

from aiforge.registry import list_modules, list_providers, list_vectorstores

print(list_modules())       # ['chat', 'classify', ..., 'summarize']
print(list_providers())     # ['openai', 'anthropic', 'ollama', 'azure', 'gemini', 'bedrock', 'mistral']
print(list_vectorstores())  # ['chroma', 'pgvector', 'pinecone', 'qdrant']

See Plugin System docs for full details.

---

Cost Tracking

Every module call returns cost information:

result = summarizer.run(text)
print(f"Cost: ${result.cost:.4f}")
print(f"Tokens: {result.input_tokens} in / {result.output_tokens} out")

# Session aggregates
from aiforge import CostTracker
tracker = CostTracker.get_session()
print(f"Total: ${tracker.total_cost:.4f}")
print(tracker.by_module)   # {"summarize": 0.012, "qa": 0.025, ...}
print(tracker.by_model)    # {"gpt-4o": 0.040, "text-embedding-3-small": 0.002}

---

License

MIT — see LICENSE for details.

---

Author

Kushang Patel — Lead Developer