Metadata-Version: 2.4
Name: aiagentrag
Version: 1.1.1
Summary: Reusable AI Agent library with RAG, streaming, and long-term memory
Author: leonidsliusar
Author-email: loslyusar@gmail.com
Requires-Python: >=3.13,<4.0
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Provides-Extra: modal
Requires-Dist: alembic (>=1.14.0,<2.0.0)
Requires-Dist: asyncpg (>=0.30.0,<0.31.0)
Requires-Dist: dishka (>=1.4.0,<2.0.0)
Requires-Dist: httpx (>=0.28.0,<0.29.0)
Requires-Dist: modal (>=0.73.0,<0.74.0) ; extra == "modal"
Requires-Dist: ollama (>=0.4.0,<0.5.0)
Requires-Dist: openai (>=1.57.0,<2.0.0)
Requires-Dist: pydantic (>=2.10.0,<3.0.0)
Requires-Dist: qdrant-client (>=1.12.0,<2.0.0)
Requires-Dist: sqlalchemy[asyncio] (>=2.0.36,<3.0.0)
Description-Content-Type: text/markdown

# AIAgentRag

Python library for an AI agent with RAG, streaming responses, and long-term memory.
Transport-independent — wire it into Telegram, FastAPI, CLI, or anything else on your side.

Compatible with [vectorizer](https://github.com/leonidsliusar/vectorizer): load documents into Qdrant via vectorizer; the agent reads the same `documents` collection.

---

## 1. Dependencies

The library does **not** start external services. You need:

| Service | Purpose | Typical setup |
|---------|---------|---------------|
| **PostgreSQL** | Conversation history | `docker compose up -d` in this repo, or your own instance |
| **Qdrant** | RAG + long-term memory | `docker compose up -d` in this repo, or your own instance |
| **LLM + embeddings** | Generation and vectorization | **Ollama**, **OpenAI API**, or **Modal RPC** |

For RAG:

- **`documents`** collection in Qdrant — created and populated by **vectorizer** (the agent is read-only).
- **`user_memory`** collection — created automatically by the agent.
- With Ollama: embedding model must match vectorizer (`nomic-embed-text`).
- With Modal: your deployed Modal app must expose compatible `embed` RPC (see below).

Python **3.13+**.

---

## 2. Installation

### From PyPI

```bash
pip install aiagentrag
pip install "aiagentrag[modal]"   # Modal RPC provider
```

```bash
poetry add aiagentrag
poetry add aiagentrag --extras modal
```

### From source (development)

```bash
git clone https://github.com/leonidsliusar/AIAgentRag
cd AIAgentRag
poetry install
```

---

## 3. Local script (`scripts/try_agent`)

Script to test the agent against real infrastructure. Lives in this repository only — not shipped on PyPI.

### One-time setup

```bash
docker compose up -d          # PostgreSQL :5432 + Qdrant :6333
ollama pull nomic-embed-text  # embeddings (same as vectorizer)
ollama pull qwen3:8b          # chat model (or set OLLAMA_MODEL)
# ingest documents into Qdrant via vectorizer (collection: documents)
```

### Run

From the repository root:

```bash
poetry run python -m scripts.try_agent -m "Your question"
```

PostgreSQL migrations run automatically inside the script — no manual step required.

### Flags

| Flag | Default | Description |
|------|---------|-------------|
| `-m`, `--message` | — | **Required.** User message |
| `--provider` | `ollama` | `ollama`, `openai`, or `modal` |
| `--user-id` | `local-user` | User identifier |
| `--database-url` | see below | PostgreSQL URL |
| `--qdrant-url` | see below | Qdrant URL |
| `--knowledge-collection` | see below | Document collection (vectorizer) |

### Environment variables

Used when the corresponding flag is not passed:

| Variable | Default |
|----------|---------|
| `DATABASE_URL` | `postgresql+asyncpg://postgres:postgres@localhost:5432/aiagentrag` |
| `QDRANT_URL` | `http://localhost:6333` |
| `QDRANT_COLLECTION` | `documents` |
| `OLLAMA_HOST` | `http://localhost:11434` |
| `OLLAMA_MODEL` | `qwen3:8b` |
| `OLLAMA_EMBED_MODEL` | `nomic-embed-text` |
| `OPENAI_API_KEY` | — (required for `--provider openai`) |
| `OPENAI_MODEL` | `gpt-4o-mini` |
| `OPENAI_EMBED_MODEL` | `text-embedding-3-small` |
| `MODAL_TOKEN_ID` | — (required unless set in `~/.modal.toml`) |
| `MODAL_TOKEN_SECRET` | — (required unless set in `~/.modal.toml`) |
| `MODAL_APP_NAME` | — (required for `--provider modal`) |
| `MODAL_ENVIRONMENT` | — (optional Modal environment) |
| `MODAL_LLM_CLS` | — (Modal class with `complete` / `stream` methods) |
| `MODAL_LLM_COMPLETE_METHOD` | `complete` |
| `MODAL_LLM_STREAM_METHOD` | `stream` |
| `MODAL_EMBED_CLS` | same as `MODAL_LLM_CLS` |
| `MODAL_EMBED_METHOD` | `embed` |
| `MODAL_LLM_COMPLETE_FUNCTION` | — (alternative to class: deployed Function name) |
| `MODAL_LLM_STREAM_FUNCTION` | — |
| `MODAL_EMBED_FUNCTION` | — |

### Examples

```bash
# minimal run
poetry run python -m scripts.try_agent -m "What do the documents say about X?"

# explicit URLs
poetry run python -m scripts.try_agent \
  -m "Hello" \
  --database-url "postgresql+asyncpg://postgres:postgres@localhost:5432/aiagentrag" \
  --qdrant-url "http://localhost:6333" \
  --knowledge-collection documents

# OpenAI
export OPENAI_API_KEY=sk-...
poetry run python -m scripts.try_agent --provider openai -m "Hello"
```

### Modal RPC (full setup)

Modal calls your **deployed** app over RPC. You need **credentials**, a **deployed app**, and **target names**.

#### Step 1 — Modal account and token

1. Register at [modal.com](https://modal.com).
2. Create an API token at [modal.com/settings](https://modal.com/settings).

#### Step 2 — Authenticate locally (pick one)

**Option A — CLI** (writes `~/.modal.toml`):

```bash
pip install modal
modal token set --token-id ak-... --token-secret as-...
```

**Option B — environment variables** (CI, Docker, no config file):

```bash
export MODAL_TOKEN_ID=ak-...
export MODAL_TOKEN_SECRET=as-...
```

Verify:

```bash
modal profile current
```

#### Step 3 — Deploy the LLM service on Modal

Uses the same credentials as above:

```bash
poetry install --extras modal
modal deploy scripts/modal_llm_service.py
```

This deploys app `aiagentrag-llm` with class `LLMService`. Replace method bodies in that file with your real inference code.

#### Step 4 — Run the agent (local machine or server)

Still needs PostgreSQL, Qdrant, and vectorizer data — only LLM/embeddings go through Modal RPC:

```bash
export MODAL_TOKEN_ID=ak-...          # skip if already in ~/.modal.toml
export MODAL_TOKEN_SECRET=as-...      # skip if already in ~/.modal.toml
export MODAL_APP_NAME=aiagentrag-llm  # must match deployed app name
export MODAL_LLM_CLS=LLMService       # must match @app.cls class name

poetry run python -m scripts.try_agent --provider modal -m "Hello"
```

If credentials are missing, the script fails immediately with a clear error (before any RPC call).

#### Modal RPC contract

Modal is used as **RPC** to your deployed app — not HTTP, not Ollama.

Your deployed Modal app must expose these RPC endpoints (see [`scripts/modal_llm_service.py`](scripts/modal_llm_service.py)).

**Class-based** (typical for GPU models with `@modal.cls`):

| Method | Input | Output |
|--------|-------|--------|
| `complete` | `list[dict]` with `role`, `content` | `str` (or `{"content": "..."}`) |
| `stream` | same | generator yielding `str` tokens |
| `embed` | `list[str]` | `list[list[float]]` |

**Function-based** alternative: set `MODAL_LLM_COMPLETE_FUNCTION`, `MODAL_LLM_STREAM_FUNCTION`, `MODAL_EMBED_FUNCTION` instead of `MODAL_LLM_CLS`.

Client-side usage:

```python
from aiagentrag.providers.modal import ModalRpcClient, ModalRpcConfig, ModalLLMProvider, ModalEmbeddingProvider

config = ModalRpcConfig(app_name="aiagentrag-llm", llm_cls="LLMService", embed_cls="LLMService")
rpc = ModalRpcClient(config.app_name)
llm = ModalLLMProvider(rpc, config)
embedding = ModalEmbeddingProvider(rpc, config)
```

---

## 4. PostgreSQL migrations

Migrations are embedded in the package. Consumer projects do **not** need `alembic.ini`, an `alembic/` folder, or manual SQL.

The PostgreSQL **database** must exist beforehand; the library creates tables.

### In this repository

```bash
export DATABASE_URL="postgresql+asyncpg://postgres:postgres@localhost:5432/aiagentrag"
poetry run aiagentrag-migrate upgrade head
```

Check status:

```bash
poetry run aiagentrag-migrate current
poetry run aiagentrag-migrate history
```

### In a project with the installed package

**CLI** (after `pip install aiagentrag`):

```bash
export DATABASE_URL="postgresql+asyncpg://user:pass@host:5432/mydb"
aiagentrag-migrate upgrade head
aiagentrag-migrate current
aiagentrag-migrate downgrade base   # rollback
```

**Python API:**

```python
from aiagentrag.storage.postgres import upgrade_head, PostgresConversationStore

DATABASE_URL = "postgresql+asyncpg://user:pass@host:5432/mydb"

# migrations only
upgrade_head(DATABASE_URL)

# migrations + ready-to-use store
store = await PostgresConversationStore.initialize(DATABASE_URL)
```

`upgrade_head` is idempotent — safe to call on every application startup.

---

## 5. Using the library in your project

The library does not start infrastructure. You connect PostgreSQL, Qdrant, and an LLM provider, build an `Agent`, and call `run()`.

Full wiring example: [`scripts/try_agent.py`](scripts/try_agent.py).

Minimal outline:

```python
import asyncio

from aiagentrag import Agent, AgentConfig, TokenEvent, FinishedEvent, ErrorEvent
from aiagentrag.knowledge.retriever import KnowledgeRetriever
from aiagentrag.memory.compressor import ConversationCompressor
from aiagentrag.memory.repository import MemoryRepository
from aiagentrag.prompt.builder import PromptBuilder
from aiagentrag.storage.postgres import PostgresConversationStore
from aiagentrag.storage.qdrant.client import QdrantVectorStore

# + your LLM provider (Ollama / OpenAI / Modal RPC)
# + QdrantVectorStore, embedding provider


async def main() -> None:
    config = AgentConfig(
        system_prompt="You are a helpful assistant.",
        knowledge_collection="documents",  # vectorizer collection
    )

    database_url = "postgresql+asyncpg://postgres:postgres@localhost:5432/mydb"
    conversation_store = await PostgresConversationStore.initialize(database_url)

    # vector_store, embedding, llm — see scripts/try_agent.py

    agent = Agent(
        config=config,
        memory_repository=MemoryRepository(conversation_store, vector_store, embedding, config),
        knowledge_retriever=KnowledgeRetriever(vector_store, embedding, config),
        prompt_builder=PromptBuilder(config),
        llm_provider=llm,
        conversation_compressor=ConversationCompressor(
            conversation_store, vector_store, embedding, llm, config
        ),
    )

    async for event in agent.run(user_id="user-1", message="Hello"):
        match event:
            case TokenEvent(content=token):
                print(token, end="", flush=True)
            case FinishedEvent(metadata=meta):
                print(f"\nchunks: {meta.knowledge_chunks_retrieved}")
            case ErrorEvent(error=err):
                print(f"Error: {err}")


asyncio.run(main())
```

### DI (optional)

```python
from aiagentrag import create_container

container = create_container(
    config=config,
    embedding_provider=embedding,
    llm_provider=llm,
    database_url=database_url,
    qdrant_url="http://localhost:6333",
    vector_size=768,
)
```

### How to create an Agent (clear, concrete examples)

Short summary: you must provide
- an EmbeddingProvider (async embed/embed_batch)
- an LLMProvider (stream(messages) -> AsyncIterator[str], complete(messages) -> str)
- PostgreSQL URL and Qdrant URL

Important: the synchronous constructors DO NOT run PostgreSQL migrations. Run migrations once manually before production:

```bash
aiagentrag-migrate upgrade head
```

Below are minimal, copy-paste-ready examples for the three supported providers in this package.

1) Ollama (local Ollama server)

```python
from aiagentrag.core.models import AgentConfig
from aiagentrag.agent.agent import Agent
from aiagentrag.providers.ollama import OllamaEmbeddingProvider, OllamaLLMProvider
from ollama import AsyncClient

# 1) create Ollama async client (point to your Ollama host)
ollama_client = AsyncClient(base_url="http://localhost:11434")

# 2) create provider adapters (these implement EmbeddingProvider / LLMProvider)
embedding = OllamaEmbeddingProvider(ollama_client, model="nomic-embed-text")
llm = OllamaLLMProvider(ollama_client, model="qwen3:8b")

# 3) Agent config
cfg = AgentConfig(system_prompt="You are a helpful assistant.")

# 4) Build Agent (synchronous constructor — does NOT run migrations)
agent = Agent.from_ollama(
    cfg,
    ollama_client=llm,             # must implement LLMProvider
    embedding_provider=embedding,  # must implement EmbeddingProvider
    database_url="postgresql+asyncpg://postgres:postgres@localhost:5432/aiagentrag",
    qdrant_url="http://localhost:6333",
    vector_size=1536,
)

# 5) run (async)
async for event in agent.run(user_id="local-user", message="Hello"):
    match event:
        case TokenEvent(content=tok): print(tok, end="", flush=True)
        case FinishedEvent(metadata=meta): print("\nDone:", meta)
        case ErrorEvent(error=err): raise RuntimeError(err)
```

2) OpenAI (official async client)

```python
from aiagentrag.providers.openai import OpenAIEmbeddingProvider, OpenAILLMProvider
from openai import AsyncOpenAI

openai_client = AsyncOpenAI(api_key="sk-...")
embedding = OpenAIEmbeddingProvider(openai_client, model="text-embedding-3-small")
llm = OpenAILLMProvider(openai_client, model="gpt-4o-mini")

agent = Agent.from_openai(
    cfg,
    openai_client=llm,             # OpenAILLMProvider instance (LLMProvider)
    embedding_provider=embedding,  # OpenAIEmbeddingProvider (EmbeddingProvider)
    database_url=DATABASE_URL,
    qdrant_url=QDRANT_URL,
)
```

3) Modal RPC (deployed Modal app)

```python
from aiagentrag.providers.modal import (
    ModalRpcConfig,
    ModalRpcClient,
    ModalEmbeddingProvider,
    ModalLLMProvider,
)

# Build Modal RPC config (or use modal_config_from_env())
rpc_cfg = ModalRpcConfig(
    app_name="aiagentrag-llm",
    llm_cls="LLMService",           # your deployed class or use functions instead
    embed_cls="LLMService",
)
rpc = ModalRpcClient(app_name=rpc_cfg.app_name, environment_name=rpc_cfg.environment_name)

embedding = ModalEmbeddingProvider(rpc, rpc_cfg)
llm = ModalLLMProvider(rpc, rpc_cfg)

agent = Agent.from_modal(
    cfg,
    modal_client=llm,               # ModalLLMProvider instance (LLMProvider)
    embedding_provider=embedding,   # ModalEmbeddingProvider (EmbeddingProvider)
    database_url=DATABASE_URL,
    qdrant_url=QDRANT_URL,
)
```

Exactly-typed signatures
- Agent.from_ollama(cfg: AgentConfig, *, ollama_client: LLMProvider, embedding_provider: EmbeddingProvider, database_url: str, qdrant_url: str, vector_size: int = 1536) -> Agent
- Agent.from_openai(cfg: AgentConfig, *, openai_client: LLMProvider, embedding_provider: EmbeddingProvider, database_url: str, qdrant_url: str, vector_size: int = 1536) -> Agent
- Agent.from_modal(cfg: AgentConfig, *, modal_client: LLMProvider, embedding_provider: EmbeddingProvider, database_url: str, qdrant_url: str, vector_size: int = 1536) -> Agent

If you already have fully constructed components (MemoryRepository, PromptBuilder, etc.) use:

```python
agent = Agent.from_components(
    config=cfg,
    memory_repository=my_memory_repo,
    knowledge_retriever=my_knowledge,
    prompt_builder=my_prompt_builder,
    llm_provider=my_llm_provider,
    conversation_compressor=my_compressor,
)
```

Notes
- Constructors require objects that satisfy the small protocols in `aiagentrag.core.interfaces`. Use the provider classes in `aiagentrag.providers.*` as drop-in adapters for common setups.
- Run migrations manually before first production run: `aiagentrag-migrate upgrade head`.
### Events

`agent.run()` yields an async stream:

| Event | When |
|-------|------|
| `StatusEvent` | Pipeline step (load history, RAG, LLM, …) |
| `TokenEvent` | Streamed response token |
| `FinishedEvent` | Success + metadata |
| `ErrorEvent` | Failure |

### Qdrant collections

| Collection | Written by | Read by |
|------------|------------|---------|
| `documents` | vectorizer | agent (RAG) |
| `user_memory` | agent | agent |

Document payload (vectorizer): `text` field plus metadata (`document_id`, `chunk_index`, …).

---

## 6. Local development

```bash
git clone https://github.com/leonidsliusar/AIAgentRag
cd AIAgentRag
poetry install

poetry install --extras modal

docker compose up -d    # PostgreSQL + Qdrant
# LLM: Ollama, OpenAI, or Modal RPC; documents via vectorizer

poetry run python -m scripts.try_agent -m "test"

poetry run pytest
poetry run ruff check src tests scripts
poetry run mypy src
poetry build              # verify migrations are included in the wheel
```

In-memory test fakes live in `tests/conftest.py` only — not used for local runs.

---

## Troubleshooting

| Error | Fix |
|-------|-----|
| `Knowledge collection 'documents' was not found` | Ingest documents via vectorizer first |
| `connection refused` on :5432 / :6333 | `docker compose up -d` |
| Ollama streaming / embedding failed | `ollama serve`, `ollama pull nomic-embed-text`, `ollama pull qwen3:8b` |
| Poor RAG quality | Same embedding vectors as used when ingesting documents |
| Modal RPC failed | Credentials set (`modal token set` or `MODAL_TOKEN_*`), app deployed (`modal deploy`) |
| `Token missing` | Set `MODAL_TOKEN_ID` + `MODAL_TOKEN_SECRET` or run `modal token set` |

