Promptimus 🧠

A PyTorch-like API for building composable LLM agents with advanced tool calling, memory management, and observability.

✨ Key Features

• 🧠 PyTorch-like Modules: Composable agent architecture with hierarchical module system
• 🔧 Tool Calling: ReACT-style and native OpenAI function calling with automatic schema generation
• 📝 Structured Output: Pydantic schema-based JSON generation with validation
• 💾 Memory Management: Conversation context with configurable memory limits
• 🔍 Embeddings: Text embedding generation with batch processing support
• 🗄️ Vector Stores: ChromaDB integration with async-first vector operations
• 🤖 RAG (Retrieval-Augmented Generation): Retrieval system with conversation memory
• 📊 Tracing: Arize Phoenix integration for observability with token usage and cost tracking
• 💾 Serialization: TOML-based save/load for prompts and module configurations
• ⚡ Async First: Built for high-performance asynchronous operations

🚀 Quick Start

Installation

pip install promptimus

Basic Example

import promptimus as pm

# Create an LLM provider
llm = pm.llms.OpenAILike(
    model_name="gpt-4",
    api_key="your-api-key"
)

# Create a simple agent with memory
agent = pm.modules.MemoryModule(
    memory_size=5,
    system_prompt="You are a helpful assistant."
).with_llm(llm)

# Have a conversation
response1 = await agent.forward("Hi, I'm Alice!")
response2 = await agent.forward("What's my name?")
print(response2.content)  # "Your name is Alice!"

LLM Providers

Any provider is interchangeable — pass it to .with_llm(...) on the root module.

# OpenAI-compatible endpoints (OpenAI, OpenRouter, vLLM, ...)
llm = pm.llms.OpenAILike(model_name="gpt-4.1-nano")  # api_key from OPENAI_API_KEY

# Anthropic Claude (wraps the official `anthropic` SDK)
llm = pm.llms.Claude(model_name="claude-opus-4-8")  # api_key from ANTHROPIC_API_KEY
# extended thinking → text lands in Message.reasoning:
llm = pm.llms.Claude(
    model_name="claude-opus-4-8",
    call_kwargs={"thinking": {"type": "adaptive"}},
)

# Local Ollama
llm = pm.llms.OllamaProvider(model_name="gemma3:4b", base_url="http://localhost:11434/v1")

The Claude provider translates the OpenAI-shaped tools / response_format kwargs emitted by OpenaiToolCallingAgent and StructuralOutput, so every module works against Claude unchanged.

🏗️ Architecture

Core Concepts

Modules: Container system for organizing prompts, submodules, and logic

class MyAgent(pm.Module):
    def __init__(self):
        super().__init__()
        self.chat = pm.Prompt("You are a helpful assistant")
        self.memory = []

    async def forward(self, message: str) -> str:
        # Custom logic here
        pass

Prompts: Parameter-like system for system prompts (similar to PyTorch parameters)

prompt = pm.Prompt("You are a {role} assistant").with_llm(llm)
response = await prompt.forward(role="helpful")

Tools: Function decoration for external capabilities

@pm.modules.Tool.decorate
def calculate(a: float, b: float, operation: str) -> float:
    """Calculate result of operation on two numbers."""
    if operation == "add":
        return a + b
    # ... more operations

Pre-built Modules

Memory Module: Conversation memory with configurable limits

agent = pm.modules.MemoryModule(
    memory_size=10,
    system_prompt="You are a helpful assistant."
).with_llm(llm)

Retrieval Module: Hybrid search with vector and text stores

retrieval = pm.modules.RetrievalModule(
    vector_store=vector_store,
    n_semantic=5,
).with_embedder(embedder)

# Insert documents
await retrieval.insert(documents)

# Search for relevant content
results = await retrieval.forward("query about AI")

RAG Module: Retrieval-Augmented Generation with conversation memory

import chromadb
from chromadb_store import ChromaVectorStore

# Setup components
embedder = pm.embedders.OpenAILikeEmbedder(model_name="text-embedding-3-small")
client = chromadb.EphemeralClient()
vector_store = ChromaVectorStore(client, "my_docs")

# Create RAG agent
rag_agent = pm.modules.RAGModule(
    vector_store=vector_store,
    n_semantic=3,
    memory_size=5
).with_llm(llm).with_embedder(embedder)

# Add documents
await rag_agent.retrieval.insert([
    "Machine learning is a subset of AI...",
    "Deep learning uses neural networks...",
    # ... more documents
])

# Query with context
response = await rag_agent.forward("What is machine learning?")

Structural Output: Pydantic schema-based JSON generation

from pydantic import BaseModel

class Person(BaseModel):
    name: str
    age: int
    occupation: str

module = pm.modules.StructuralOutput(Person).with_llm(llm)
result = await module.forward("Extract info about John, a 30-year-old engineer")

Tool Calling Agents: Agents that can use tools autonomously

agent = pm.modules.ToolCallingAgent([
    calculate,
    # ... more tools
]).with_llm(llm)

result = await agent.forward("What is 15 + 27?")

🔧 Advanced Features

Serialization

Save and load module configurations:

agent.save("my_agent.toml")
loaded_agent = pm.modules.MemoryModule().load("my_agent.toml")

Config Composition

Split complex TOML configs into composable files using OmegaConf-inspired resolver syntax:

# main.toml — lean orchestrator config
analysis_format = """
<query>{query}</query>
"""

# Reference submodule configs from installed packages
query_decomposer = "${pkg:my_lib.prompts.query_decomposer.toml}"

# Or from relative file paths
series_analyzer = "${file:./series_analyzer.toml}"

References are resolved at load time and inlined on save — save() always produces a single flat TOML.

Tracing with Phoenix

import phoenix as px
from phoenix_tracer import trace

px.launch_app()
trace(agent, "my_agent", project_name="my_project")

# With cost tracking
trace(
    agent, "my_agent",
    pricing={
        "gpt-4.1": (2.0, 8.0),       # (input $/M tokens, output $/M tokens)
        "gpt-4.1-mini": (0.4, 1.6),
    },
    project_name="my_project",
)

Vector Stores

import chromadb
from chromadb_store import ChromaVectorStore

# Setup ChromaDB vector store
client = chromadb.PersistentClient(path="./chroma_db")
vector_store = ChromaVectorStore(client, "my_collection")

# Create embedder
embedder = pm.embedders.OpenAILikeEmbedder(
    model_name="text-embedding-3-small"
)

# Build RAG system
rag = pm.modules.RAGModule(
    vector_store=vector_store,
    n_semantic=5,
    memory_size=10
).with_llm(llm).with_embedder(embedder)

# Add documents
await rag.retrieval.insert([
    "Document 1 content...",
    "Document 2 content...",
])

# Query with retrieval-augmented generation
response = await rag.forward("What information do you have about X?")

Custom Embedders

embedder = pm.embedders.OpenAILikeEmbedder(
    model_name="text-embedding-3-small"
)

embeddings = await embedder.aembed_batch([
    "Hello world",
    "How are you?"
])

📖 Documentation

Tutorials

Explore our comprehensive notebook tutorials:

1. LLM Providers & Embedders - Getting started with providers
2. Prompts & Modules - Core architecture concepts
3. Pre-built Modules - Ready-to-use components including RAG
4. Custom Agents - Tool calling and advanced agents
5. Tracing - Observability with Phoenix

API Reference

• pm.Module: Base class for all modules
• pm.Prompt: System prompt management
• pm.llms.*: LLM provider implementations
• pm.embedders.*: Embedding provider implementations
• pm.vectore_store.*: Vector store protocols and implementations
• pm.modules.*: Pre-built module components
  • MemoryModule: Conversation memory management
  • RAGModule: Retrieval-Augmented Generation
  • RetrievalModule: Vector database operations
  • StructuralOutput: Schema-based JSON generation
  • ToolCallingAgent: Tool-augmented agents

🛠️ Installation Options

Basic Installation

pip install promptimus

With Optional Dependencies

# Phoenix tracing support
pip install promptimus[phoenix]

# ChromaDB vector store for RAG
pip install promptimus[chromadb]

# All optional dependencies
pip install promptimus[all]

Development Setup

git clone https://github.com/AIladin/promptimus.git
cd promptimus
pip install -e .[dev]

🙏 Acknowledgments

• Inspired by PyTorch's modular architecture
• Built on top of modern Python async patterns
• Integrated with Arize Phoenix for tracing
• Compatible with OpenAI and OpenAI-compatible APIs
• Vector store support powered by ChromaDB

Ready to build your next LLM agent? Check out our tutorials to get started! 🚀