AIP - AI Protocol

A modern, stateful communication protocol for conversational AI models. AIP provides an efficient WebSocket-based alternative to traditional stateless HTTP REST APIs for AI interactions.

Features

• Stateful Sessions: Maintains conversation context automatically
• Real-time Streaming: Token-by-token response streaming for better UX
• Developer-Friendly: Simple SDK abstracts WebSocket complexity
• Efficient: WebSocket connections reduce overhead compared to HTTP polling
• Extensible: Easy to add new message types and features
• Multiple Providers: Support for both OpenAI and local Ollama models

Two Ways to Use AIP

Option 1: Local Models with Ollama (FREE)

Run AI models completely on your machine - no API key needed!

✅ Free - No API costs ✅ Private - Data never leaves your machine ✅ Fast - No network latency ✅ Offline - Works without internet

📖 See OLLAMA_GUIDE.md for complete instructions

Quick Start:

# 1. Make sure Ollama is running (with your models)
ollama list

# 2. Start AIP server with Ollama
./run_with_model.sh
# or: python server_ollama.py

# 3. Run examples
python example_ollama.py

Option 2: OpenAI API (Cloud)

Use OpenAI's GPT models (requires API key and costs money)

Architecture

The AIP MVP consists of two main components:

1. Server: FastAPI-based WebSocket server that manages sessions
   • server.py - OpenAI version (requires API key)
   • server_ollama.py - Ollama version (free, local)
2. Client SDK (aip_sdk.py): Python library that provides a simple interface for connecting and communicating

Quick Start (OpenAI)

Prerequisites

• Python 3.8 or higher
• OpenAI API key

Installation

1. Clone or download this repository

2. Install dependencies:

pip install -r requirements.txt

3. Set your OpenAI API key:

export OPENAI_API_KEY="your-api-key-here"

Running the Server

Start the AIP server:

python server.py

The server will be available at ws://localhost:8000/aip

You can verify it's running by visiting http://localhost:8000 in your browser.

Running the Example

In a new terminal, run the example client:

python example.py

This will demonstrate:

• Connecting to the server
• Asking questions with streamed responses
• Context retention across multiple questions
• Clean disconnection

Usage

Basic Example

import asyncio
from aip_sdk import AIPClient

async def main():
    # Connect to the server
    client = await AIPClient.connect("ws://localhost:8000/aip")

    # Ask a question and print the streamed response
    await client.ask(
        "What is Python?",
        on_token=lambda token: print(token, end="", flush=True)
    )

    # Ask a follow-up (context is maintained)
    await client.ask(
        "What are its main uses?",
        on_token=lambda token: print(token, end="", flush=True)
    )

    # Disconnect
    await client.disconnect()

asyncio.run(main())

Using Context Manager

For automatic cleanup, use the async context manager:

async def main():
    async with await AIPClient.connect("ws://localhost:8000/aip") as client:
        await client.ask("Hello!", on_token=print)
        # Automatically disconnected when exiting the block

Accumulating Response

Instead of printing tokens, you can accumulate them:

response_tokens = []

await client.ask(
    "Tell me a story",
    on_token=response_tokens.append
)

full_response = "".join(response_tokens)
print(full_response)

Protocol Specification

Connection Flow

1. Client initiates WebSocket connection to /aip
2. Server accepts connection and generates unique Session ID
3. Server sends SESSION_ID message to client
4. Connection is ready for use

Message Types

Client → Server

ASK: Send a prompt to the AI

{
  "type": "ASK",
  "payload": "Your question here"
}

Server → Client

SESSION_ID: Initial message containing session identifier

{
  "type": "SESSION_ID",
  "payload": "550e8400-e29b-41d4-a716-446655440000"
}

TOKEN: A single token from the streaming response

{
  "type": "TOKEN",
  "payload": "Hello"
}

DONE: Indicates the response stream is complete

{
  "type": "DONE",
  "payload": ""
}

ERROR: Error message

{
  "type": "ERROR",
  "payload": "Error description"
}

Project Structure

aprotoc/
├── server.py           # FastAPI WebSocket server
├── aip_sdk.py          # Python client SDK
├── example.py          # Usage examples
├── requirements.txt    # Python dependencies
└── README.md          # This file

SDK API Reference

AIPClient

await AIPClient.connect(host: str) -> AIPClient

Connect to an AIP server.

• Parameters:
  • host: WebSocket URL (e.g., ws://localhost:8000/aip)
• Returns: Authenticated AIPClient instance
• Raises: Exception if connection fails

await client.ask(prompt: str, on_token: Callable[[str], None]) -> None

Send a prompt and receive streamed response.

• Parameters:
  • prompt: The question or prompt to send
  • on_token: Callback function called for each token (optional)
• Raises: Exception if an error occurs

await client.disconnect() -> None

Close the WebSocket connection.

Properties

• client.session_id: Get the current session ID
• client.is_connected: Check if client is connected

Configuration

Server Configuration

Edit server.py to customize:

• Port: Change in uvicorn.run() (default: 8000)
• Host: Change in uvicorn.run() (default: 0.0.0.0)
• Model: Change model parameter in OpenAI call (default: gpt-4o)
• Temperature: Change temperature parameter (default: 0.7)

Session Management

Currently uses in-memory storage. For production, consider:

• Redis for distributed session storage
• Database for persistent conversation history
• TTL-based session expiration

Development

Running Tests

(Tests to be implemented)

pytest tests/

Code Style

Follow PEP 8 guidelines. Format with:

black server.py aip_sdk.py example.py

Roadmap

Future enhancements:

• Authentication and authorization
• Multiple AI provider support (Anthropic, Google, etc.)
• Persistent session storage
• Rate limiting
• Client SDKs in other languages (JavaScript, Go, etc.)
• Message history retrieval
• Custom system prompts per session
• Conversation branching
• File/image upload support

Troubleshooting

Connection Refused

• Ensure the server is running: python server.py
• Check the server URL in your client code
• Verify no firewall is blocking port 8000

OpenAI API Errors

• Verify your API key is set: echo $OPENAI_API_KEY
• Check your OpenAI account has credits
• Ensure you have access to the gpt-4o model

Import Errors

• Install dependencies: pip install -r requirements.txt
• Verify you're using Python 3.8+: python --version

Contributing

Contributions are welcome! Please:

1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Add tests if applicable
5. Submit a pull request

License

MIT License - feel free to use this in your own projects!

Acknowledgments

Built with:

• FastAPI - Modern web framework
• OpenAI - AI model provider
• websockets - WebSocket client library

---

Note: This is an MVP (Minimum Viable Product). It's designed for demonstration and development purposes. For production use, consider adding authentication, rate limiting, persistent storage, and comprehensive error handling.