# DeepWiki WebSocket API Implementation for Real-Time Chat Streaming

> Discover how the DeepWiki WebSocket API streams real-time chat using an async generator and RAG pipeline. Learn about incremental text updates and provider streaming.

- Repository: [ASYNCFUNC/deepwiki-open](https://github.com/asyncfuncai/deepwiki-open)
- Tags: api-reference
- Published: 2026-02-16

---

**The DeepWiki WebSocket API streams real-time chat responses through an async generator in [`api/websocket_wiki.py`](https://github.com/AsyncFuncAI/deepwiki-open/blob/main/api/websocket_wiki.py), accepting JSON requests, initializing a per-request RAG pipeline, and iterating over provider-specific streaming chunks to deliver incremental text updates.**

The DeepWiki WebSocket API enables low-latency, bidirectional communication for AI-powered chat completions in the `AsyncFuncAI/deepwiki-open` repository. This implementation handles connection lifecycle management, retrieval-augmented generation (RAG) context injection, and multi-provider streaming through a single endpoint.

## WebSocket Endpoint Architecture

The core implementation resides in [`api/websocket_wiki.py`](https://github.com/AsyncFuncAI/deepwiki-open/blob/main/api/websocket_wiki.py), specifically within the `handle_websocket_chat` function. This async handler manages the entire lifecycle of a WebSocket connection, from initial handshake to final disconnection, while coordinating between the client, the RAG retrieval system, and various large language model (LLM) providers.

## Connection Lifecycle and Request Handling

### Accepting the Connection and Receiving Payloads

The handler begins by establishing the bidirectional channel:

```python
await websocket.accept()

```

Once connected, the server awaits a JSON payload containing the chat completion parameters:

```python
request = ChatCompletionRequest(**await websocket.receive_json())

```

This payload is validated against the `ChatCompletionRequest` Pydantic model, ensuring type safety for fields like `provider`, `model`, and `messages`.

### Token Count Guardrails

Before processing, the implementation checks prompt size to prevent context window overflow:

```python
token_count = count_tokens(request.messages)
if token_count > 8000:
    logger.warning(f"Large prompt detected: {token_count} tokens")

```

The `count_tokens` utility analyzes the message list and logs warnings when prompts exceed approximately 8,000 tokens, allowing clients to adjust their requests before hitting provider-specific limits.

## RAG Pipeline Initialization

For each incoming request, the system instantiates a fresh `RAG` object to handle context retrieval:

```python
rag = RAG(provider=request.provider, model=request.model)

```

The retriever is configured with repository-specific filters to narrow the search space:

```python
retriever = rag.get_retriever(
    repo_name=request.repo_name,
    filters={"source": request.source_type}
)

```

This per-request isolation ensures that conversation context remains separate across concurrent connections, while the retrieval system fetches relevant documentation chunks based on the current query.

## Message Validation and Prompt Construction

### Validating Conversation State

The handler enforces structural requirements on the message list:

```python
if not request.messages:
    raise ValueError("At least one user message required")
    
if request.messages[-1].role != "user":
    raise ValueError("Last message must be from user")

```

These checks ensure that the LLM receives a valid conversation thread ending with a user query, preventing malformed requests from reaching the provider APIs.

### Building the Contextual Prompt

The final prompt aggregates multiple context sources:

```python
prompt_parts = [
    system_prompt,
    conversation_history,
    file_content if request.file_path else "",
    retrieved_context
]
full_prompt = "\n\n".join(filter(None, prompt_parts))

```

This construction concatenates the system instructions, previous conversation turns, optional file attachments, and RAG-retrieved documentation into a single coherent context string for the LLM.

## Provider Selection and Streaming Implementation

### Branching to Provider-Specific Clients

The implementation supports seven distinct LLM providers through a branching structure in [`api/websocket_wiki.py`](https://github.com/AsyncFuncAI/deepwiki-open/blob/main/api/websocket_wiki.py):

```python
if request.provider == "ollama":
    client = OllamaClient()
elif request.provider == "openrouter":
    client = OpenRouterClient()
elif request.provider == "openai":
    client = OpenAIClient()

# ... additional providers: bedrock, azure, dashscope, google

```

Each client is initialized with model-specific kwargs retrieved from [`api/config.py`](https://github.com/AsyncFuncAI/deepwiki-open/blob/main/api/config.py), ensuring proper authentication and endpoint configuration.

### Converting Inputs and Initiating Streams

Before calling the model, inputs are converted to provider-specific API arguments:

```python
api_kwargs = client.convert_inputs_to_api_kwargs(
    prompt=full_prompt,
    model=request.model,
    temperature=request.temperature
)

```

The `acall` method returns an async generator that yields response chunks:

```python
stream = await model.acall(
    api_kwargs=api_kwargs,
    model_type=ModelType.LLM
)

async for chunk in stream:
    text = extract_text(chunk)
    await websocket.send_text(text)

```

The handler iterates over this generator, extracts the textual content from each chunk, strips internal markers, and immediately forwards the text to the client via `websocket.send_text()`, achieving true real-time streaming.

## Summary

- The DeepWiki WebSocket API implementation centers on [`api/websocket_wiki.py`](https://github.com/AsyncFuncAI/deepwiki-open/blob/main/api/websocket_wiki.py), specifically the `handle_websocket_chat` function.
- Each connection initializes an isolated **RAG pipeline** to retrieve contextually relevant documentation before generation begins.
- The system enforces **token limits** and **message validation** to prevent malformed requests from reaching LLM providers.
- **Seven provider backends** (Ollama, OpenRouter, OpenAI, Bedrock, Azure, DashScope, Google) are supported through a unified async streaming interface.
- Responses stream in real-time via an **async generator** pattern, with chunks immediately forwarded to the client over the persistent WebSocket connection.

## Frequently Asked Questions

### What file contains the main WebSocket handler for DeepWiki?

The primary implementation resides in [`api/websocket_wiki.py`](https://github.com/AsyncFuncAI/deepwiki-open/blob/main/api/websocket_wiki.py), specifically within the `handle_websocket_chat` async function that manages the full connection lifecycle from handshake to disconnection.

### How does DeepWiki handle context retrieval for chat requests?

Each incoming request instantiates a fresh `RAG` object configured with repository-specific filters. The retriever fetches relevant documentation chunks based on the current query and repository name, injecting this context into the final prompt sent to the LLM.

### Which LLM providers are supported by the WebSocket API?

The implementation supports seven providers: Ollama, OpenRouter, OpenAI, AWS Bedrock, Azure OpenAI, DashScope, and Google Generative AI. Each provider has a dedicated client class that handles authentication and API-specific argument conversion.

### How is real-time streaming achieved in the WebSocket implementation?

The system uses an async generator pattern where `model.acall()` returns a stream object. The handler iterates over this generator with `async for`, extracts text from each chunk, and immediately sends it to the client via `websocket.send_text()`, ensuring low-latency incremental delivery without waiting for the full response.