Metadata-Version: 2.4
Name: ai-query
Version: 1.1.0
Summary: A unified Python SDK for querying AI models from multiple providers
Project-URL: Homepage, https://github.com/Abdulmumin1/ai-query
Project-URL: Repository, https://github.com/Abdulmumin1/ai-query
License: MIT License
        
        Copyright (c) 2026 Abdulmumin Yaqeen
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Requires-Python: >=3.13
Requires-Dist: aiohttp>=3.9.0
Description-Content-Type: text/markdown

# ai-query

A unified Python SDK for querying AI models from multiple providers with a consistent interface.

## Installation

```bash
uv add ai-query
# or
pip install ai-query
```

## Quick Start

```python
import asyncio
from ai_query import generate_text, openai

async def main():
    result = await generate_text(
        model=openai("gpt-4o"),
        prompt="What is the capital of France?"
    )
    print(result.text)

asyncio.run(main())
```

## Streaming

```python
from ai_query import stream_text, google

async def main():
    result = stream_text(
        model=google("gemini-2.0-flash"),
        prompt="Write a short story."
    )

    async for chunk in result.text_stream:
        print(chunk, end="", flush=True)

    usage = await result.usage
    print(f"\nTokens: {usage.total_tokens}")
```

## Tool Calling

Define tools and let the model use them automatically. The library handles the execution loop. Tools can be defined using the `@tool` decorator with type hints and the `Field` class for descriptions.

```python
from ai_query import generate_text, google, tool, Field, step_count_is

# Define tools using decorators
@tool(description="Get the current weather for a location")
async def get_weather(
    location: str = Field(description="City name")
) -> str:
    # Function implementation
    return f"Weather in {location}: 72°F, Sunny"

@tool(description="Perform math calculations")
def calculate(
    expression: str = Field(description="Math expression")
) -> str:
    return str(eval(expression))

async def main():
    result = await generate_text(
        model=google("gemini-2.0-flash"),
        prompt="What's the weather in Paris? Also, what is 25 * 4?",
        tools={
            "weather": get_weather,
            "calculator": calculate
        },
        stop_when=step_count_is(5),  # Max 5 model calls
    )
    print(result.text)
    print(f"Steps: {len(result.response['steps'])}")
```

### Stop Conditions

Control when the tool execution loop stops:

```python
from ai_query import step_count_is, has_tool_call

# Stop after N model calls
stop_when=step_count_is(5)

# Stop when a specific tool is called
stop_when=has_tool_call("final_answer")

# Multiple conditions (stops when any is true)
stop_when=[step_count_is(10), has_tool_call("done")]
```

## Step Callbacks

Monitor and react to each step in the execution loop with `on_step_start` and `on_step_finish`.

```python
from ai_query import generate_text, google, StepStartEvent, StepFinishEvent

def on_start(event: StepStartEvent):
    print(f"Step {event.step_number} starting...")
    print(f"  Messages: {len(event.messages)}")
    # event.messages can be modified before the model call

def on_finish(event: StepFinishEvent):
    print(f"Step {event.step_number} finished")

    # Current step details
    if event.step.tool_calls:
        for tc in event.step.tool_calls:
            print(f"  Called: {tc.name}({tc.arguments})")
    if event.step.tool_results:
        for tr in event.step.tool_results:
            print(f"  Result: {tr.result}")

    # Accumulated state
    print(f"  Total tokens: {event.usage.total_tokens}")
    print(f"  Text so far: {event.text[:50]}...")

result = await generate_text(
    model=google("gemini-2.0-flash"),
    prompt="What's the weather in Tokyo?",
    tools={"weather": get_weather},
    on_step_start=on_start,
    on_step_finish=on_finish,
)
```

### StepStartEvent

| Field | Type | Description |
|-------|------|-------------|
| `step_number` | `int` | 1-indexed step number |
| `messages` | `list[Message]` | Conversation history (modifiable) |
| `tools` | `ToolSet \| None` | Available tools |

### StepFinishEvent

| Field | Type | Description |
|-------|------|-------------|
| `step_number` | `int` | 1-indexed step number |
| `step` | `StepResult` | Current step (text, tool_calls, tool_results) |
| `text` | `str` | Accumulated text from all steps |
| `usage` | `Usage` | Accumulated token usage |
| `steps` | `list[StepResult]` | All completed steps |

Both callbacks support sync and async functions.

## Providers

Built-in support for:

- **OpenAI**: `openai("gpt-4o")` - uses `OPENAI_API_KEY`
- **Anthropic**: `anthropic("claude-sonnet-4-20250514")` - uses `ANTHROPIC_API_KEY`
- **Google**: `google("gemini-2.0-flash")` - uses `GOOGLE_API_KEY`

Pass API keys directly if needed:

```python
model = google("gemini-2.0-flash", api_key="your_key")
```

## Provider Options

Pass provider-specific parameters:

```python
result = await generate_text(
    model=google("gemini-2.0-flash"),
    prompt="Tell me a story",
    provider_options={
        "google": {
            "safety_settings": {"HARM_CATEGORY_VIOLENCE": "BLOCK_NONE"}
        }
    }
)
```

## Examples

See the [examples/](examples/) folder for agent implementations