# animuz-core

Core shared utilities for Animuz RAG (Retrieval-Augmented Generation) system.

## Features

- **Unified RAG API**: 3-knob interface (`prompt`, `llm`, `tools`) for chat, ingestion, and retrieval
- **LLM Clients**: OpenAI (Responses API + Chat Completions), Anthropic Claude, Ollama
- **Tool System**: `@tool` decorator, `qdrant_retriever()` factory, MCP server support
- **RAG Pipelines**: Simple and Agentic RAG implementations (lower-level building blocks)
- **Vector Database**: Qdrant integration with hybrid search (dense + sparse)
- **Embedding Clients**: Multiple providers (local server, Modal, S3/SageMaker)
- **Document Ingestion**: Azure Document Intelligence, Unstructured, PDF extraction, structured text parsing
- **CloudWatch Logging**: Structured JSON logging with watchtower

## Requirements

- Python >= 3.10

## Installation

Install the core package (minimal dependencies only):

```bash
pip install animuz-core
```

Then install only the extras you need:

```bash
# Single extra
pip install animuz-core[openai]

# Multiple extras
pip install animuz-core[openai,qdrant,aws]

# Everything
pip install animuz-core[all]
```

Works with uv too:

```bash
uv pip install animuz-core[openai,qdrant]
```

### Available Extras

| Extra       | What it installs                                  | Use when you need                                   |
| ----------- | ------------------------------------------------- | --------------------------------------------------- |
| `openai`    | `openai`                                          | OpenAI GPT models                                   |
| `anthropic` | `anthropic`                                       | Anthropic Claude models                             |
| `ollama`    | `ollama`                                          | Local LLMs via Ollama                               |
| `qdrant`    | `qdrant-client`                                   | Qdrant vector database                              |
| `aws`       | `boto3`, `aiobotocore`, `watchtower`, `sagemaker` | S3, SageMaker embeddings, CloudWatch logging        |
| `azure`     | `azure-ai-documentintelligence`                   | Azure Document Intelligence for PDF ingestion       |
| `ingest`    | `unstructured-client`, `PyMuPDF`                  | Document parsing (Unstructured API, PDF extraction) |
| `fastapi`   | `fastapi`                                         | Streaming SSE endpoints                             |
| `all`       | All of the above                                  | Everything                                          |
| `dev`       | `all` + `pytest`, `black`, `ruff`, `mypy`         | Development and testing                             |

## Usage

### Unified RAG API (recommended)

The `RAG` class is the main entry point. It has 3 knobs:

- **`prompt`** — callable `(team_id, assistant_id) -> dict` that fetches an assistant config
- **`llm`** — model name string (provider auto-detected: `"gpt-*"` → OpenAI, `"claude-*"` → Anthropic)
- **`tools`** — `list[ToolSpec]` for local tools, `MCP(url=...)` for MCP server, or `None` for plain chat

```python
from animuz_core import RAG, qdrant_retriever, tool

# Define how to fetch the assistant config
def my_fetcher(team_id, assistant_id):
    return {"prompt": "You are a helpful assistant.", "model": "gpt-4o-mini"}

# Create RAG with local retriever tool
rag = RAG(
    prompt=my_fetcher,
    llm="gpt-4o-mini",
    tools=[
        qdrant_retriever(host="localhost", port=6333, collection="animuz"),
    ],
)

# Chat (returns frontend-ready output)
output = await rag.chat("team1", "asst1", [{"role": "user", "content": "Hello"}])

# Ingest a document
await rag.add_doc("docs/intro.md", user_chat_id="team1|asst1")

# Retrieve documents
texts, points = await rag.retrieve("what is this?", user_chat_id="team1|asst1")
```

#### With MCP tools (Lambda / cloud)

```python
from animuz_core import RAG, MCP

rag = RAG(
    prompt=ddb_fetcher,
    llm="gpt-4o-mini",
    tools=MCP(url=os.environ["MCP_URL"], api_key=os.environ.get("MCP_API_KEY")),
)
output = await rag.chat(team_id, assistant_id, messages, user_context=ctx)
```

#### Plain chat (no tools)

```python
rag = RAG(prompt=my_fetcher, llm="gpt-4o-mini")
output = await rag.chat(team_id, assistant_id, messages)
```

#### Custom OpenAI base URL (proxy / gateway)

Route OpenAI API calls through a proxy or API gateway by passing `base_url`:

```python
rag = RAG(
    prompt=my_fetcher,
    llm="gpt-4o-mini",
    base_url="https://your-proxy.example.com/openai/v1",
    tools=[qdrant_retriever(...)],
)
```

When omitted, the default OpenAI endpoint (`api.openai.com`) is used.

#### Custom tools with @tool decorator

```python
from animuz_core import RAG, tool, qdrant_retriever

@tool(description="Get weather for a city")
async def weather(city: str) -> str:
    return await fetch_weather(city)

rag = RAG(
    prompt=my_fetcher,
    llm="gpt-4o-mini",
    tools=[qdrant_retriever(...), weather],
)
```

### Lower-level APIs

#### LLM Clients

```python
from animuz_core import OpenAIAgentClientResponses

# OpenAI Responses API agent with tool loop (recommended for production)
agent = OpenAIAgentClientResponses(user_chat_id="tenant-123", tools=tool_dict, model="gpt-4o-mini")
result = await agent.get_reply_frontend(messages, system_prompt)

# With a custom base URL (proxy / gateway)
agent = OpenAIAgentClientResponses(model="gpt-4o-mini", base_url="https://your-proxy.example.com/openai/v1")
```

#### RAG Pipelines

```python
from animuz_core import SimpleRAG

# Simple RAG — always retrieves then generates
pipeline = SimpleRAG(
    embedding_client=embedding_client,
    db_client=qdrant_client,
    LLM=llm_client,
)
await pipeline.add_doc("document.pdf", user_chat_id="tenant-123")
result = await pipeline.query("What is RAG?", user_chat_id="tenant-123")
```

#### Vector Database

```python
from animuz_core import QdrantDBClient

client = QdrantDBClient(host="localhost", port=6333, collection_name="animuz")
results = await client.hybrid_search(dense_vec, indices, values, user_chat_id="tenant-123")
```

#### Embedding

```python
from animuz_core import EmbeddingClient, ModalEmbeddingClient

# Local embedding server
client = EmbeddingClient(host="localhost", port=12081)
result = await client.get_embedding("Some text")

# Modal-hosted embeddings
client = ModalEmbeddingClient()
result = client.get_embedding("Some text")
```

## Development

```bash
# Clone and install in editable mode with dev dependencies
git clone <repo-url>
cd animuz-core
pip install -e ".[dev]"

# Run tests
pytest tests/

# Run integration tests (requires external services + env vars)
pytest -m integration tests/integration/
pytest -m integration tests/integration/test_e2e_rag_wrapper_simple.py

# Format
black src/
ruff check src/
```

## Publishing to PyPI

1. Bump the version in `pyproject.toml` and `__init__.py`.

2. Build the package:

```bash
uv pip install --upgrade build
uv run python -m build
```

3. (Optional) Verify the artifacts:

```bash
uv pip install --upgrade twine
uv run python -m twine check dist/*
```

4. Upload to TestPyPI first:

```bash
uv run python -m twine upload -r testpypi dist/*
```

5. Upload to PyPI:

```bash
uv run python -m twine upload dist/*
```

Notes:

- Create a PyPI API token and set `TWINE_USERNAME=__token__` and `TWINE_PASSWORD=<your-token>`.
- If you upload to TestPyPI, install with `pip install -i https://test.pypi.org/simple animuz-core` to verify.

### Integration Test Setup (Qdrant)

Use Docker Compose to run Qdrant locally:

```bash
docker compose -f docker-compose-qdrant.yml up -d qdrant
```

Then set the Qdrant env vars (example):

```bash
export QDRANT_HOST=localhost
export QDRANT_PORT=6333
```

## Environment Variables

The package reads configuration from environment variables (loaded via `python-dotenv`):

| Variable                                               | Used by                     |
| ------------------------------------------------------ | --------------------------- |
| `OPENAI_API_KEY`                                       | OpenAI client               |
| `ANTHROPIC_API_KEY`                                    | Anthropic client            |
| `QDRANT_HOST`, `QDRANT_PORT`, `QDRANT_COLLECTION_NAME` | Qdrant client               |
| `QDRANT_CLOUD_API_KEY`                                 | Qdrant Cloud                |
| `EMBEDDING_HOST`, `EMBEDDING_PORT`                     | Embedding client            |
| `AZURE_DOCAI_KEY`, `AZURE_DOCAI_ENDPOINT`              | Azure Document Intelligence |
| `UNSTRUCTURED_ENDPOINT`, `UNSTRUCTURED_API_KEY`        | Unstructured client         |
| `S3_BUCKET_NAME`, `S3_DOWNLOAD_DIR`                    | S3 operations               |
| `MCP_API_KEY`                                          | MCP tool server             |

## License

MIT
